refactor: remove VertexStorageProtocol

Author: glevco

hathor/consensus/block_consensus.py

@@ -432,7 +432,7 @@ def remove_first_block_markers(self, block: Block) -> None:
         storage = block.storage
 
         from hathor.transaction.storage.traversal import BFSTimestampWalk
-        bfs = BFSTimestampWalk(storage, is_dag_verifications=True, is_left_to_right=False)
+        bfs = BFSTimestampWalk(storage.get_vertex, is_dag_verifications=True, is_left_to_right=False)
         for tx in bfs.run(block, skip_root=True):
             if tx.is_block:
                 bfs.skip_neighbors(tx)
@@ -469,7 +469,7 @@ def _score_block_dfs(self, block: BaseTransaction, used: set[bytes],
 
             else:
                 from hathor.transaction.storage.traversal import BFSTimestampWalk
-                bfs = BFSTimestampWalk(storage, is_dag_verifications=True, is_left_to_right=False)
+                bfs = BFSTimestampWalk(storage.get_vertex, is_dag_verifications=True, is_left_to_right=False)
                 for tx in bfs.run(parent, skip_root=False):
                     assert not tx.is_block
 

hathor/consensus/transaction_consensus.py

@@ -344,7 +344,9 @@ def remove_voided_by(self, tx: Transaction, voided_hash: bytes) -> bool:
 
         self.log.debug('remove_voided_by', tx=tx.hash_hex, voided_hash=voided_hash.hex())
 
-        bfs = BFSTimestampWalk(tx.storage, is_dag_funds=True, is_dag_verifications=True, is_left_to_right=True)
+        bfs = BFSTimestampWalk(
+            tx.storage.get_vertex, is_dag_funds=True, is_dag_verifications=True, is_left_to_right=True
+        )
         check_list: list[BaseTransaction] = []
         for tx2 in bfs.run(tx, skip_root=False):
             assert tx2.storage is not None
@@ -400,7 +402,7 @@ def add_voided_by(self, tx: Transaction, voided_hash: bytes) -> bool:
             is_dag_verifications = False
 
         from hathor.transaction.storage.traversal import BFSTimestampWalk
-        bfs = BFSTimestampWalk(tx.storage, is_dag_funds=True, is_dag_verifications=is_dag_verifications,
+        bfs = BFSTimestampWalk(tx.storage.get_vertex, is_dag_funds=True, is_dag_verifications=is_dag_verifications,
                                is_left_to_right=True)
         check_list: list[Transaction] = []
         for tx2 in bfs.run(tx, skip_root=False):

hathor/indexes/mempool_tips_index.py

@@ -209,7 +209,7 @@ def iter(self, tx_storage: 'TransactionStorage', max_timestamp: Optional[float]
 
     def iter_all(self, tx_storage: 'TransactionStorage') -> Iterator[Transaction]:
         from hathor.transaction.storage.traversal import BFSTimestampWalk
-        bfs = BFSTimestampWalk(tx_storage, is_dag_verifications=True, is_left_to_right=False)
+        bfs = BFSTimestampWalk(tx_storage.get_vertex, is_dag_verifications=True, is_left_to_right=False)
         for tx in bfs.run(self.iter(tx_storage), skip_root=False):
             assert isinstance(tx, Transaction)
             if tx.get_metadata().first_block is not None:

hathor/p2p/sync_v2/streamers.py

@@ -232,7 +232,9 @@ def __init__(self,
             assert tx.get_metadata().first_block == self.first_block.hash
 
         self.current_block: Optional[Block] = self.first_block
-        self.bfs = BFSOrderWalk(self.tx_storage, is_dag_verifications=True, is_dag_funds=True, is_left_to_right=False)
+        self.bfs = BFSOrderWalk(
+            self.tx_storage.get_vertex, is_dag_verifications=True, is_dag_funds=True, is_left_to_right=False
+        )
         self.iter = self.get_iter()
 
     def _stop_streaming_server(self, response_code: StreamEnd) -> None:

hathor/reward_lock/reward_lock.py

@@ -12,18 +12,20 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from __future__ import annotations
+
 from typing import TYPE_CHECKING, Iterator, Optional
 
 from hathor.conf.settings import HathorSettings
 from hathor.transaction import Block
 from hathor.util import not_none
 
 if TYPE_CHECKING:
-    from hathor.transaction.storage.vertex_storage_protocol import VertexStorageProtocol
+    from hathor.transaction.storage import TransactionStorage
     from hathor.transaction.transaction import RewardLockedInfo, Transaction
 
 
-def iter_spent_rewards(tx: 'Transaction', storage: 'VertexStorageProtocol') -> Iterator[Block]:
+def iter_spent_rewards(tx: 'Transaction', storage: TransactionStorage) -> Iterator[Block]:
     """Iterate over all the rewards being spent, assumes tx has been verified."""
     for input_tx in tx.inputs:
         spent_tx = storage.get_vertex(input_tx.tx_id)
@@ -41,7 +43,7 @@ def is_spent_reward_locked(settings: HathorSettings, tx: 'Transaction') -> bool:
 def get_spent_reward_locked_info(
     settings: HathorSettings,
     tx: 'Transaction',
-    storage: 'VertexStorageProtocol',
+    storage: TransactionStorage,
 ) -> Optional['RewardLockedInfo']:
     """Check if any input block reward is locked, returning the locked information if any, or None if they are all
     unlocked."""
@@ -54,7 +56,7 @@ def get_spent_reward_locked_info(
     return None
 
 
-def get_minimum_best_height(storage: 'VertexStorageProtocol') -> int:
+def get_minimum_best_height(storage: TransactionStorage) -> int:
     """Return the height of the current best block that shall be used for `min_height` verification."""
     import math
 

hathor/transaction/base_transaction.py

@@ -672,7 +672,9 @@ def update_accumulated_weight(self, *, stop_value: float = inf, save_file: bool
         # directly verified by a block.
 
         from hathor.transaction.storage.traversal import BFSTimestampWalk
-        bfs_walk = BFSTimestampWalk(self.storage, is_dag_funds=True, is_dag_verifications=True, is_left_to_right=True)
+        bfs_walk = BFSTimestampWalk(
+            self.storage.get_vertex, is_dag_funds=True, is_dag_verifications=True, is_left_to_right=True
+        )
         for tx in bfs_walk.run(self, skip_root=True):
             accumulated_weight += weight_to_work(tx.weight)
             if accumulated_weight > stop_value:

hathor/transaction/block.py

@@ -360,7 +360,9 @@ def iter_transactions_in_this_block(self) -> Iterator[BaseTransaction]:
         """Return an iterator of the transactions that have this block as meta.first_block."""
         from hathor.transaction.storage.traversal import BFSOrderWalk
         assert self.storage is not None
-        bfs = BFSOrderWalk(self.storage, is_dag_verifications=True, is_dag_funds=True, is_left_to_right=False)
+        bfs = BFSOrderWalk(
+            self.storage.get_vertex, is_dag_verifications=True, is_dag_funds=True, is_left_to_right=False
+        )
         for tx in bfs.run(self, skip_root=True):
             tx_meta = tx.get_metadata()
             if tx_meta.first_block != self.hash:

hathor/transaction/storage/__init__.py

@@ -15,7 +15,6 @@
 from hathor.transaction.storage.cache_storage import TransactionCacheStorage
 from hathor.transaction.storage.memory_storage import TransactionMemoryStorage
 from hathor.transaction.storage.transaction_storage import TransactionStorage
-from hathor.transaction.storage.vertex_storage_protocol import VertexStorageProtocol
 
 try:
     from hathor.transaction.storage.rocksdb_storage import TransactionRocksDBStorage
@@ -27,5 +26,4 @@
     'TransactionMemoryStorage',
     'TransactionCacheStorage',
     'TransactionRocksDBStorage',
-    'VertexStorageProtocol'
 ]

hathor/transaction/storage/transaction_storage.py

@@ -1012,7 +1012,7 @@ def iter_mempool_from_tx_tips(self) -> Iterator[Transaction]:
         from hathor.transaction.storage.traversal import BFSTimestampWalk
 
         root = self.iter_mempool_tips_from_tx_tips()
-        walk = BFSTimestampWalk(self, is_dag_funds=True, is_dag_verifications=True, is_left_to_right=False)
+        walk = BFSTimestampWalk(self.get_vertex, is_dag_funds=True, is_dag_verifications=True, is_left_to_right=False)
         for tx in walk.run(root):
             tx_meta = tx.get_metadata()
             # XXX: skip blocks and tx-tips that have already been confirmed

hathor/transaction/storage/traversal.py

@@ -18,11 +18,10 @@
 from abc import ABC, abstractmethod
 from collections import deque
 from itertools import chain
-from typing import TYPE_CHECKING, Iterable, Iterator, Optional, Union
+from typing import TYPE_CHECKING, Callable, Iterable, Iterator, Optional, Union
 
 if TYPE_CHECKING:
-    from hathor.transaction import BaseTransaction  # noqa: F401
-    from hathor.transaction.storage import VertexStorageProtocol
+    from hathor.transaction import BaseTransaction, Vertex  # noqa: F401
     from hathor.types import VertexId
 
 
@@ -50,7 +49,7 @@ class GenericWalk(ABC):
 
     def __init__(
         self,
-        storage: VertexStorageProtocol,
+        vertex_getter: Callable[[VertexId], Vertex],
         *,
         is_dag_funds: bool = False,
         is_dag_verifications: bool = False,
@@ -64,7 +63,7 @@ def __init__(
         :param is_dag_verifications: Add neighbors from the DAG of verifications
         :param is_left_to_right: Decide which side of the DAG we will walk to
         """
-        self.storage = storage
+        self.vertex_getter = vertex_getter
         self.seen = set()
 
         self.is_dag_funds = is_dag_funds
@@ -119,7 +118,7 @@ def add_neighbors(self, tx: 'BaseTransaction') -> None:
         for _hash in it:
             if _hash not in self.seen:
                 self.seen.add(_hash)
-                neighbor = self.storage.get_vertex(_hash)
+                neighbor = self.vertex_getter(_hash)
                 self._push_visit(neighbor)
 
     def skip_neighbors(self, tx: 'BaseTransaction') -> None:
@@ -164,14 +163,14 @@ class BFSTimestampWalk(GenericWalk):
 
     def __init__(
         self,
-        storage: VertexStorageProtocol,
+        vertex_getter: Callable[[VertexId], Vertex],
         *,
         is_dag_funds: bool = False,
         is_dag_verifications: bool = False,
         is_left_to_right: bool = True,
     ) -> None:
         super().__init__(
-            storage,
+            vertex_getter=vertex_getter,
             is_dag_funds=is_dag_funds,
             is_dag_verifications=is_dag_verifications,
             is_left_to_right=is_left_to_right
@@ -200,14 +199,14 @@ class BFSOrderWalk(GenericWalk):
 
     def __init__(
         self,
-        storage: VertexStorageProtocol,
+        vertex_getter: Callable[[VertexId], Vertex],
         *,
         is_dag_funds: bool = False,
         is_dag_verifications: bool = False,
         is_left_to_right: bool = True,
     ) -> None:
         super().__init__(
-            storage,
+            vertex_getter=vertex_getter,
             is_dag_funds=is_dag_funds,
             is_dag_verifications=is_dag_verifications,
             is_left_to_right=is_left_to_right
@@ -231,14 +230,14 @@ class DFSWalk(GenericWalk):
 
     def __init__(
         self,
-        storage: VertexStorageProtocol,
+        vertex_getter: Callable[[VertexId], Vertex],
         *,
         is_dag_funds: bool = False,
         is_dag_verifications: bool = False,
         is_left_to_right: bool = True,
     ) -> None:
         super().__init__(
-            storage,
+            vertex_getter=vertex_getter,
             is_dag_funds=is_dag_funds,
             is_dag_verifications=is_dag_verifications,
             is_left_to_right=is_left_to_right

hathor/transaction/storage/vertex_storage_protocol.py

@@ -197,7 +197,7 @@ def test_exceeds_streaming_and_mempool_limits(self) -> None:
         blk = manager1.tx_storage.get_best_block()
         tx_parents = [manager1.tx_storage.get_transaction(x) for x in blk.parents[1:]]
         self.assertEqual(len(tx_parents), 2)
-        dfs = DFSWalk(manager1.tx_storage, is_dag_verifications=True, is_left_to_right=False)
+        dfs = DFSWalk(manager1.tx_storage.get_vertex, is_dag_verifications=True, is_left_to_right=False)
         cnt = 0
         for tx in dfs.run(tx_parents):
             if tx.get_metadata().first_block == blk.hash:

tests/p2p/test_sync_v2.py

@@ -89,7 +89,7 @@ def test_right_to_left(self):
 
 class BaseBFSTimestampWalkTestCase(_TraversalTestCase):
     def gen_walk(self, **kwargs):
-        return BFSTimestampWalk(self.manager.tx_storage, **kwargs)
+        return BFSTimestampWalk(self.manager.tx_storage.get_vertex, **kwargs)
 
     def _run_lr(self, walk, skip_root=True):
         seen = set()
@@ -120,7 +120,7 @@ class SyncV2BFSTimestampWalkTestCase(unittest.SyncV2Params, BaseBFSTimestampWalk
 
 class BaseBFSOrderWalkTestCase(_TraversalTestCase):
     def gen_walk(self, **kwargs):
-        return BFSOrderWalk(self.manager.tx_storage, **kwargs)
+        return BFSOrderWalk(self.manager.tx_storage.get_vertex, **kwargs)
 
     def _run_lr(self, walk, skip_root=True):
         seen = set()
@@ -168,7 +168,7 @@ class SyncBridgeBFSOrderWalkTestCase(unittest.SyncBridgeParams, SyncV2BFSOrderWa
 
 class BaseDFSWalkTestCase(_TraversalTestCase):
     def gen_walk(self, **kwargs):
-        return DFSWalk(self.manager.tx_storage, **kwargs)
+        return DFSWalk(self.manager.tx_storage.get_vertex, **kwargs)
 
     def _run_lr(self, walk, skip_root=True):
         seen = set()