Add type checking

Author: richardkiss

chia/consensus/consensus_store_protocol.py

@@ -2,9 +2,8 @@
 
 from collections.abc import Collection
 from contextlib import AbstractAsyncContextManager
-from typing import Optional, Protocol
+from typing import Any, Optional, Protocol
 
-import aiosqlite
 from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment, SubEpochSummary
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import uint32, uint64
@@ -19,8 +18,7 @@ class ConsensusStoreProtocol(Protocol):
     the consensus-related data in the blockchain.
     """
 
-    # Block store methods
-    def transaction(self) -> AbstractAsyncContextManager[aiosqlite.Connection]: ...
+    def transaction(self) -> AbstractAsyncContextManager[Any]: ...
 
     async def get_block_records_close_to_peak(
         self, blocks_n: int
@@ -46,7 +44,6 @@ async def get_sub_epoch_challenge_segments(
     async def get_generator(self, header_hash: bytes32) -> Optional[bytes]: ...
     async def get_generators_at(self, heights: set[uint32]) -> dict[uint32, bytes]: ...
 
-    # Coin store methods
     async def get_coin_records(self, names: Collection[bytes32]) -> list[CoinRecord]: ...
     async def rollback_to_block(self, block_index: int) -> dict[bytes32, CoinRecord]: ...
     async def new_block(
@@ -60,7 +57,6 @@ async def new_block(
     async def get_coins_added_at_height(self, height: uint32) -> list[CoinRecord]: ...
     async def get_coins_removed_at_height(self, height: uint32) -> list[CoinRecord]: ...
 
-    # Height map methods
     def get_ses_heights(self) -> list[uint32]: ...
     def get_ses(self, height: uint32) -> SubEpochSummary: ...
     def contains_height(self, height: uint32) -> bool: ...

chia/full_node/consensus_store_sqlite3.py

@@ -3,25 +3,23 @@
 import dataclasses
 from collections.abc import Collection
 from contextlib import AbstractAsyncContextManager
-from pathlib import Path
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 import aiosqlite
 from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment, SubEpochSummary
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import uint32, uint64
+from typing_extensions import Self
 
 from chia.consensus.block_height_map import BlockHeightMap
-from chia.consensus.consensus_store_protocol import ConsensusStoreProtocol
 from chia.full_node.block_store import BlockStore
 from chia.full_node.coin_store import CoinStore
 from chia.types.blockchain_format.coin import Coin
 from chia.types.coin_record import CoinRecord
-from chia.util.db_wrapper import DBWrapper2
 
 
 @dataclasses.dataclass
-class ConsensusStoreSQLite3(ConsensusStoreProtocol):
+class ConsensusStoreSQLite3:
     """
     Consensus store that combines block_store, coin_store, and height_map functionality.
     """
@@ -36,7 +34,7 @@ async def create(
         block_store: BlockStore,
         coin_store: CoinStore,
         height_map: BlockHeightMap,
-    ) -> ConsensusStoreSQLite3:
+    ) -> Self:
         """Create a new ConsensusStore instance from existing sub-stores.
 
         This factory does not create sub-stores. Construct BlockStore, CoinStore,
@@ -158,23 +156,12 @@ async def maybe_flush_height_map(self) -> None:
         await self.height_map.maybe_flush()
 
 
-async def build_consensus_store(
-    db_wrapper: DBWrapper2,
-    blockchain_dir: Path,
-    selected_network: Optional[str] = None,
-    *,
-    use_cache: bool = True,
-) -> ConsensusStoreSQLite3:
-    """Convenience utility to construct a ConsensusStore from DB and path.
+if TYPE_CHECKING:
+    # here we validate that ConsensusStoreSQLite3 implements ConsensusStoreProtocol
+    from typing import cast
 
-    This keeps ConsensusStore.create minimal (only accepts sub-stores) while still
-    providing a one-stop helper for callers that need to create the sub-stores.
-    """
-    block_store = await BlockStore.create(db_wrapper, use_cache=use_cache)
-    coin_store = await CoinStore.create(db_wrapper)
-    height_map = await BlockHeightMap.create(blockchain_dir, db_wrapper, selected_network)
-    return await ConsensusStoreSQLite3.create(block_store, coin_store, height_map)
+    from chia.consensus.consensus_store_protocol import ConsensusStoreProtocol
 
+    def _protocol_check(_: ConsensusStoreProtocol): ...
 
-# Backwards-compatible alias
-ConsensusStore = ConsensusStoreSQLite3
+    _protocol_check(cast(ConsensusStoreSQLite3, None))