chore: add type hints (#157)

Co-authored-by: CodeAgentCN <codeagentcn@protonmail.com>

Author: CodeAgentCN

setup.py

@@ -46,6 +46,7 @@
     include_package_data=True,
     install_requires=[
         "eth-hash>=0.1.0",
+        "eth-typing>=2.0.0,<3",
         "eth-utils>=2.0.0",
         "hexbytes>=0.2.3",
         "rlp>=3",

trie/binary.py

@@ -1,7 +1,17 @@
+from typing import (
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
+
 from eth_hash.auto import (
     keccak,
 )
-
 from trie.constants import (
     BLANK_HASH,
     BRANCH_TYPE,
@@ -30,12 +40,12 @@
 
 
 class BinaryTrie:
-    def __init__(self, db, root_hash=BLANK_HASH):
+    def __init__(self, db: Dict[bytes, bytes], root_hash: bytes = BLANK_HASH) -> None:
         self.db = db
         validate_is_bytes(root_hash)
         self.root_hash = root_hash
 
-    def get(self, key):
+    def get(self, key: bytes) -> Optional[bytes]:
         """
         Fetches the value with a given keypath from the given node.
 
@@ -45,7 +55,7 @@ def get(self, key):
 
         return self._get(self.root_hash, encode_to_bin(key))
 
-    def _get(self, node_hash, keypath):
+    def _get(self, node_hash: bytes, keypath: bytes) -> Optional[bytes]:
         """
         Note: keypath should be in binary array format, i.e., encoded by encode_to_bin()
         """
@@ -76,7 +86,7 @@ def _get(self, node_hash, keypath):
             else:
                 return self._get(right_child, keypath[1:])
 
-    def set(self, key, value):
+    def set(self, key: bytes, value: bytes) -> None:
         """
         Sets the value at the given keypath from the given node
 
@@ -87,7 +97,11 @@ def set(self, key, value):
 
         self.root_hash = self._set(self.root_hash, encode_to_bin(key), value)
 
-    def _set(self, node_hash, keypath, value, if_delete_subtrie=False):
+    def _set(self,
+             node_hash: bytes,
+             keypath: bytes,
+             value: bytes,
+             if_delete_subtrie: bool = False) -> bytes:
         """
         If if_delete_subtrie is set to True, what it will do is that it take in a
         keypath and traverse til the end of keypath, then delete the whole subtrie
@@ -155,14 +169,14 @@ def _set(self, node_hash, keypath, value, if_delete_subtrie=False):
 
     def _set_kv_node(
         self,
-        keypath,
-        node_hash,
-        node_type,
-        left_child,
-        right_child,
-        value,
-        if_delete_subtrie=False,
-    ):
+        keypath: bytes,
+        node_hash: bytes,
+        node_type: int,
+        left_child: bytes,
+        right_child: bytes,
+        value: bytes,
+        if_delete_subtrie: bool = False,
+    ) -> bytes:
         # Keypath prefixes match
         if if_delete_subtrie:
             if len(keypath) < len(left_child) and keypath == left_child[: len(keypath)]:
@@ -257,13 +271,13 @@ def _set_kv_node(
 
     def _set_branch_node(
         self,
-        keypath,
-        node_type,
-        left_child,
-        right_child,
-        value,
-        if_delete_subtrie=False,
-    ):
+        keypath: bytes,
+        node_type: int,
+        left_child: bytes,
+        right_child: bytes,
+        value: bytes,
+        if_delete_subtrie: bool = False,
+    ) -> bytes:
         # Which child node to update? Depends on first bit in keypath
         if keypath[:1] == BYTE_0:
             new_left_child = self._set(
@@ -303,20 +317,20 @@ def _set_branch_node(
                 encode_branch_node(new_left_child, new_right_child)
             )
 
-    def exists(self, key):
+    def exists(self, key: bytes) -> bool:
         validate_is_bytes(key)
 
         return self.get(key) is not None
 
-    def delete(self, key):
+    def delete(self, key: bytes) -> None:
         """
         Equals to setting the value to None
         """
         validate_is_bytes(key)
 
         self.root_hash = self._set(self.root_hash, encode_to_bin(key), b"")
 
-    def delete_subtrie(self, key):
+    def delete_subtrie(self, key: bytes) -> None:
         """
         Given a key prefix, delete the whole subtrie that starts with the key prefix.
 
@@ -337,19 +351,19 @@ def delete_subtrie(self, key):
     # Convenience
     #
     @property
-    def root_node(self):
+    def root_node(self) -> bytes:
         return self.db[self.root_hash]
 
     @root_node.setter
-    def root_node(self, node):
+    def root_node(self, node: bytes) -> None:
         validate_is_bin_node(node)
 
         self.root_hash = self._hash_and_save(node)
 
     #
     # Utils
     #
-    def _hash_and_save(self, node):
+    def _hash_and_save(self, node: bytes) -> bytes:
         """
         Saves a node into the database and returns its hash
         """
@@ -362,14 +376,14 @@ def _hash_and_save(self, node):
     #
     # Dictionary API
     #
-    def __getitem__(self, key):
+    def __getitem__(self, key: bytes) -> Optional[bytes]:
         return self.get(key)
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: bytes, value: bytes) -> None:
         return self.set(key, value)
 
-    def __delitem__(self, key):
+    def __delitem__(self, key: bytes) -> None:
         return self.delete(key)
 
-    def __contains__(self, key):
+    def __contains__(self, key: bytes) -> bool:
         return self.exists(key)

trie/branches.py

@@ -1,3 +1,10 @@
+from typing import (
+    Dict,
+    Iterator,
+    Optional,
+    Tuple,
+)
+
 from eth_hash.auto import (
     keccak,
 )
@@ -27,7 +34,7 @@
 )
 
 
-def check_if_branch_exist(db, root_hash, key_prefix):
+def check_if_branch_exist(db: Dict[bytes, bytes], root_hash: bytes, key_prefix: bytes) -> bool:
     """
     Given a key prefix, return whether this prefix is
     the prefix of an existing key in the trie.
@@ -37,7 +44,7 @@ def check_if_branch_exist(db, root_hash, key_prefix):
     return _check_if_branch_exist(db, root_hash, encode_to_bin(key_prefix))
 
 
-def _check_if_branch_exist(db, node_hash, key_prefix):
+def _check_if_branch_exist(db: Dict[bytes, bytes], node_hash: bytes, key_prefix: bytes) -> bool:
     # Empty trie
     if node_hash == BLANK_HASH:
         return False
@@ -70,7 +77,7 @@ def _check_if_branch_exist(db, node_hash, key_prefix):
         raise Exception("Invariant: unreachable code path")
 
 
-def get_branch(db, root_hash, key):
+def get_branch(db: Dict[bytes, bytes], root_hash: bytes, key: bytes) -> Tuple[bytes, ...]:
     """
     Get a long-format Merkle branch
     """
@@ -79,7 +86,7 @@ def get_branch(db, root_hash, key):
     return tuple(_get_branch(db, root_hash, encode_to_bin(key)))
 
 
-def _get_branch(db, node_hash, keypath):
+def _get_branch(db: Dict[bytes, bytes], node_hash: bytes, keypath: bytes) -> Iterator[bytes]:
     if node_hash == BLANK_HASH:
         return
     node = db[node_hash]
@@ -110,7 +117,7 @@ def _get_branch(db, node_hash, keypath):
         raise Exception("Invariant: unreachable code path")
 
 
-def if_branch_valid(branch, root_hash, key, value):
+def if_branch_valid(branch: Tuple[bytes, ...], root_hash: bytes, key: bytes, value: Optional[bytes]) -> bool:
     # value being None means the key is not in the trie
     if value is not None:
         validate_is_bytes(key)
@@ -124,14 +131,14 @@ def if_branch_valid(branch, root_hash, key, value):
     return True
 
 
-def get_trie_nodes(db, node_hash):
+def get_trie_nodes(db: Dict[bytes, bytes], node_hash: bytes) -> Tuple[bytes, ...]:
     """
     Get full trie of a given root node
     """
     return tuple(_get_trie_nodes(db, node_hash))
 
 
-def _get_trie_nodes(db, node_hash):
+def _get_trie_nodes(db: Dict[bytes, bytes], node_hash: bytes) -> Iterator[bytes]:
     if node_hash in db:
         node = db[node_hash]
     else:
@@ -150,7 +157,7 @@ def _get_trie_nodes(db, node_hash):
         raise Exception("Invariant: unreachable code path")
 
 
-def get_witness_for_key_prefix(db, node_hash, key):
+def get_witness_for_key_prefix(db: Dict[bytes, bytes], node_hash: bytes, key: bytes) -> Tuple[bytes, ...]:
     """
     Get all witness given a keypath prefix.
     Include
@@ -163,7 +170,7 @@ def get_witness_for_key_prefix(db, node_hash, key):
     return tuple(_get_witness_for_key_prefix(db, node_hash, encode_to_bin(key)))
 
 
-def _get_witness_for_key_prefix(db, node_hash, keypath):
+def _get_witness_for_key_prefix(db: Dict[bytes, bytes], node_hash: bytes, keypath: bytes) -> Iterator[bytes]:
     if not keypath:
         yield from get_trie_nodes(db, node_hash)
     if node_hash in db:

trie/constants.py

@@ -1,35 +1,39 @@
-BLANK_NODE = b""
+from typing import (
+    Tuple,
+)
+
+BLANK_NODE: bytes = b""
 # keccak(b'')
-BLANK_HASH = b"\xc5\xd2F\x01\x86\xf7#<\x92~}\xb2\xdc\xc7\x03\xc0\xe5\x00\xb6S\xca\x82';{\xfa\xd8\x04]\x85\xa4p"  # noqa: E501
+BLANK_HASH: bytes = b"\xc5\xd2F\x01\x86\xf7#<\x92~}\xb2\xdc\xc7\x03\xc0\xe5\x00\xb6S\xca\x82';{\xfa\xd8\x04]\\x85\xa4p"  # noqa: E501
 # keccak(rlp.encode(b''))
-BLANK_NODE_HASH = b"V\xe8\x1f\x17\x1b\xccU\xa6\xff\x83E\xe6\x92\xc0\xf8n\x5bH\xe0\x1b\x99l\xad\xc0\x01b/\xb5\xe3c\xb4!"  # noqa: E501
+BLANK_NODE_HASH: bytes = b"V\xe8\x1f\x17\x1b\xccU\xa6\xff\x83E\xe6\x92\xc0\xf8n\x5bH\xe0\x1b\x99l\xad\xc0\x01b/\xb5\xe3c\xb4!"  # noqa: E501
 
 
-NIBBLE_TERMINATOR = 16
+NIBBLE_TERMINATOR: int = 16
 
-HP_FLAG_2 = 2
-HP_FLAG_0 = 0
+HP_FLAG_2: int = 2
+HP_FLAG_0: int = 0
 
 
-NODE_TYPE_BLANK = 0
-NODE_TYPE_LEAF = 1
-NODE_TYPE_EXTENSION = 2
-NODE_TYPE_BRANCH = 3
+NODE_TYPE_BLANK: int = 0
+NODE_TYPE_LEAF: int = 1
+NODE_TYPE_EXTENSION: int = 2
+NODE_TYPE_BRANCH: int = 3
 
 # Constants for Binary Trie
-EXP = tuple(reversed(tuple(2**i for i in range(8))))
-
-TWO_BITS = [bytes([0, 0]), bytes([0, 1]), bytes([1, 0]), bytes([1, 1])]
-PREFIX_00 = bytes([0, 0])
-PREFIX_100000 = bytes([1, 0, 0, 0, 0, 0])
-
-KV_TYPE = 0
-BRANCH_TYPE = 1
-LEAF_TYPE = 2
-BINARY_TRIE_NODE_TYPES = (0, 1, 2)
-KV_TYPE_PREFIX = bytes([0])
-BRANCH_TYPE_PREFIX = bytes([1])
-LEAF_TYPE_PREFIX = bytes([2])
-
-BYTE_1 = bytes([1])
-BYTE_0 = bytes([0])
+EXP: Tuple[int, ...] = tuple(reversed(tuple(2**i for i in range(8))))
+
+TWO_BITS: Tuple[bytes, bytes, bytes, bytes] = (bytes([0, 0]), bytes([0, 1]), bytes([1, 0]), bytes([1, 1]))
+PREFIX_00: bytes = bytes([0, 0])
+PREFIX_100000: bytes = bytes([1, 0, 0, 0, 0, 0])
+
+KV_TYPE: int = 0
+BRANCH_TYPE: int = 1
+LEAF_TYPE: int = 2
+BINARY_TRIE_NODE_TYPES: Tuple[int, int, int] = (0, 1, 2)
+KV_TYPE_PREFIX: bytes = bytes([0])
+BRANCH_TYPE_PREFIX: bytes = bytes([1])
+LEAF_TYPE_PREFIX: bytes = bytes([2])
+
+BYTE_1: bytes = bytes([1])
+BYTE_0: bytes = bytes([0])

trie/exceptions.py

@@ -65,7 +65,7 @@ def __init__(
         requested_key: bytes,
         prefix: Nibbles = None,
         *args,
-    ):
+    ) -> None:
         if not isinstance(missing_node_hash, bytes):
             raise TypeError(
                 "Missing node hash must be bytes, was: %r" % missing_node_hash

trie/hexary.py

@@ -5,9 +5,17 @@
 import functools
 import itertools
 from typing import (
+    Any,
     Callable,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Set,
     Tuple,
     TypeVar,
+    Union,
     cast,
 )
 
@@ -96,7 +104,7 @@ class HexaryTrie:
     BLANK_NODE_HASH = BLANK_NODE_HASH
     BLANK_NODE = BLANK_NODE
 
-    def __init__(self, db, root_hash=BLANK_NODE_HASH, prune=False, ref_count=None):
+    def __init__(self, db: Dict, root_hash: bytes = BLANK_NODE_HASH, prune: bool = False, ref_count: Dict = None) -> None:
         """
         Important note about Pruning:
 
@@ -128,7 +136,7 @@ def __init__(self, db, root_hash=BLANK_NODE_HASH, prune=False, ref_count=None):
                 )
         self._pending_prune_keys = None
 
-    def get(self, key):
+    def get(self, key: bytes) -> bytes:
         validate_is_bytes(key)
 
         trie_key = bytes_to_nibbles(key)

trie/py.typed

@@ -0,0 +1 @@
+