review changes

Author: glevco

hathor/transaction/base_transaction.py

@@ -825,7 +825,7 @@ def serialize_output(tx: BaseTransaction, tx_out: TxOutput) -> dict[str, Any]:
 
         return ret
 
-    def clone(self, *, include_metadata: bool = True) -> 'BaseTransaction':
+    def clone(self, *, include_metadata: bool = True, include_storage: bool = True) -> 'BaseTransaction':
         """Return exact copy without sharing memory, including metadata if loaded.
 
         :return: Transaction or Block copy
@@ -834,7 +834,8 @@ def clone(self, *, include_metadata: bool = True) -> 'BaseTransaction':
         if hasattr(self, '_metadata') and include_metadata:
             assert self._metadata is not None  # FIXME: is this actually true or do we have to check if not None
             new_tx._metadata = self._metadata.clone()
-        new_tx.storage = self.storage
+        if include_storage:
+            new_tx.storage = self.storage
         return new_tx
 
     @abstractmethod

hathor/transaction/storage/simple_memory_storage.py

@@ -12,22 +12,13 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from dataclasses import dataclass
-from typing import Any
-
-from hathor.transaction import Block, Transaction, TransactionMetadata
-from hathor.transaction.base_transaction import BaseTransaction, tx_or_block_from_bytes
+from hathor.transaction import Block, Transaction
+from hathor.transaction.base_transaction import BaseTransaction
 from hathor.transaction.storage import TransactionStorage
 from hathor.transaction.storage.exceptions import TransactionDoesNotExist
 from hathor.types import VertexId
 
 
-@dataclass(frozen=True, slots=True)
-class _SimpleMemoryRecord:
-    vertex_bytes: bytes
-    vertex_metadata: dict[str, Any]
-
-
 class SimpleMemoryStorage:
     """
     Instances of this class simply facilitate storing some data in memory, specifically for pre-fetched verification
@@ -36,33 +27,30 @@ class SimpleMemoryStorage:
     __slots__ = ('_blocks', '_transactions',)
 
     def __init__(self) -> None:
-        self._blocks: dict[VertexId, _SimpleMemoryRecord] = {}
-        self._transactions: dict[VertexId, _SimpleMemoryRecord] = {}
+        self._blocks: dict[VertexId, BaseTransaction] = {}
+        self._transactions: dict[VertexId, BaseTransaction] = {}
 
     @property
-    def _vertices(self) -> dict[VertexId, _SimpleMemoryRecord]:
+    def _vertices(self) -> dict[VertexId, BaseTransaction]:
         """Blocks and Transactions together."""
         return {**self._blocks, **self._transactions}
 
     def get_block(self, block_id: VertexId) -> Block:
         """Return a block from the storage, throw if it's not found."""
-        block = self._get_record(self._blocks, block_id)
+        block = self._get_vertex(self._blocks, block_id)
         assert isinstance(block, Block)
         return block
 
     def get_transaction(self, tx_id: VertexId) -> Transaction:
         """Return a transaction from the storage, throw if it's not found."""
-        tx = self._get_record(self._transactions, tx_id)
+        tx = self._get_vertex(self._transactions, tx_id)
         assert isinstance(tx, Transaction)
         return tx
 
     @staticmethod
-    def _get_record(storage: dict[VertexId, _SimpleMemoryRecord], vertex_id: VertexId) -> BaseTransaction:
-        """Return a record from a storage, throw if it's not found."""
-        if record := storage.get(vertex_id):
-            vertex = tx_or_block_from_bytes(record.vertex_bytes)
-            metadata = TransactionMetadata.create_from_json(record.vertex_metadata)
-            vertex._metadata = metadata
+    def _get_vertex(storage: dict[VertexId, BaseTransaction], vertex_id: VertexId) -> BaseTransaction:
+        """Return a vertex from a storage, throw if it's not found."""
+        if vertex := storage.get(vertex_id):
             return vertex
 
         raise TransactionDoesNotExist(f'Vertex "{vertex_id.hex()}" does not exist in this SimpleMemoryStorage.')
@@ -90,16 +78,14 @@ def add_vertex_from_storage(self, storage: TransactionStorage, vertex_id: Vertex
             return
 
         vertex = storage.get_transaction(vertex_id)
-        vertex_bytes = vertex.get_struct()
-        metadata = vertex.get_metadata().to_json()
-        record = _SimpleMemoryRecord(vertex_bytes, metadata)
+        clone = vertex.clone(include_metadata=True, include_storage=False)
 
         if isinstance(vertex, Block):
-            self._blocks[vertex_id] = record
+            self._blocks[vertex_id] = clone
             return
 
         if isinstance(vertex, Transaction):
-            self._transactions[vertex_id] = record
+            self._transactions[vertex_id] = clone
             return
 
         raise NotImplementedError