Enhance Type Hinting for WebSocket Connection (#437)

michaelhly · Jun 11, 2024 · 7f2f3cb · 7f2f3cb
1 parent 27a3468
commit 7f2f3cb
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/solana/rpc/websocket_api.py b/src/solana/rpc/websocket_api.py
@@ -416,4 +416,11 @@ def __init__(self, uri: str = "ws://localhost:8900", **kwargs: Any) -> None:
             uri: The websocket endpoint.
             **kwargs: Keyword arguments for ``websockets.legacy.client.connect``
         """
-        super().__init__(uri, **kwargs, create_protocol=SolanaWsClientProtocol)
+        # Ensure that create_protocol explicitly creates a SolanaWsClientProtocol
+        kwargs.setdefault("create_protocol", SolanaWsClientProtocol)
+        super().__init__(uri, **kwargs)
+
+    async def __aenter__(self) -> SolanaWsClientProtocol:
+        """Overrides to specify the type of protocol explicitly."""
+        protocol = await super().__aenter__()
+        return cast(SolanaWsClientProtocol, protocol)