typing: strengthen types in networking and discovery transport (#452)

* typing: strengthen types in networking and discovery transport

* fix(discovery): use RequestId types in transport internal state and tests

* fix(tests): update RequestId equality assertions in test_transport

* style: clean up stray inline comments in transport.py

Author: morelucks

src/lean_spec/subspecs/networking/client/event_source.py

@@ -874,7 +874,8 @@ async def _dial_quic(self, multiaddr: str) -> QuicConnection:
             QuicTransportError: If connection fails.
         """
         await self._ensure_quic_manager()
-        return await self.quic_manager.connect(multiaddr)  # type: ignore[union-attr]
+        assert self.quic_manager is not None
+        return await self.quic_manager.connect(multiaddr)
 
     async def listen(self, multiaddr: str) -> None:
         """Start listening for incoming connections.
@@ -905,7 +906,8 @@ async def _listen_quic(self, multiaddr: str) -> None:
             multiaddr: QUIC address to listen on.
         """
         await self._ensure_quic_manager()
-        await self.quic_manager.listen(  # type: ignore[union-attr]
+        assert self.quic_manager is not None
+        await self.quic_manager.listen(
             multiaddr,
             on_connection=self._handle_inbound_quic_connection,
         )

src/lean_spec/subspecs/networking/discovery/transport.py

@@ -68,7 +68,7 @@
 class PendingRequest:
     """Tracks a pending request awaiting response."""
 
-    request_id: bytes
+    request_id: RequestId
     """Request ID for matching responses."""
 
     dest_node_id: NodeId
@@ -95,7 +95,7 @@ class PendingMultiRequest:
     across UDP packets when results exceed MTU.
     """
 
-    request_id: bytes
+    request_id: RequestId
     """Request ID for matching responses."""
 
     dest_node_id: NodeId
@@ -193,8 +193,8 @@ def __init__(
 
         self._protocol: DiscoveryProtocol | None = None
         self._transport: asyncio.DatagramTransport | None = None
-        self._pending_requests: dict[bytes, PendingRequest] = {}
-        self._pending_multi_requests: dict[bytes, PendingMultiRequest] = {}
+        self._pending_requests: dict[RequestId, PendingRequest] = {}
+        self._pending_multi_requests: dict[RequestId, PendingMultiRequest] = {}
         self._node_addresses: dict[NodeId, tuple[str, int]] = {}
 
         self._message_handler: (
@@ -384,12 +384,11 @@ async def _send_multi_response_request(
 
         # Create collector for multiple responses.
         loop = asyncio.get_running_loop()
-        request_id_bytes = bytes(message.request_id)
 
         # Use a queue to collect multiple responses.
         response_queue: asyncio.Queue[DiscoveryMessage] = asyncio.Queue()
         pending = PendingMultiRequest(
-            request_id=request_id_bytes,
+            request_id=message.request_id,
             dest_node_id=dest_node_id,
             sent_at=loop.time(),
             nonce=nonce,
@@ -398,7 +397,7 @@ async def _send_multi_response_request(
             expected_total=None,
             received_count=0,
         )
-        self._pending_multi_requests[request_id_bytes] = pending
+        self._pending_multi_requests[message.request_id] = pending
 
         # Send packet.
         self._transport.sendto(packet, dest_addr)
@@ -434,7 +433,7 @@ async def _send_multi_response_request(
                     break
 
         finally:
-            self._pending_multi_requests.pop(request_id_bytes, None)
+            self._pending_multi_requests.pop(message.request_id, None)
 
         return responses
 
@@ -494,17 +493,16 @@ async def _send_request(
         # Create pending request.
         loop = asyncio.get_running_loop()
         future: asyncio.Future[DiscoveryMessage | None] = loop.create_future()
-
-        request_id_bytes = bytes(message.request_id)
+        request_id = message.request_id
         pending = PendingRequest(
-            request_id=request_id_bytes,
+            request_id=request_id,
             dest_node_id=dest_node_id,
             sent_at=loop.time(),
             nonce=nonce,
             message=message,
             future=future,
         )
-        self._pending_requests[request_id_bytes] = pending
+        self._pending_requests[request_id] = pending
 
         # Send packet.
         self._transport.sendto(packet, dest_addr)
@@ -518,7 +516,7 @@ async def _send_request(
         except asyncio.TimeoutError:
             return None
         finally:
-            self._pending_requests.pop(request_id_bytes, None)
+            self._pending_requests.pop(request_id, None)
 
     def _build_message_packet(
         self,
@@ -790,7 +788,7 @@ async def _handle_decoded_message(
         self._session_cache.touch(remote_node_id, ip, Port(port))
 
         # Check if this is a response to a pending request.
-        request_id = bytes(message.request_id)
+        request_id = message.request_id
 
         # Check for multi-response requests first (e.g., FINDNODE -> NODES).
         multi_pending = self._pending_multi_requests.get(request_id)

src/lean_spec/subspecs/networking/transport/quic/connection.py

@@ -324,7 +324,7 @@ def __init__(self, *args, **kwargs) -> None:
         """Initialize the libp2p QUIC protocol handler."""
         super().__init__(*args, **kwargs)
         self.connection: QuicConnection | None = None
-        self.peer_identity: bytes | None = None
+        self.peer_identity: PeerId | None = None
         self.handshake_complete = asyncio.Event()
         self._buffered_events: list[QuicEvent] = []
 
@@ -532,7 +532,8 @@ async def connect(self, multiaddr: str) -> QuicConnection:
             self._context_managers.append(cm)
 
             # Cast to our protocol type to access custom attributes.
-            protocol: LibP2PQuicProtocol = base_protocol  # type: ignore[assignment]
+            assert isinstance(base_protocol, LibP2PQuicProtocol)
+            protocol: LibP2PQuicProtocol = base_protocol
 
             # Wait for handshake to complete.
             await protocol.handshake_complete.wait()

tests/lean_spec/subspecs/networking/discovery/test_transport.py

@@ -276,7 +276,7 @@ async def test_stop_cancels_pending_requests(self, local_node_id, local_private_
         loop = asyncio.get_running_loop()
         future: asyncio.Future = loop.create_future()
         pending = PendingRequest(
-            request_id=b"\x01\x02\x03\x04",
+            request_id=RequestId(data=b"\x01\x02\x03\x04"),
             dest_node_id=NodeId(bytes(32)),
             sent_at=loop.time(),
             nonce=Nonce(bytes(12)),
@@ -315,15 +315,15 @@ def test_create_pending_request(self):
         message = Ping(request_id=RequestId(data=b"\x01"), enr_seq=SeqNumber(1))
 
         pending = PendingRequest(
-            request_id=b"\x01\x02\x03\x04",
+            request_id=RequestId(data=b"\x01\x02\x03\x04"),
             dest_node_id=NodeId(bytes(32)),
             sent_at=123.456,
             nonce=Nonce(bytes(12)),
             message=message,
             future=future,
         )
 
-        assert pending.request_id == b"\x01\x02\x03\x04"
+        assert pending.request_id == RequestId(data=b"\x01\x02\x03\x04")
         assert pending.dest_node_id == bytes(32)
         assert pending.sent_at == 123.456
         assert pending.nonce == bytes(12)
@@ -407,7 +407,7 @@ def test_pending_multi_request_creation(self, local_node_id, local_private_key,
         queue: asyncio.Queue = asyncio.Queue()
 
         pending = PendingMultiRequest(
-            request_id=b"\x01\x02\x03\x04",
+            request_id=RequestId(data=b"\x01\x02\x03\x04"),
             dest_node_id=NodeId(bytes(32)),
             sent_at=123.456,
             nonce=Nonce(bytes(12)),
@@ -417,7 +417,7 @@ def test_pending_multi_request_creation(self, local_node_id, local_private_key,
             received_count=0,
         )
 
-        assert pending.request_id == b"\x01\x02\x03\x04"
+        assert pending.request_id == RequestId(data=b"\x01\x02\x03\x04")
         assert pending.expected_total is None
         assert pending.received_count == 0
 
@@ -430,7 +430,7 @@ def test_pending_multi_request_expected_total_tracking(self):
         queue: asyncio.Queue = asyncio.Queue()
 
         pending = PendingMultiRequest(
-            request_id=b"\x01\x02\x03\x04",
+            request_id=RequestId(data=b"\x01\x02\x03\x04"),
             dest_node_id=NodeId(bytes(32)),
             sent_at=123.456,
             nonce=Nonce(bytes(12)),
@@ -466,7 +466,7 @@ async def test_queue():
             queue: asyncio.Queue = asyncio.Queue()
 
             pending = PendingMultiRequest(
-                request_id=b"\x01\x02\x03\x04",
+                request_id=RequestId(data=b"\x01\x02\x03\x04"),
                 dest_node_id=NodeId(bytes(32)),
                 sent_at=123.456,
                 nonce=Nonce(bytes(12)),
@@ -572,7 +572,7 @@ def test_pending_request_stores_request_id(self):
         message = Ping(request_id=RequestId(data=b"\x01\x02\x03\x04"), enr_seq=SeqNumber(1))
 
         pending = PendingRequest(
-            request_id=b"\x01\x02\x03\x04",
+            request_id=RequestId(data=b"\x01\x02\x03\x04"),
             dest_node_id=NodeId(bytes(32)),
             sent_at=123.456,
             nonce=Nonce(bytes(12)),
@@ -581,8 +581,8 @@ def test_pending_request_stores_request_id(self):
         )
 
         # Request ID should be stored for matching.
-        assert pending.request_id == b"\x01\x02\x03\x04"
-        assert bytes(pending.message.request_id) == b"\x01\x02\x03\x04"
+        assert pending.request_id == RequestId(data=b"\x01\x02\x03\x04")
+        assert pending.message.request_id == RequestId(data=b"\x01\x02\x03\x04")
 
         loop.close()
 
@@ -596,7 +596,7 @@ async def test_future():
 
             message = Ping(request_id=RequestId(data=b"\x01"), enr_seq=SeqNumber(1))
             pending = PendingRequest(
-                request_id=b"\x01",
+                request_id=RequestId(data=b"\x01"),
                 dest_node_id=NodeId(bytes(32)),
                 sent_at=loop.time(),
                 nonce=Nonce(bytes(12)),
@@ -632,7 +632,7 @@ def test_pending_request_future_cancellation(self):
 
         message = Ping(request_id=RequestId(data=b"\x01"), enr_seq=SeqNumber(1))
         pending = PendingRequest(
-            request_id=b"\x01",
+            request_id=RequestId(data=b"\x01"),
             dest_node_id=NodeId(bytes(32)),
             sent_at=loop.time(),
             nonce=Nonce(bytes(12)),
@@ -663,7 +663,7 @@ def test_request_id_bytes_for_dict_lookup(self):
         message2 = Ping(request_id=RequestId(data=request_id_2), enr_seq=SeqNumber(2))
 
         pending1 = PendingRequest(
-            request_id=request_id_1,
+            request_id=RequestId(data=request_id_1),
             dest_node_id=NodeId(bytes(32)),
             sent_at=loop.time(),
             nonce=Nonce(bytes(12)),
@@ -672,7 +672,7 @@ def test_request_id_bytes_for_dict_lookup(self):
         )
 
         pending2 = PendingRequest(
-            request_id=request_id_2,
+            request_id=RequestId(data=request_id_2),
             dest_node_id=NodeId(bytes(32)),
             sent_at=loop.time(),
             nonce=Nonce(bytes(12)),
@@ -703,7 +703,7 @@ def test_pending_request_stores_nonce_for_whoareyou_matching(self):
 
         message = Ping(request_id=RequestId(data=b"\x01"), enr_seq=SeqNumber(1))
         pending = PendingRequest(
-            request_id=b"\x01",
+            request_id=RequestId(data=b"\x01"),
             dest_node_id=NodeId(bytes(32)),
             sent_at=loop.time(),
             nonce=nonce,
@@ -724,7 +724,7 @@ def test_pending_request_stores_message_for_retransmission(self):
 
         message = Ping(request_id=RequestId(data=b"\x01"), enr_seq=SeqNumber(42))
         pending = PendingRequest(
-            request_id=b"\x01",
+            request_id=RequestId(data=b"\x01"),
             dest_node_id=NodeId(bytes(32)),
             sent_at=loop.time(),
             nonce=Nonce(bytes(12)),
@@ -782,7 +782,7 @@ async def test_pending_requests_cleared_on_stop(
         for i in range(3):
             future: asyncio.Future = loop.create_future()
             pending = PendingRequest(
-                request_id=bytes([i]),
+                request_id=RequestId(data=bytes([i])),
                 dest_node_id=NodeId(bytes(32)),
                 sent_at=loop.time(),
                 nonce=Nonce(bytes(12)),
@@ -825,7 +825,7 @@ async def test_pending_request_futures_cancelled_on_stop(
             future: asyncio.Future = loop.create_future()
             futures.append(future)
             pending = PendingRequest(
-                request_id=bytes([i]),
+                request_id=RequestId(data=bytes([i])),
                 dest_node_id=NodeId(bytes(32)),
                 sent_at=loop.time(),
                 nonce=Nonce(bytes(12)),
@@ -1086,7 +1086,7 @@ async def test_response_completes_pending_request_future(
 
         loop = asyncio.get_running_loop()
         future: asyncio.Future[Pong | None] = loop.create_future()
-        request_id = b"\x01\x02\x03\x04"
+        request_id = RequestId(data=b"\x01\x02\x03\x04")
 
         pending = PendingRequest(
             request_id=request_id,
@@ -1099,7 +1099,7 @@ async def test_response_completes_pending_request_future(
         transport._pending_requests[request_id] = pending
 
         pong = Pong(
-            request_id=RequestId(data=request_id),
+            request_id=request_id,
             enr_seq=SeqNumber(1),
             recipient_ip=IPv4(b"\x7f\x00\x00\x01"),
             recipient_port=Port(9000),
@@ -1121,7 +1121,7 @@ async def test_response_enqueued_for_multi_request(
             local_enr=local_enr,
         )
 
-        request_id = b"\x01\x02\x03\x04"
+        request_id = RequestId(data=b"\x01\x02\x03\x04")
         queue: asyncio.Queue = asyncio.Queue()
 
         multi_pending = PendingMultiRequest(
@@ -1137,7 +1137,7 @@ async def test_response_enqueued_for_multi_request(
         transport._pending_multi_requests[request_id] = multi_pending
 
         nodes = Nodes(
-            request_id=RequestId(data=request_id),
+            request_id=request_id,
             total=Uint8(1),
             enrs=[b"enr1"],
         )