Hint GetBeaconBlock message type with TypedDict

Author: jannikluhn

trinity/protocol/bcc/commands.py

@@ -1,10 +1,25 @@
+from typing import (
+    Union,
+)
+
 from rlp import sedes
 
+from mypy_extensions import (
+    TypedDict,
+)
+
+from eth_typing import (
+    Hash32,
+)
+
 from p2p.protocol import (
     Command,
 )
 
-from trinity.rlp.sedes import HashOrNumber
+from trinity.rlp.sedes import (
+    HashOrNumber,
+)
+
 from eth.beacon.types.blocks import BaseBeaconBlock
 from eth.beacon.types.attestation_records import AttestationRecord
 
@@ -19,6 +34,12 @@ class Status(Command):
     ]
 
 
+GetBeaconBlocksMessage = TypedDict("GetBeaconBlocksMessage", {
+    "block_slot_or_hash": Union[int, Hash32],
+    "max_blocks": int,
+})
+
+
 class GetBeaconBlocks(Command):
     _cmd_id = 1
     structure = [

trinity/protocol/bcc/servers.py

@@ -1,8 +1,6 @@
 import itertools
 from typing import (
     cast,
-    Any,
-    Dict,
     Iterable,
     Set,
     Type,
@@ -33,7 +31,10 @@
 from eth.beacon.types.blocks import BaseBeaconBlock
 
 from trinity.protocol.common.servers import BaseRequestServer
-from trinity.protocol.bcc import commands
+from trinity.protocol.bcc.commands import (
+    GetBeaconBlocks,
+    GetBeaconBlocksMessage,
+)
 from trinity.protocol.bcc.peer import (
     BCCPeer,
     BCCPeerPool,
@@ -42,7 +43,7 @@
 
 class BCCRequestServer(BaseRequestServer):
     subscription_msg_types: Set[Type[Command]] = {
-        commands.GetBeaconBlocks,
+        GetBeaconBlocks,
     }
 
     def __init__(self,
@@ -56,16 +57,16 @@ async def _handle_msg(self, base_peer: BasePeer, cmd: Command,
                           msg: protocol._DecodedMsgType) -> None:
         peer = cast(BCCPeer, base_peer)
 
-        if isinstance(cmd, commands.GetBeaconBlocks):
-            await self._handle_get_beacon_blocks(peer, cast(Dict[str, Any], msg))
+        if isinstance(cmd, GetBeaconBlocks):
+            await self._handle_get_beacon_blocks(peer, cast(GetBeaconBlocksMessage, msg))
         else:
             raise Exception("Invariant: Only subscribed to GetBeaconBlocks")
 
-    async def _handle_get_beacon_blocks(self, peer: BCCPeer, msg: Dict[str, Any]) -> None:
+    async def _handle_get_beacon_blocks(self, peer: BCCPeer, msg: GetBeaconBlocksMessage) -> None:
         if not peer.is_operational:
             return
 
-        max_blocks = cast(int, msg["max_blocks"])
+        max_blocks = msg["max_blocks"]
         block_slot_or_hash = msg["block_slot_or_hash"]
 
         try: