Enable complete type hinting for eth.db

Author: Bhargavasomu

eth/beacon/db/chain.py

@@ -243,7 +243,7 @@ def _get_canonical_head(cls, db: BaseDB) -> BaseBeaconBlock:
             canonical_head_hash = db[SchemaV1.make_canonical_head_hash_lookup_key()]
         except KeyError:
             raise CanonicalHeadNotFound("No canonical head set for this chain")
-        return cls._get_block_by_hash(db, canonical_head_hash)
+        return cls._get_block_by_hash(db, Hash32(canonical_head_hash))
 
     def get_block_by_hash(self, block_hash: Hash32) -> BaseBeaconBlock:
         return self._get_block_by_hash(self.db, block_hash)

eth/db/account.py

@@ -28,6 +28,9 @@
     BLANK_ROOT_HASH,
     EMPTY_SHA3,
 )
+from eth.db.backends.base import (
+    BaseDB,
+)
 from eth.db.batch import (
     BatchDB,
 )
@@ -71,7 +74,7 @@ def __init__(self) -> None:
 
     @property
     @abstractmethod
-    def state_root(self):
+    def state_root(self) -> Hash32:
         raise NotImplementedError("Must be implemented by subclasses")
 
     @abstractmethod
@@ -82,11 +85,11 @@ def has_root(self, state_root: bytes) -> bool:
     # Storage
     #
     @abstractmethod
-    def get_storage(self, address, slot):
+    def get_storage(self, address: Address, slot: int) -> int:
         raise NotImplementedError("Must be implemented by subclasses")
 
     @abstractmethod
-    def set_storage(self, address, slot, value):
+    def set_storage(self, address: Address, slot: int, value: int) -> None:
         raise NotImplementedError("Must be implemented by subclasses")
 
     #
@@ -104,40 +107,40 @@ def set_nonce(self, address: Address, nonce: int) -> None:
     # Balance
     #
     @abstractmethod
-    def get_balance(self, address):
+    def get_balance(self, address: Address) -> int:
         raise NotImplementedError("Must be implemented by subclasses")
 
     @abstractmethod
-    def set_balance(self, address, balance):
+    def set_balance(self, address: Address, balance: int) -> None:
         raise NotImplementedError("Must be implemented by subclasses")
 
-    def delta_balance(self, address, delta):
+    def delta_balance(self, address: Address, delta: int) -> None:
         self.set_balance(address, self.get_balance(address) + delta)
 
     #
     # Code
     #
     @abstractmethod
-    def set_code(self, address, code):
+    def set_code(self, address: Address, code: bytes) -> None:
         raise NotImplementedError("Must be implemented by subclasses")
 
     @abstractmethod
-    def get_code(self, address):
+    def get_code(self, address: Address) -> bytes:
         raise NotImplementedError("Must be implemented by subclasses")
 
     @abstractmethod
-    def get_code_hash(self, address):
+    def get_code_hash(self, address: Address) -> Hash32:
         raise NotImplementedError("Must be implemented by subclasses")
 
     @abstractmethod
-    def delete_code(self, address):
+    def delete_code(self, address: Address) -> None:
         raise NotImplementedError("Must be implemented by subclasses")
 
     #
     # Account Methods
     #
     @abstractmethod
-    def account_is_empty(self, address):
+    def account_is_empty(self, address: Address) -> bool:
         raise NotImplementedError("Must be implemented by subclass")
 
     #
@@ -177,7 +180,7 @@ class AccountDB(BaseAccountDB):
 
     logger = cast(TraceLogger, logging.getLogger('eth.db.account.AccountDB'))
 
-    def __init__(self, db, state_root=BLANK_ROOT_HASH):
+    def __init__(self, db: BaseDB, state_root: Hash32=BLANK_ROOT_HASH) -> None:
         r"""
         Internal implementation details (subject to rapid change):
         Database entries go through several pipes, like so...
@@ -225,11 +228,11 @@ def __init__(self, db, state_root=BLANK_ROOT_HASH):
         self._journaltrie = JournalDB(self._trie_cache)
 
     @property
-    def state_root(self):
+    def state_root(self) -> Hash32:
         return self._trie.root_hash
 
     @state_root.setter
-    def state_root(self, value):
+    def state_root(self, value: Hash32) -> None:
         self._trie_cache.reset_cache()
         self._trie.root_hash = value
 
@@ -239,7 +242,7 @@ def has_root(self, state_root: bytes) -> bool:
     #
     # Storage
     #
-    def get_storage(self, address, slot, from_journal=True):
+    def get_storage(self, address: Address, slot: int, from_journal: bool=True) -> int:
         validate_canonical_address(address, title="Storage Address")
         validate_uint256(slot, title="Storage Slot")
 
@@ -254,7 +257,7 @@ def get_storage(self, address, slot, from_journal=True):
         else:
             return 0
 
-    def set_storage(self, address, slot, value):
+    def set_storage(self, address: Address, slot: int, value: int) -> None:
         validate_uint256(value, title="Storage Value")
         validate_uint256(slot, title="Storage Slot")
         validate_canonical_address(address, title="Storage Address")
@@ -272,7 +275,7 @@ def set_storage(self, address, slot, value):
 
         self._set_account(address, account.copy(storage_root=storage.root_hash))
 
-    def delete_storage(self, address):
+    def delete_storage(self, address: Address) -> None:
         validate_canonical_address(address, title="Storage Address")
 
         account = self._get_account(address)
@@ -281,13 +284,13 @@ def delete_storage(self, address):
     #
     # Balance
     #
-    def get_balance(self, address):
+    def get_balance(self, address: Address) -> int:
         validate_canonical_address(address, title="Storage Address")
 
         account = self._get_account(address)
         return account.balance
 
-    def set_balance(self, address, balance):
+    def set_balance(self, address: Address, balance: int) -> None:
         validate_canonical_address(address, title="Storage Address")
         validate_uint256(balance, title="Account Balance")
 
@@ -297,35 +300,35 @@ def set_balance(self, address, balance):
     #
     # Nonce
     #
-    def get_nonce(self, address):
+    def get_nonce(self, address: Address) -> int:
         validate_canonical_address(address, title="Storage Address")
 
         account = self._get_account(address)
         return account.nonce
 
-    def set_nonce(self, address, nonce):
+    def set_nonce(self, address: Address, nonce: int) -> None:
         validate_canonical_address(address, title="Storage Address")
         validate_uint256(nonce, title="Nonce")
 
         account = self._get_account(address)
         self._set_account(address, account.copy(nonce=nonce))
 
-    def increment_nonce(self, address):
+    def increment_nonce(self, address: Address) -> None:
         current_nonce = self.get_nonce(address)
         self.set_nonce(address, current_nonce + 1)
 
     #
     # Code
     #
-    def get_code(self, address):
+    def get_code(self, address: Address) -> bytes:
         validate_canonical_address(address, title="Storage Address")
 
         try:
             return self._journaldb[self.get_code_hash(address)]
         except KeyError:
             return b""
 
-    def set_code(self, address, code):
+    def set_code(self, address: Address, code: bytes) -> None:
         validate_canonical_address(address, title="Storage Address")
         validate_is_bytes(code, title="Code")
 
@@ -335,13 +338,13 @@ def set_code(self, address, code):
         self._journaldb[code_hash] = code
         self._set_account(address, account.copy(code_hash=code_hash))
 
-    def get_code_hash(self, address):
+    def get_code_hash(self, address: Address) -> Hash32:
         validate_canonical_address(address, title="Storage Address")
 
         account = self._get_account(address)
         return account.code_hash
 
-    def delete_code(self, address):
+    def delete_code(self, address: Address) -> None:
         validate_canonical_address(address, title="Storage Address")
 
         account = self._get_account(address)
@@ -350,40 +353,40 @@ def delete_code(self, address):
     #
     # Account Methods
     #
-    def account_has_code_or_nonce(self, address):
+    def account_has_code_or_nonce(self, address: Address) -> bool:
         return self.get_nonce(address) != 0 or self.get_code_hash(address) != EMPTY_SHA3
 
-    def delete_account(self, address):
+    def delete_account(self, address: Address) -> None:
         validate_canonical_address(address, title="Storage Address")
 
         del self._journaltrie[address]
 
-    def account_exists(self, address):
+    def account_exists(self, address: Address) -> bool:
         validate_canonical_address(address, title="Storage Address")
 
         return self._journaltrie.get(address, b'') != b''
 
-    def touch_account(self, address):
+    def touch_account(self, address: Address) -> None:
         validate_canonical_address(address, title="Storage Address")
 
         account = self._get_account(address)
         self._set_account(address, account)
 
-    def account_is_empty(self, address):
+    def account_is_empty(self, address: Address) -> bool:
         return not self.account_has_code_or_nonce(address) and self.get_balance(address) == 0
 
     #
     # Internal
     #
-    def _get_account(self, address, from_journal=True):
+    def _get_account(self, address: Address, from_journal: bool=True) -> Account:
         rlp_account = (self._journaltrie if from_journal else self._trie_cache).get(address, b'')
         if rlp_account:
             account = rlp.decode(rlp_account, sedes=Account)
         else:
             account = Account()
         return account
 
-    def _set_account(self, address, account):
+    def _set_account(self, address: Address, account: Account) -> None:
         rlp_account = rlp.encode(account, sedes=Account)
         self._journaltrie[address] = rlp_account
 
@@ -424,7 +427,7 @@ def _log_pending_accounts(self) -> None:
                     continue
                 else:
                     accounts_displayed.add(address)
-                    account = self._get_account(address)
+                    account = self._get_account(Address(address))
                     self.logger.trace(
                         "Account %s: balance %d, nonce %d, storage root %s, code hash %s",
                         encode_hex(address),

eth/db/atomic.py

@@ -1,6 +1,9 @@
 from contextlib import contextmanager
 import logging
-from typing import Generator
+from typing import (
+    Generator,
+    Iterator,
+)
 
 from eth_utils import (
     ValidationError,
@@ -110,14 +113,14 @@ def _exists(self, key: bytes) -> bool:
 
     @classmethod
     @contextmanager
-    def _commit_unless_raises(cls, write_target_db):
+    def _commit_unless_raises(cls, write_target_db: BaseDB) -> Iterator['AtomicDBWriteBatch']:
         """
         Commit all writes inside the context, unless an exception was raised.
 
         Although this is technically an external API, it (and this whole class) is only intended
         to be used by AtomicDB.
         """
-        readable_write_batch = cls(write_target_db)
+        readable_write_batch = cls(write_target_db)     # type: AtomicDBWriteBatch
         try:
             yield readable_write_batch
         except Exception:

eth/db/backends/base.py

@@ -6,8 +6,18 @@
     MutableMapping,
 )
 
+from typing import (
+    Any,
+    TYPE_CHECKING
+)
+
+if TYPE_CHECKING:
+    MM = MutableMapping[bytes, bytes]
+else:
+    MM = MutableMapping
+
 
-class BaseDB(MutableMapping, ABC):
+class BaseDB(MM, ABC):
     """
     This is an abstract key/value lookup with all :class:`bytes` values,
     with some convenience methods for databases. As much as possible,
@@ -35,9 +45,10 @@ def set(self, key: bytes, value: bytes) -> None:
     def exists(self, key: bytes) -> bool:
         return self.__contains__(key)
 
-    def __contains__(self, key):
+    def __contains__(self, key: bytes) -> bool:     # type: ignore # Breaks LSP
         if hasattr(self, '_exists'):
-            return self._exists(key)
+            # Classes which inherit this class would have `_exists` attr
+            return self._exists(key)    # type: ignore
         else:
             return super().__contains__(key)
 
@@ -47,10 +58,10 @@ def delete(self, key: bytes) -> None:
         except KeyError:
             return None
 
-    def __iter__(self):
-        raise NotImplementedError("By default, DB classes cannot by iterated.")
+    def __iter__(self) -> None:
+        raise NotImplementedError("By default, DB classes cannot be iterated.")
 
-    def __len__(self):
+    def __len__(self) -> int:
         raise NotImplementedError("By default, DB classes cannot return the total number of keys.")
 
 
@@ -80,5 +91,5 @@ class BaseAtomicDB(BaseDB):
             # or neither will
     """
     @abstractmethod
-    def atomic_batch(self):
+    def atomic_batch(self) -> Any:
         raise NotImplementedError

eth/db/batch.py

@@ -37,7 +37,7 @@ def __exit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
             self.clear()
             self.logger.exception("Unexpected error occurred during batch update")
 
-    def clear(self):
+    def clear(self) -> None:
         self._track_diff = DBDiffTracker()
 
     def commit(self, apply_deletes: bool = True) -> None:

eth/db/cache.py

@@ -1,5 +1,7 @@
 from lru import LRU
 
+from typing import Any
+
 from eth.db.backends.base import BaseDB
 
 
@@ -8,24 +10,24 @@ class CacheDB(BaseDB):
     Set and get decoded RLP objects, where the underlying db stores
     encoded objects.
     """
-    def __init__(self, db, cache_size=2048):
+    def __init__(self, db: BaseDB, cache_size: int=2048) -> None:
         self._db = db
         self._cache_size = cache_size
         self.reset_cache()
 
-    def reset_cache(self):
+    def reset_cache(self) -> None:
         self._cached_values = LRU(self._cache_size)
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: Any) -> Any:
         if key not in self._cached_values:
             self._cached_values[key] = self._db[key]
         return self._cached_values[key]
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: Any, value: Any) -> None:
         self._cached_values[key] = value
         self._db[key] = value
 
-    def __delitem__(self, key):
+    def __delitem__(self, key: Any) -> None:
         if key in self._cached_values:
             del self._cached_values[key]
         del self._db[key]