Improve typing (#79)

src/providers/beacon_node.py

@@ -526,14 +526,18 @@ async def publish_attestations(
             headers={"Eth-Consensus-Version": fork_version.value},
         )
 
-    async def prepare_beacon_committee_subscriptions(self, data: list[dict]) -> None:  # type: ignore[type-arg]
+    async def prepare_beacon_committee_subscriptions(
+        self, data: list[SchemaBeaconAPI.SubscribeToBeaconCommitteeSubnetRequestBody]
+    ) -> None:
         await self._make_request(
             method="POST",
             endpoint="/eth/v1/validator/beacon_committee_subscriptions",
             data=self.json_encoder.encode(data),
         )
 
-    async def prepare_sync_committee_subscriptions(self, data: list[dict]) -> None:  # type: ignore[type-arg]
+    async def prepare_sync_committee_subscriptions(
+        self, data: list[SchemaBeaconAPI.SubscribeToSyncCommitteeSubnetRequestBody]
+    ) -> None:
         await self._make_request(
             method="POST",
             endpoint="/eth/v1/validator/sync_committee_subscriptions",

src/schemas/beacon_api.py

@@ -56,6 +56,40 @@ class ForkVersion(Enum):
     ELECTRA = "electra"
 
 
+class AttestationPhase0(msgspec.Struct):
+    aggregation_bits: str
+    data: dict  # type: ignore[type-arg]
+    signature: str
+
+
+class SingleAttestation(msgspec.Struct):
+    committee_index: str
+    attester_index: str
+    data: dict  # type: ignore[type-arg]
+    signature: str
+
+
+class SubscribeToBeaconCommitteeSubnetRequestBody(msgspec.Struct):
+    validator_index: str
+    committee_index: str
+    committees_at_slot: str
+    slot: str
+    is_aggregator: bool
+
+
+class SyncCommitteeSignature(msgspec.Struct):
+    slot: str
+    beacon_block_root: str
+    validator_index: str
+    signature: str
+
+
+class SubscribeToSyncCommitteeSubnetRequestBody(msgspec.Struct):
+    validator_index: str
+    sync_committee_indices: list[str]
+    until_epoch: str
+
+
 class GetAggregatedAttestationV2Response(msgspec.Struct):
     version: ForkVersion
     data: dict  # type: ignore[type-arg]

src/services/attestation.py

@@ -228,7 +228,9 @@ async def attest_if_not_yet_attested(
             )
 
             # Sign the attestation data
-            attestations_objects_to_publish: list[dict] = []  # type: ignore[type-arg]
+            attestations_objects_to_publish: list[
+                SchemaBeaconAPI.AttestationPhase0 | SchemaBeaconAPI.SingleAttestation
+            ] = []
 
             def _att_data_for_committee_idx(
                 _orig_att_data_obj: dict,  # type: ignore[type-arg]
@@ -288,7 +290,7 @@ def _att_data_for_committee_idx(
                         aggregation_bits[int(duty.validator_committee_index)] = True
 
                         attestations_objects_to_publish.append(
-                            dict(
+                            SchemaBeaconAPI.AttestationPhase0(
                                 aggregation_bits=aggregation_bits.to_obj(),
                                 data=_att_data_for_committee_idx(
                                     att_data_obj,
@@ -300,7 +302,7 @@ def _att_data_for_committee_idx(
                     elif _fork_version == SchemaBeaconAPI.ForkVersion.ELECTRA:
                         # SingleAttestation object from the CL spec
                         attestations_objects_to_publish.append(
-                            dict(
+                            SchemaBeaconAPI.SingleAttestation(
                                 committee_index=duty.committee_index,
                                 attester_index=duty.validator_index,
                                 data=att_data_obj,
@@ -573,7 +575,7 @@ async def _prep_and_schedule_duties(
 
         # Prepare beacon node subnet subscriptions for aggregation duties
         beacon_committee_subscriptions_data = [
-            dict(
+            SchemaBeaconAPI.SubscribeToBeaconCommitteeSubnetRequestBody(
                 validator_index=duty.validator_index,
                 committee_index=duty.committee_index,
                 committees_at_slot=duty.committees_at_slot,

src/services/sync_committee.py

@@ -181,7 +181,7 @@ async def produce_sync_message_if_not_yet_produced(
             for validator in sync_committee_members
         ]
 
-        sync_messages_to_publish = []
+        sync_messages_to_publish: list[SchemaBeaconAPI.SyncCommitteeSignature] = []
         for coro in asyncio.as_completed(coroutines):
             try:
                 msg, sig, pubkey = await coro
@@ -194,7 +194,7 @@ async def produce_sync_message_if_not_yet_produced(
                 continue
 
             sync_messages_to_publish.append(
-                dict(
+                SchemaBeaconAPI.SyncCommitteeSignature(
                     beacon_block_root=msg.sync_committee_message.beacon_block_root,
                     slot=str(msg.sync_committee_message.slot),
                     validator_index=next(
@@ -527,7 +527,7 @@ async def _update_duties(self) -> None:
                 sync_period + 1
             ) * self.beacon_chain.spec.EPOCHS_PER_SYNC_COMMITTEE_PERIOD
             sync_committee_subscriptions_data = [
-                dict(
+                SchemaBeaconAPI.SubscribeToSyncCommitteeSubnetRequestBody(
                     validator_index=duty.validator_index,
                     sync_committee_indices=duty.validator_sync_committee_indices,
                     until_epoch=str(until_epoch),

tests/conftest.py

@@ -224,16 +224,14 @@ async def multi_beacon_node(
 @pytest.fixture(scope="session")
 def genesis(spec: SpecElectra) -> Genesis:
     # Fake genesis 1 hour ago
-    return Genesis.from_obj(  # type: ignore[no-any-return]
-        dict(
-            genesis_time=int(
-                (
-                    datetime.datetime.now(tz=datetime.UTC) - datetime.timedelta(hours=1)
-                ).timestamp()
-            ),
-            genesis_validators_root="0x9143aa7c615a7f7115e2b6aac319c03529df8242ae705fba9df39b79c59fa8b1",
-            genesis_fork_version=spec.GENESIS_FORK_VERSION,
+    return Genesis(
+        genesis_time=int(
+            (
+                datetime.datetime.now(tz=datetime.UTC) - datetime.timedelta(hours=1)
+            ).timestamp()
         ),
+        genesis_validators_root="0x9143aa7c615a7f7115e2b6aac319c03529df8242ae705fba9df39b79c59fa8b1",
+        genesis_fork_version=spec.GENESIS_FORK_VERSION,
     )
 
 

tests/mock_api/beacon_node.py

@@ -334,26 +334,34 @@ def _mocked_beacon_api_endpoints_post(url: URL, **kwargs: Any) -> CallbackResult
             )
 
         if re.match("/eth/v1/validator/beacon_committee_subscriptions", url.raw_path):
+            _ = msgspec.json.decode(
+                kwargs["data"].decode(),
+                type=list[SchemaBeaconAPI.SubscribeToBeaconCommitteeSubnetRequestBody],
+            )
             return CallbackResult(status=200)
 
         if re.match("/eth/v2/beacon/pool/attestations", url.raw_path):
-            data_list = msgspec.json.decode(kwargs["data"])
-            assert len(data_list) == 1
-            data = data_list[0]
-            assert (
-                data["data"]["beacon_block_root"]
-                == "0x9f19cc6499596bdf19be76d80b878ee3326e68cf2ed69cbada9a1f4fe13c51b3"
-            )
-
             if beacon_chain.current_fork_version == ForkVersion.ELECTRA:
-                assert "committee_index" in data
-                assert "attester_index" in data
+                attestations = msgspec.json.decode(
+                    kwargs["data"].decode(),
+                    type=list[SchemaBeaconAPI.SingleAttestation],
+                )
+                attestation = attestations[0]
             elif beacon_chain.current_fork_version == ForkVersion.DENEB:
-                assert data["aggregation_bits"] == "0x000201"
-                assert "committee_bits" not in data
+                attestations = msgspec.json.decode(
+                    kwargs["data"].decode(),
+                    type=list[SchemaBeaconAPI.AttestationPhase0],
+                )
+                attestation = attestations[0]
+                assert attestation.aggregation_bits == "0x000201"  # type: ignore[attr-defined]
             else:
                 raise ValueError(f"Unsupported spec: {spec}")
 
+            assert (
+                attestation.data["beacon_block_root"]
+                == "0x9f19cc6499596bdf19be76d80b878ee3326e68cf2ed69cbada9a1f4fe13c51b3"
+            )
+
             return CallbackResult(status=200)
 
         if re.match("/eth/v2/validator/aggregate_and_proofs", url.raw_path):