update: str to bytes

olivmath · Dec 23, 2023 · 5211160 · 5211160
1 parent a0fbd4e
commit 5211160
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/merkly/mtree.py b/merkly/mtree.py
@@ -30,16 +30,17 @@ class MerkleTree:
     def __init__(
         self,
         leafs: List[str],
-        hash_function: Callable[[str], str] = lambda x, y: keccak(x + y),
+        hash_function: Callable[[bytes, bytes], bytes] = lambda x, y: keccak(x + y),
     ) -> None:
         validate_leafs(leafs)
         validate_hash_function(hash_function)
+        self.hash_function: Callable[[bytes, bytes], bytes] = hash_function
         self.raw_leafs: List[str] = leafs
         self.leafs: List[str] = self.__hash_leafs(leafs)
         self.short_leafs: List[str] = self.short(self.leafs)
 
     def __hash_leafs(self, leafs: List[str]) -> List[str]:
-        return list(map(lambda x: self.hash_function(x, ""), leafs))
+        return list(map(lambda x: self.hash_function(x.encode(), b""), leafs))
 
     def __repr__(self) -> str:
         return f"""MerkleTree(\nraw_leafs: {self.raw_leafs}\nleafs: {self.leafs}\nshort_leafs: {self.short(self.leafs)})"""
@@ -48,8 +49,8 @@ def short(self, data: List[str]) -> List[str]:
         return [f"{x[:4]}..." for x in data]
 
     @property
-    def root(self) -> str:
-        return self.make_root(self.leafs)[0]
+    def root(self) -> bytes:
+        return self.make_root(self.leafs)
 
     def proof(self, raw_leaf: str) -> List[Node]:
         return self.make_proof(self.leafs, [], self.hash_function(raw_leaf, ""))

diff --git a/merkly/utils.py b/merkly/utils.py
@@ -23,7 +23,7 @@ def __init__(self) -> None:
         super().__init__(self.message)
 
 
-def keccak(data: str) -> str:
+def keccak(data: bytes) -> bytes:
     """
     # Hash `data: str` using keccak256
     - params `data: str`
@@ -45,9 +45,9 @@ def keccak(data: str) -> str:
     """
 
     keccak_256 = cryptodome_keccak.new(digest_bits=256)
-    keccak_256.update(data.encode())
+    keccak_256.update(data)
 
-    return keccak_256.hexdigest()
+    return keccak_256.digest()
 
 
 def half(list_item: List[int]) -> Tuple[int, int]:
@@ -102,7 +102,7 @@ def validate_hash_function(hash_function: Callable[[bytes, bytes], bytes]):
     a = isinstance(hash_function, types.FunctionType)
     b = callable(hash_function)
     try:
-        c = isinstance(hash_function(str(), str()), str)
+        c = isinstance(hash_function(bytes(), bytes()), bytes)
     except TypeError:
         c = False