updates for hash_tree

Author: pacrob

ssz/hash_tree.py

@@ -2,9 +2,6 @@
     partial,
 )
 import itertools
-from numbers import (
-    Integral,
-)
 from typing import (
     Any,
     Callable,
@@ -14,6 +11,7 @@
     Tuple,
     Union,
     cast,
+    overload,
 )
 
 # `transform` comes from a non-public API which is considered stable, but future changes
@@ -40,6 +38,7 @@
 from pyrsistent.typing import (
     PMap,
     PVector,
+    PVectorEvolver,
 )
 
 from ssz.constants import (
@@ -163,7 +162,9 @@ def extend(self, value: Iterable[Hash32]) -> "HashTree":
     def __add__(self, other: Iterable[Hash32]) -> "HashTree":
         return self.extend(other)
 
-    def __mul__(self, times: int) -> "HashTree":
+    # we override __mul__ to allow for a more natural syntax
+    # when using the evolver
+    def __mul__(self, times: int) -> "HashTree":  # type: ignore[override]
         if times <= 0:
             raise ValueError(f"Multiplier must be greater or equal to 1, got {times}")
 
@@ -202,7 +203,7 @@ def remove(self, value: Hash32) -> "HashTree":
         return self.__class__.compute(chunks, self.chunk_count)
 
 
-class HashTreeEvolver:
+class HashTreeEvolver(PVectorEvolver[Hash32]):
     def __init__(self, hash_tree: "HashTree") -> None:
         self.original_hash_tree = hash_tree
         self.updated_chunks: PMap[int, Hash32] = pmap()
@@ -211,7 +212,18 @@ def __init__(self, hash_tree: "HashTree") -> None:
     #
     # Getters
     #
+    @overload
     def __getitem__(self, index: int) -> Hash32:
+        ...
+
+    @overload
+    def __getitem__(self, index: slice) -> "HashTreeEvolver":
+        ...
+
+    def __getitem__(self, index: Union[int, slice]) -> Union[Hash32, "HashTreeEvolver"]:
+        if isinstance(index, slice):
+            raise NotImplementedError("Slicing not implemented.")
+
         if index < 0:
             index += len(self)
 
@@ -235,33 +247,34 @@ def is_dirty(self) -> bool:
     #
     # Setters
     #
-    def set(self, index: Integral, value: Hash32) -> None:
+    def set(self, index: int, value: Hash32) -> "HashTreeEvolver":
         self[index] = value
+        return self
 
-    def __setitem__(self, index: Integral, value: Hash32) -> None:
-        idx = int(index)
-
-        if idx < 0:
-            idx += len(self)
+    def __setitem__(self, index: int, value: Hash32) -> None:
+        if index < 0:
+            index += len(self)
 
-        if 0 <= idx < len(self.original_hash_tree):
-            self.updated_chunks = self.updated_chunks.set(idx, value)
-        elif idx < len(self):
-            index_in_appendix = idx - len(self.original_hash_tree)
+        if 0 <= index < len(self.original_hash_tree):
+            self.updated_chunks = self.updated_chunks.set(index, value)
+        elif index < len(self):
+            index_in_appendix = index - len(self.original_hash_tree)
             self.appended_chunks = self.appended_chunks.set(index_in_appendix, value)
         else:
-            raise IndexError(f"Index out of bounds: {idx}")
+            raise IndexError(f"Index out of bounds: {index}")
 
     #
     # Length changing modifiers
     #
-    def append(self, value: Hash32) -> None:
+    def append(self, value: Hash32) -> "HashTreeEvolver":
         self.appended_chunks = self.appended_chunks.append(value)
         self._check_chunk_count()
+        return self
 
-    def extend(self, values: Iterable[Hash32]) -> None:
+    def extend(self, values: Iterable[Hash32]) -> "HashTreeEvolver":
         self.appended_chunks = self.appended_chunks.extend(values)
         self._check_chunk_count()
+        return self
 
     def _check_chunk_count(self) -> None:
         chunk_count = self.original_hash_tree.chunk_count
@@ -271,10 +284,10 @@ def _check_chunk_count(self) -> None:
     #
     # Not implemented
     #
-    def delete(self, index: int, stop: Optional[int] = None) -> None:
+    def delete(self, index: int, stop: Optional[int] = None) -> None:  # type: ignore[override] # noqa: E501
         raise NotImplementedError()
 
-    def __delitem__(self, index: int) -> None:
+    def __delitem__(self, index: Union[int, slice]) -> None:
         raise NotImplementedError()
 
     def remove(self, value: Hash32) -> None:
@@ -432,7 +445,7 @@ def set_chunk_in_tree(hash_tree: RawHashTree, index: int, chunk: Hash32) -> RawH
         for layer_index, hash_index in zip(parent_layer_indices, parent_hash_indices)
     )
 
-    hash_tree_with_updated_branch = pipe(
+    hash_tree_with_updated_branch: PVector[PVector[Any]] = pipe(
         hash_tree_with_updated_chunk, *update_functions
     )