Updates websocket + bumps version

Author: ibraheem-abe

bittensor_cli/src/bittensor/async_substrate_interface.py

@@ -4,7 +4,7 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from hashlib import blake2b
-from typing import Optional, Any, Union, Callable, Awaitable, cast
+from typing import Optional, Any, Union, Callable, Awaitable, cast, TYPE_CHECKING
 
 from bt_decode import PortableRegistry, decode as decode_by_type_string, MetadataV15
 from async_property import async_property
@@ -19,7 +19,11 @@
     BlockNotFound,
 )
 from substrateinterface.storage import StorageKey
-import websockets
+from websockets.asyncio.client import connect
+from websockets.exceptions import ConnectionClosed
+
+if TYPE_CHECKING:
+    from websockets.asyncio.client import ClientConnection
 
 ResultHandler = Callable[[dict, Any], Awaitable[tuple[dict, bool]]]
 
@@ -433,7 +437,7 @@ def add_item(
             self.block_hashes[block_hash] = runtime
 
     def retrieve(
-        self, block: Optional[int], block_hash: Optional[str]
+        self, block: Optional[int] = None, block_hash: Optional[str] = None
     ) -> Optional["Runtime"]:
         if block is not None:
             return self.blocks.get(block)
@@ -624,7 +628,7 @@ def __init__(
         # TODO allow setting max concurrent connections and rpc subscriptions per connection
         # TODO reconnection logic
         self.ws_url = ws_url
-        self.ws: Optional[websockets.WebSocketClientProtocol] = None
+        self.ws: Optional["ClientConnection"] = None
         self.id = 0
         self.max_subscriptions = max_subscriptions
         self.max_connections = max_connections
@@ -646,15 +650,12 @@ async def __aenter__(self):
                 self._exit_task.cancel()
             if not self._initialized:
                 self._initialized = True
-                await self._connect()
+                self.ws = await asyncio.wait_for(
+                    connect(self.ws_url, **self._options), timeout=10
+                )
                 self._receiving_task = asyncio.create_task(self._start_receiving())
         return self
 
-    async def _connect(self):
-        self.ws = await asyncio.wait_for(
-            websockets.connect(self.ws_url, **self._options), timeout=10
-        )
-
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         async with self._lock:
             self._in_use -= 1
@@ -695,9 +696,7 @@ async def shutdown(self):
 
     async def _recv(self) -> None:
         try:
-            response = json.loads(
-                await cast(websockets.WebSocketClientProtocol, self.ws).recv()
-            )
+            response = json.loads(await self.ws.recv())
             async with self._lock:
                 self._open_subscriptions -= 1
             if "id" in response:
@@ -706,7 +705,7 @@ async def _recv(self) -> None:
                 self._received[response["params"]["subscription"]] = response
             else:
                 raise KeyError(response)
-        except websockets.ConnectionClosed:
+        except ConnectionClosed:
             raise
         except KeyError as e:
             raise e
@@ -717,7 +716,7 @@ async def _start_receiving(self):
                 await self._recv()
         except asyncio.CancelledError:
             pass
-        except websockets.ConnectionClosed:
+        except ConnectionClosed:
             # TODO try reconnect, but only if it's needed
             raise
 
@@ -734,7 +733,7 @@ async def send(self, payload: dict) -> int:
         try:
             await self.ws.send(json.dumps({**payload, **{"id": original_id}}))
             return original_id
-        except websockets.ConnectionClosed:
+        except ConnectionClosed:
             raise
 
     async def retrieve(self, item_id: int) -> Optional[dict]:
@@ -775,7 +774,6 @@ def __init__(
             chain_endpoint,
             options={
                 "max_size": 2**32,
-                "read_limit": 2**16,
                 "write_limit": 2**16,
             },
         )
@@ -1135,7 +1133,7 @@ async def create_storage_key(
         -------
         StorageKey
         """
-        runtime = await self.init_runtime(block_hash=block_hash)
+        await self.init_runtime(block_hash=block_hash)
 
         return StorageKey.create_from_storage_function(
             pallet,
@@ -1555,7 +1553,7 @@ async def _process_response(
         self,
         response: dict,
         subscription_id: Union[int, str],
-        value_scale_type: Optional[str],
+        value_scale_type: Optional[str] = None,
         storage_item: Optional[ScaleType] = None,
         runtime: Optional[Runtime] = None,
         result_handler: Optional[ResultHandler] = None,

bittensor_cli/src/commands/subnets.py

@@ -629,7 +629,7 @@ def format_liquidity_cell(
             tempo_cell = (
                 (f"{subnet.blocks_since_last_step}/{subnet.tempo}{block_change_text}")
                 if netuid != 0
-                else "[red]N/A[/red]"
+                else "-/-"
             )
 
             rows.append(

requirements.txt

@@ -15,6 +15,6 @@ rich~=13.7
 scalecodec==1.2.11
 substrate-interface~=1.7.9
 typer~=0.12
-websockets==13.0
+websockets>=14.1
 bittensor-wallet>=2.0.2
 bt-decode==0.2.0a0