diff --git a/async_substrate_interface/async_substrate.py b/async_substrate_interface/async_substrate.py index 53de792..ecf9360 100644 --- a/async_substrate_interface/async_substrate.py +++ b/async_substrate_interface/async_substrate.py @@ -24,12 +24,7 @@ import asyncstdlib as a from bittensor_wallet.keypair import Keypair from bittensor_wallet.utils import SS58_FORMAT -from bt_decode import ( - MetadataV15, - PortableRegistry, - decode as decode_by_type_string, - encode as encode_by_type_string, -) +from bt_decode import MetadataV15, PortableRegistry, decode as decode_by_type_string from scalecodec.base import ScaleBytes, ScaleType, RuntimeConfigurationObject from scalecodec.types import ( GenericCall, @@ -805,6 +800,50 @@ async def load_registry(self): ) self.registry = PortableRegistry.from_metadata_v15(self.metadata_v15) + async def _wait_for_registry(self, _attempt: int = 1, _retries: int = 3) -> None: + async def _waiter(): + while self.registry is None: + await asyncio.sleep(0.1) + return + + try: + if not self.registry: + await asyncio.wait_for(_waiter(), timeout=10) + except TimeoutError: + # indicates that registry was never loaded + if not self._initializing: + raise AttributeError( + "Registry was never loaded. This did not occur during initialization, which usually indicates " + "you must first initialize the AsyncSubstrateInterface object, either with " + "`await AsyncSubstrateInterface.initialize()` or running with `async with`" + ) + elif _attempt < _retries: + await self.load_registry() + return await self._wait_for_registry(_attempt + 1, _retries) + else: + raise AttributeError( + "Registry was never loaded. This occurred during initialization, which usually indicates a " + "connection or node error." + ) + + async def encode_scale( + self, type_string, value: Any, _attempt: int = 1, _retries: int = 3 + ) -> bytes: + """ + Helper function to encode arbitrary data into SCALE-bytes for given RUST type_string + + Args: + type_string: the type string of the SCALE object for decoding + value: value to encode + _attempt: the current number of attempts to load the registry needed to encode the value + _retries: the maximum number of attempts to load the registry needed to encode the value + + Returns: + encoded bytes + """ + await self._wait_for_registry(_attempt, _retries) + return self._encode_scale(type_string, value) + async def decode_scale( self, type_string: str, @@ -812,7 +851,7 @@ async def decode_scale( _attempt=1, _retries=3, return_scale_obj=False, - ) -> Any: + ) -> Union[ScaleObj, Any]: """ Helper function to decode arbitrary SCALE-bytes (e.g. 0x02000000) according to given RUST type_string (e.g. BlockNumber). The relevant versioning information of the type (if defined) will be applied if block_hash @@ -828,95 +867,19 @@ async def decode_scale( Returns: Decoded object """ - - async def _wait_for_registry(): - while self.registry is None: - await asyncio.sleep(0.1) - return - if scale_bytes == b"\x00": obj = None if type_string == "scale_info::0": # Is an AccountId # Decode AccountId bytes to SS58 address return bytes.fromhex(ss58_decode(scale_bytes, SS58_FORMAT)) else: - try: - if not self.registry: - await asyncio.wait_for(_wait_for_registry(), timeout=10) - obj = decode_by_type_string(type_string, self.registry, scale_bytes) - except TimeoutError: - # indicates that registry was never loaded - if not self._initializing: - raise AttributeError( - "Registry was never loaded. This did not occur during initialization, which usually indicates " - "you must first initialize the AsyncSubstrateInterface object, either with " - "`await AsyncSubstrateInterface.initialize()` or running with `async with`" - ) - elif _attempt < _retries: - await self.load_registry() - return await self.decode_scale( - type_string, scale_bytes, _attempt + 1 - ) - else: - raise AttributeError( - "Registry was never loaded. This occurred during initialization, which usually indicates a " - "connection or node error." - ) + await self._wait_for_registry(_attempt, _retries) + obj = decode_by_type_string(type_string, self.registry, scale_bytes) if return_scale_obj: return ScaleObj(obj) else: return obj - async def encode_scale(self, type_string, value: Any) -> bytes: - """ - Helper function to encode arbitrary data into SCALE-bytes for given RUST type_string - - Args: - type_string: the type string of the SCALE object for decoding - value: value to encode - - Returns: - encoded SCALE bytes - """ - if value is None: - result = b"\x00" - else: - if type_string == "scale_info::0": # Is an AccountId - # encode string into AccountId - ## AccountId is a composite type with one, unnamed field - return bytes.fromhex(ss58_decode(value, SS58_FORMAT)) - - elif type_string == "scale_info::151": # Vec - if not isinstance(value, (list, tuple)): - value = [value] - - # Encode length - length = len(value) - if length < 64: - result = bytes([length << 2]) # Single byte mode - else: - raise ValueError("Vector length too large") - - # Encode each AccountId - for account in value: - if isinstance(account, bytes): - result += account # Already encoded - else: - result += bytes.fromhex( - ss58_decode(value, SS58_FORMAT) - ) # SS58 string - return result - - if isinstance(value, ScaleType): - if value.data.data is not None: - # Already encoded - return bytes(value.data.data) - else: - value = value.value # Unwrap the value of the type - - result = bytes(encode_by_type_string(type_string, self.registry, value)) - return result - async def _first_initialize_runtime(self): """ TODO docstring diff --git a/async_substrate_interface/sync_substrate.py b/async_substrate_interface/sync_substrate.py index 68b0774..1aa45db 100644 --- a/async_substrate_interface/sync_substrate.py +++ b/async_substrate_interface/sync_substrate.py @@ -6,12 +6,7 @@ from bittensor_wallet.keypair import Keypair from bittensor_wallet.utils import SS58_FORMAT -from bt_decode import ( - MetadataV15, - PortableRegistry, - decode as decode_by_type_string, - encode as encode_by_type_string, -) +from bt_decode import MetadataV15, PortableRegistry, decode as decode_by_type_string from scalecodec import ( GenericCall, GenericExtrinsic, @@ -630,56 +625,6 @@ def decode_scale( else: return obj - def encode_scale(self, type_string, value: Any) -> bytes: - """ - Helper function to encode arbitrary data into SCALE-bytes for given RUST type_string - - Args: - type_string: the type string of the SCALE object for decoding - value: value to encode - - Returns: - encoded SCALE bytes - """ - if value is None: - result = b"\x00" - else: - if type_string == "scale_info::0": # Is an AccountId - # encode string into AccountId - ## AccountId is a composite type with one, unnamed field - return bytes.fromhex(ss58_decode(value, SS58_FORMAT)) - - elif type_string == "scale_info::151": # Vec - if not isinstance(value, (list, tuple)): - value = [value] - - # Encode length - length = len(value) - if length < 64: - result = bytes([length << 2]) # Single byte mode - else: - raise ValueError("Vector length too large") - - # Encode each AccountId - for account in value: - if isinstance(account, bytes): - result += account # Already encoded - else: - result += bytes.fromhex( - ss58_decode(value, SS58_FORMAT) - ) # SS58 string - return result - - if isinstance(value, ScaleType): - if value.data.data is not None: - # Already encoded - return bytes(value.data.data) - else: - value = value.value # Unwrap the value of the type - - result = bytes(encode_by_type_string(type_string, self.registry, value)) - return result - def _first_initialize_runtime(self): """ TODO docstring @@ -2933,3 +2878,5 @@ def close(self): self.ws.shutdown() except AttributeError: pass + + encode_scale = SubstrateMixin._encode_scale diff --git a/async_substrate_interface/types.py b/async_substrate_interface/types.py index 93ba00b..a2237f3 100644 --- a/async_substrate_interface/types.py +++ b/async_substrate_interface/types.py @@ -6,7 +6,8 @@ from datetime import datetime from typing import Optional, Union, Any -from bt_decode import PortableRegistry +from bt_decode import PortableRegistry, encode as encode_by_type_string +from bittensor_wallet.utils import SS58_FORMAT from scalecodec import ss58_encode, ss58_decode, is_valid_ss58_address from scalecodec.base import RuntimeConfigurationObject, ScaleBytes from scalecodec.type_registry import load_type_registry_preset @@ -705,3 +706,53 @@ def make_payload(id_: str, method: str, params: list) -> dict: "id": id_, "payload": {"jsonrpc": "2.0", "method": method, "params": params}, } + + def _encode_scale(self, type_string, value: Any) -> bytes: + """ + Helper function to encode arbitrary data into SCALE-bytes for given RUST type_string + + Args: + type_string: the type string of the SCALE object for decoding + value: value to encode + + Returns: + encoded bytes + """ + if value is None: + result = b"\x00" + else: + if type_string == "scale_info::0": # Is an AccountId + # encode string into AccountId + ## AccountId is a composite type with one, unnamed field + return bytes.fromhex(ss58_decode(value, SS58_FORMAT)) + + elif type_string == "scale_info::151": # Vec + if not isinstance(value, (list, tuple)): + value = [value] + + # Encode length + length = len(value) + if length < 64: + result = bytes([length << 2]) # Single byte mode + else: + raise ValueError("Vector length too large") + + # Encode each AccountId + for account in value: + if isinstance(account, bytes): + result += account # Already encoded + else: + result += bytes.fromhex( + ss58_decode(value, SS58_FORMAT) + ) # SS58 string + return result + + if isinstance(value, ScaleType): + if value.data.data is not None: + # Already encoded + return bytes(value.data.data) + else: + value = value.value # Unwrap the value of the type + + result = bytes(encode_by_type_string(type_string, self.registry, value)) + return result diff --git a/tests/unit_tests/asyncio/test_substrate_interface.py b/tests/unit_tests/asyncio/test_substrate_interface.py index c9d28b8..b1ee98b 100644 --- a/tests/unit_tests/asyncio/test_substrate_interface.py +++ b/tests/unit_tests/asyncio/test_substrate_interface.py @@ -1,10 +1,10 @@ import unittest.mock import pytest -import scalecodec.base from websockets.exceptions import InvalidURI from async_substrate_interface.async_substrate import AsyncSubstrateInterface +from async_substrate_interface.types import ScaleObj @pytest.mark.asyncio @@ -59,7 +59,7 @@ async def test_runtime_call(monkeypatch): "SubstrateMethod", ) - assert isinstance(result, scalecodec.base.ScaleType) + assert isinstance(result, ScaleObj) assert result.value is substrate.decode_scale.return_value substrate.rpc_request.assert_called_once_with( diff --git a/tests/unit_tests/sync/test_substrate_interface.py b/tests/unit_tests/sync/test_substrate_interface.py index 0373c6a..18e85ea 100644 --- a/tests/unit_tests/sync/test_substrate_interface.py +++ b/tests/unit_tests/sync/test_substrate_interface.py @@ -1,8 +1,7 @@ import unittest.mock -import scalecodec.base - from async_substrate_interface.sync_substrate import SubstrateInterface +from async_substrate_interface.types import ScaleObj def test_runtime_call(monkeypatch): @@ -45,7 +44,7 @@ def test_runtime_call(monkeypatch): "SubstrateMethod", ) - assert isinstance(result, scalecodec.base.ScaleType) + assert isinstance(result, ScaleObj) assert result.value is substrate.decode_scale.return_value substrate.rpc_request.assert_called_once_with(