Pass block_class explictly in BeaconChainDB APIs

Author: hwwhww

eth/beacon/db/chain.py

@@ -53,38 +53,40 @@
 
 class BaseBeaconChainDB(ABC):
     db = None  # type: BaseAtomicDB
-    block_class = None  # type: Type[BaseBeaconBlock]
-
-    @abstractmethod
-    def set_block_class(self, block_class: Type[BaseBeaconBlock]) -> None:
-        pass
 
     #
     # Block API
     #
     @abstractmethod
-    def persist_block(self,
-                      block: BaseBeaconBlock) -> Tuple[Tuple[bytes, ...], Tuple[bytes, ...]]:
+    def persist_block(
+            self,
+            block: BaseBeaconBlock,
+            block_class: Type[BaseBeaconBlock]
+    ) -> Tuple[Tuple[bytes, ...], Tuple[bytes, ...]]:
         pass
 
     @abstractmethod
     def get_canonical_block_root(self, slot: int) -> Hash32:
         pass
 
     @abstractmethod
-    def get_canonical_block_by_slot(self, slot: int) -> BaseBeaconBlock:
+    def get_canonical_block_by_slot(self,
+                                    slot: int,
+                                    block_class: Type[BaseBeaconBlock]) -> BaseBeaconBlock:
         pass
 
     @abstractmethod
     def get_canonical_block_root_by_slot(self, slot: int) -> Hash32:
         pass
 
     @abstractmethod
-    def get_canonical_head(self) -> BaseBeaconBlock:
+    def get_canonical_head(self, block_class: Type[BaseBeaconBlock]) -> BaseBeaconBlock:
         pass
 
     @abstractmethod
-    def get_block_by_root(self, block_root: Hash32) -> BaseBeaconBlock:
+    def get_block_by_root(self,
+                          block_root: Hash32,
+                          block_class: Type[BaseBeaconBlock]) -> BaseBeaconBlock:
         pass
 
     @abstractmethod
@@ -97,8 +99,9 @@ def block_exists(self, block_root: Hash32) -> bool:
 
     @abstractmethod
     def persist_block_chain(
-        self,
-        blocks: Iterable[BaseBeaconBlock]
+            self,
+            blocks: Iterable[BaseBeaconBlock],
+            block_class: Type[BaseBeaconBlock]
     ) -> Tuple[Tuple[BaseBeaconBlock, ...], Tuple[BaseBeaconBlock, ...]]:
         pass
 
@@ -127,20 +130,19 @@ def get(self, key: bytes) -> bytes:
 
 
 class BeaconChainDB(BaseBeaconChainDB):
-    def __init__(self, db: BaseAtomicDB, block_class: Type[BaseBeaconBlock]) -> None:
+    def __init__(self, db: BaseAtomicDB) -> None:
         self.db = db
-        self.block_class = block_class
-
-    def set_block_class(self, block_class: Type[BaseBeaconBlock]) -> None:
-        self.block_class = block_class
 
-    def persist_block(self,
-                      block: BaseBeaconBlock) -> Tuple[Tuple[bytes, ...], Tuple[bytes, ...]]:
+    def persist_block(
+            self,
+            block: BaseBeaconBlock,
+            block_class: Type[BaseBeaconBlock]
+    ) -> Tuple[Tuple[bytes, ...], Tuple[bytes, ...]]:
         """
         Persist the given block.
         """
         with self.db.atomic_batch() as db:
-            return self._persist_block(db, block, self.block_class)
+            return self._persist_block(db, block, block_class)
 
     @classmethod
     def _persist_block(
@@ -188,14 +190,16 @@ def _get_canonical_block_root(db: BaseDB, slot: int) -> Hash32:
         else:
             return rlp.decode(encoded_key, sedes=rlp.sedes.binary)
 
-    def get_canonical_block_by_slot(self, slot: int) -> BaseBeaconBlock:
+    def get_canonical_block_by_slot(self,
+                                    slot: int,
+                                    block_class: Type[BaseBeaconBlock]) -> BaseBeaconBlock:
         """
         Return the block with the given slot in the canonical chain.
 
         Raise BlockNotFound if there's no block with the given slot in the
         canonical chain.
         """
-        return self._get_canonical_block_by_slot(self.db, slot, self.block_class)
+        return self._get_canonical_block_by_slot(self.db, slot, block_class)
 
     @classmethod
     def _get_canonical_block_by_slot(
@@ -223,11 +227,11 @@ def _get_canonical_block_root_by_slot(
         validate_slot(slot)
         return cls._get_canonical_block_root(db, slot)
 
-    def get_canonical_head(self) -> BaseBeaconBlock:
+    def get_canonical_head(self, block_class: Type[BaseBeaconBlock]) -> BaseBeaconBlock:
         """
         Return the current block at the head of the chain.
         """
-        return self._get_canonical_head(self.db, self.block_class)
+        return self._get_canonical_head(self.db, block_class)
 
     @classmethod
     def _get_canonical_head(cls,
@@ -239,8 +243,10 @@ def _get_canonical_head(cls,
             raise CanonicalHeadNotFound("No canonical head set for this chain")
         return cls._get_block_by_root(db, Hash32(canonical_head_root), block_class)
 
-    def get_block_by_root(self, block_root: Hash32) -> BaseBeaconBlock:
-        return self._get_block_by_root(self.db, block_root, self.block_class)
+    def get_block_by_root(self,
+                          block_root: Hash32,
+                          block_class: Type[BaseBeaconBlock]) -> BaseBeaconBlock:
+        return self._get_block_by_root(self.db, block_root, block_class)
 
     @staticmethod
     def _get_block_by_root(db: BaseDB,
@@ -280,15 +286,16 @@ def _block_exists(db: BaseDB, block_root: Hash32) -> bool:
         return block_root in db
 
     def persist_block_chain(
-        self,
-        blocks: Iterable[BaseBeaconBlock]
+            self,
+            blocks: Iterable[BaseBeaconBlock],
+            block_class: Type[BaseBeaconBlock]
     ) -> Tuple[Tuple[BaseBeaconBlock, ...], Tuple[BaseBeaconBlock, ...]]:
         """
         Return two iterable of blocks, the first containing the new canonical blocks,
         the second containing the old canonical headers
         """
         with self.db.atomic_batch() as db:
-            return self._persist_block_chain(db, blocks, self.block_class)
+            return self._persist_block_chain(db, blocks, block_class)
 
     @classmethod
     def _set_block_scores_to_db(

eth/beacon/state_machines/base.py

@@ -71,10 +71,10 @@ def import_block(self, block: BaseBeaconBlock) -> Tuple[BeaconState, BaseBeaconB
 class BeaconStateMachine(BaseBeaconStateMachine):
     def __init__(self,
                  chaindb: BaseBeaconChainDB,
-                 block_root: Hash32) -> None:
+                 block_root: Hash32,
+                 block_class: Type[BaseBeaconBlock]) -> None:
         self.chaindb = chaindb
-        self.chaindb.set_block_class(self.get_block_class())
-        self.block = self.chaindb.get_block_by_root(block_root)
+        self.block = self.chaindb.get_block_by_root(block_root, block_class)
 
     @property
     def state(self) -> BeaconState:

eth/beacon/types/blocks.py

@@ -173,7 +173,7 @@ def from_root(cls, root: Hash32, chaindb: 'BaseBeaconChainDB') -> 'BeaconBlock':
         """
         Returns the block denoted by the given block header.
         """
-        block = chaindb.get_block_by_root(root)
+        block = chaindb.get_block_by_root(root, cls)
         body = cls.block_body_class(
             proposer_slashings=block.body.proposer_slashings,
             casper_slashings=block.body.casper_slashings,

tests/beacon/db/test_beacon_chaindb.py

@@ -33,7 +33,7 @@
 
 @pytest.fixture
 def chaindb(base_db):
-    return BeaconChainDB(base_db, BeaconBlock)
+    return BeaconChainDB(base_db)
 
 
 @pytest.fixture(params=[0, 10, 999])
@@ -52,33 +52,33 @@ def state(sample_beacon_state_params):
 def test_chaindb_add_block_number_to_root_lookup(chaindb, block):
     block_slot_to_root_key = SchemaV1.make_block_slot_to_root_lookup_key(block.slot)
     assert not chaindb.exists(block_slot_to_root_key)
-    chaindb.persist_block(block)
+    chaindb.persist_block(block, block.__class__)
     assert chaindb.exists(block_slot_to_root_key)
 
 
 def test_chaindb_persist_block_and_slot_to_root(chaindb, block):
     with pytest.raises(BlockNotFound):
-        chaindb.get_block_by_root(block.root)
+        chaindb.get_block_by_root(block.root, block.__class__)
     slot_to_root_key = SchemaV1.make_block_root_to_score_lookup_key(block.root)
     assert not chaindb.exists(slot_to_root_key)
 
-    chaindb.persist_block(block)
+    chaindb.persist_block(block, block.__class__)
 
-    assert chaindb.get_block_by_root(block.root) == block
+    assert chaindb.get_block_by_root(block.root, block.__class__) == block
     assert chaindb.exists(slot_to_root_key)
 
 
 @given(seed=st.binary(min_size=32, max_size=32))
 def test_chaindb_persist_block_and_unknown_parent(chaindb, block, seed):
     n_block = block.copy(parent_root=hash_eth2(seed))
     with pytest.raises(ParentNotFound):
-        chaindb.persist_block(n_block)
+        chaindb.persist_block(n_block, n_block.__class__)
 
 
 def test_chaindb_persist_block_and_block_to_root(chaindb, block):
     block_to_root_key = SchemaV1.make_block_root_to_score_lookup_key(block.root)
     assert not chaindb.exists(block_to_root_key)
-    chaindb.persist_block(block)
+    chaindb.persist_block(block, block.__class__)
     assert chaindb.exists(block_to_root_key)
 
 
@@ -87,7 +87,7 @@ def test_chaindb_get_score(chaindb, sample_beacon_block_params):
         parent_root=GENESIS_PARENT_HASH,
         slot=0,
     )
-    chaindb.persist_block(genesis)
+    chaindb.persist_block(genesis, genesis.__class__)
 
     genesis_score_key = SchemaV1.make_block_root_to_score_lookup_key(genesis.root)
     genesis_score = rlp.decode(chaindb.db.get(genesis_score_key), sedes=rlp.sedes.big_endian_int)
@@ -98,7 +98,7 @@ def test_chaindb_get_score(chaindb, sample_beacon_block_params):
         parent_root=genesis.root,
         slot=1,
     )
-    chaindb.persist_block(block1)
+    chaindb.persist_block(block1, block1.__class__)
 
     block1_score_key = SchemaV1.make_block_root_to_score_lookup_key(block1.root)
     block1_score = rlp.decode(chaindb.db.get(block1_score_key), sedes=rlp.sedes.big_endian_int)
@@ -107,13 +107,13 @@ def test_chaindb_get_score(chaindb, sample_beacon_block_params):
 
 
 def test_chaindb_get_block_by_root(chaindb, block):
-    chaindb.persist_block(block)
-    result_block = chaindb.get_block_by_root(block.root)
+    chaindb.persist_block(block, block.__class__)
+    result_block = chaindb.get_block_by_root(block.root, block.__class__)
     validate_rlp_equal(result_block, block)
 
 
 def test_chaindb_get_canonical_block_root(chaindb, block):
-    chaindb.persist_block(block)
+    chaindb.persist_block(block, block.__class__)
     block_root = chaindb.get_canonical_block_root(block.slot)
     assert block_root == block.root
 

tests/beacon/state_machines/test_demo.py

@@ -36,7 +36,7 @@ def test_demo(base_db,
               config,
               privkeys,
               pubkeys):
-    chaindb = BeaconChainDB(base_db, SerenityBeaconBlock)
+    chaindb = BeaconChainDB(base_db)
     state = genesis_state
     block = SerenityBeaconBlock(**sample_beacon_block_params).copy(
         slot=state.slot + 2,
@@ -68,11 +68,11 @@ def test_demo(base_db,
     )
 
     # Store in chaindb
-    chaindb.persist_block(block)
+    chaindb.persist_block(block, SerenityBeaconBlock)
     chaindb.persist_state(state)
 
     # Get state machine instance
-    sm = fixture_sm_class(chaindb, block.root)
+    sm = fixture_sm_class(chaindb, block.root, SerenityBeaconBlock)
     result_state, _ = sm.import_block(block)
 
     assert state.slot == 0