feat: str to bytes

merkly/mtree.py

@@ -7,11 +7,12 @@
 
 from merkly.node import Node, Side
 from merkly.utils import (
-    hash_function_type_checking,
+    validate_hash_function,
     is_power_2,
     slice_in_pairs,
     keccak,
     half,
+    validate_leafs,
 )
 
 
@@ -29,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:
-        hash_function_type_checking(hash_function)
-        self.hash_function: Callable[[str], str] = hash_function
+        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)})"""
@@ -47,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, ""))

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]:
@@ -88,11 +88,21 @@ def slice_in_pairs(list_item: list):
     return [list_item[i : i + 2] for i in range(0, len(list_item), 2)]
 
 
-def hash_function_type_checking(hash_function: Callable[[str], str]) -> bool:
+def validate_leafs(leafs: List[str]):
+    if len(leafs) < 2:
+        raise Exception("Invalid size, need > 2")
+
+    a = isinstance(leafs, List)
+    b = all(isinstance(leaf, str) for leaf in leafs)
+    if not (a and b):
+        raise Exception("Invalid type of leafs")
+
+
+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
 

test/test_merkle_tree.py

@@ -17,18 +17,21 @@
     tree = MerkleTree(leafs)
 
     assert tree.raw_leafs == leafs
-    assert tree.short_leafs == [
-        "3ac2...",
-        "b555...",
-        "0b42...",
-        "f191...",
-    ]
+    for i, j in zip(tree.short_leafs, [
+        bytes.fromhex("3ac2"),
+        bytes.fromhex("b555"),
+        bytes.fromhex("0b42"),
+        bytes.fromhex("f191"),
+    ]):
+        assert i == j
     assert tree.leafs == [
-        "3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb",
-        "b5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
-        "0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2",
-        "f1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3",
+        bytes.fromhex("3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb"),
+        bytes.fromhex("b5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510"),
+        bytes.fromhex("0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2"),
+        bytes.fromhex("f1918e8562236eb17adc8502332f4c9c82bc14e19bfc0aa10ab674ff75b3d2f3"),
     ]
+
+    assert tree.root.hex() == "68203f90e9d07dc5859259d7536e87a6ba9d345f2552b5b9de2999ddce9ce1bf"
 
 
 @mark.parametrize(