wip

Author: pacrob

ssz/hash_tree.py

@@ -7,10 +7,13 @@
 )
 from typing import (
     Any,
+    Callable,
     Generator,
     Iterable,
     Optional,
+    Tuple,
     Union,
+    cast,
 )
 
 # `transform` comes from a non-public API which is considered stable, but future changes
@@ -105,8 +108,13 @@ def chunks(self) -> RawHashTreeLayer:
     def root(self) -> Hash32:
         return self.raw_hash_tree[-1][0]
 
-    def transform(self, *transformations):
-        return transform(self, transformations)
+    def transform(
+        self,
+        *transformations: Tuple[
+            Tuple[Tuple[int, ...], Union[Any, Callable[[Any], Any]]], ...
+        ],
+    ) -> "HashTree":
+        return cast("HashTree", transform(self, transformations))
 
     def evolver(self) -> "HashTreeEvolver":
         return HashTreeEvolver(self)
@@ -133,7 +141,7 @@ def __len__(self) -> int:
     def __getitem__(self, index: Union[int, slice]) -> Hash32:
         return self.chunks[index]
 
-    def index(self, value: Hash32, *args, **kwargs) -> Hash32:
+    def index(self, value: Hash32, *args: Any, **kwargs: Any) -> Hash32:
         return self.chunks.index(value, *args, **kwargs)
 
     def count(self, value: Hash32) -> int:
@@ -222,7 +230,7 @@ def __len__(self) -> int:
         return len(self.original_hash_tree) + len(self.appended_chunks)
 
     def is_dirty(self) -> bool:
-        return self.updated_chunks or self.appended_chunks
+        return any([self.updated_chunks, self.appended_chunks])
 
     #
     # Setters
@@ -231,16 +239,18 @@ def set(self, index: Integral, value: Hash32) -> None:
         self[index] = value
 
     def __setitem__(self, index: Integral, value: Hash32) -> None:
-        if index < 0:
-            index += len(self)
+        idx = int(index)
+
+        if idx < 0:
+            idx += len(self)
 
-        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)
+        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)
             self.appended_chunks = self.appended_chunks.set(index_in_appendix, value)
         else:
-            raise IndexError(f"Index out of bounds: {index}")
+            raise IndexError(f"Index out of bounds: {idx}")
 
     #
     # Length changing modifiers