From 13d730ae5c76f2efe2729624512068deef5323b9 Mon Sep 17 00:00:00 2001 From: Akash Mondal Date: Mon, 9 Jun 2025 19:10:15 +0000 Subject: [PATCH] fix: improve types according to new typecheck --- examples/doc-examples/example_net_stream.py | 5 +++-- libp2p/network/stream/net_stream.py | 11 ++++------- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/examples/doc-examples/example_net_stream.py b/examples/doc-examples/example_net_stream.py index 62dff05e..d8842bea 100644 --- a/examples/doc-examples/example_net_stream.py +++ b/examples/doc-examples/example_net_stream.py @@ -193,7 +193,7 @@ async def run_enhanced_demo( print("🖥️ ENHANCED ECHO SERVER MODE") print("=" * 60) - # Set the enhanced stream handler + # type: ignore: Stream is type of NetStream host.set_stream_handler(PROTOCOL_ID, enhanced_echo_handler) print( @@ -217,7 +217,8 @@ async def run_enhanced_demo( # Create stream and run enhanced demo stream = await host.new_stream(info.peer_id, [PROTOCOL_ID]) - await enhanced_client_demo(stream) + if isinstance(stream, NetStream): + await enhanced_client_demo(stream) print("\n" + "=" * 60) print("CLIENT DEMO COMPLETE") diff --git a/libp2p/network/stream/net_stream.py b/libp2p/network/stream/net_stream.py index 073abfb4..b54fdda4 100644 --- a/libp2p/network/stream/net_stream.py +++ b/libp2p/network/stream/net_stream.py @@ -1,9 +1,6 @@ from enum import ( Enum, ) -from typing import ( - Optional, -) import trio @@ -102,11 +99,11 @@ class NetStream(INetStream): """ muxed_stream: IMuxedStream - protocol_id: Optional[TProtocol] + protocol_id: TProtocol | None __stream_state: StreamState def __init__( - self, muxed_stream: IMuxedStream, nursery: Optional[trio.Nursery] = None + self, muxed_stream: IMuxedStream, nursery: trio.Nursery | None = None ) -> None: super().__init__() @@ -142,7 +139,7 @@ class NetStream(INetStream): async with self._state_lock: return self.__stream_state - async def read(self, n: Optional[int] = None) -> bytes: + async def read(self, n: int | None = None) -> bytes: """ Read from stream. @@ -279,7 +276,7 @@ class NetStream(INetStream): if hasattr(swarm, "refs") and hasattr(swarm.refs, "done"): swarm.refs.done() - def get_remote_address(self) -> Optional[tuple[str, int]]: + def get_remote_address(self) -> tuple[str, int] | None: """Delegate to the underlying muxed stream.""" return self.muxed_stream.get_remote_address()