Another pass: rough shape of tree and tests

It isn't working yet, the tests don't pass.

But it is on the way!

merkle_tree_stream/generator.py

@@ -5,11 +5,10 @@ from typing import Any, Callable, List, Optional
 import attr
 from flat_tree import FlatTreeAccessor
 
-__all__ = ['MerkleTreeGenerator', 'MerkleTreeNode']
-
-
 Hash = str
 
+__all__ = ['MerkleTreeGenerator', 'MerkleTreeNode']
+
 flat_tree = FlatTreeAccessor()
 
 
@@ -17,17 +16,17 @@ flat_tree = FlatTreeAccessor()
 class MerkleTreeNode:
     """A node in a merkle tree.
 
-    :param index: TODO
+    :param index: The index of node
-    :param parent: TODO
+    :param parent: The parent of the node
-    :param size: TODO
+    :param size: The size of the data
-    :param data: TODO
+    :param data: The data of the node
-    :param hash: TODO
+    :param hash: The hash of the data
     """
 
     index: int
-    parent: Optional[int]
+    parent: int
     size: int
-    data: bytes
+    data: Optional[bytes]
     hash: Optional[str] = None
 
     def __attrs_post_init__(self) -> Any:
@@ -39,30 +38,69 @@ class MerkleTreeNode:
 class MerkleTreeGenerator:
     """A stream that generates a merkle tree based on the incoming data.
 
-    :param leaf: TODO
+    :param leaf: The leaf hash generation function
-    :param parent: TODO
+    :param parent: The parent hash generation function
-    :param roots: TODO
+    :param roots: The tree roots
-    :param blocks: TODO
     """
 
-    leaf: Callable[[bytes], Hash]
+    leaf: Callable[[MerkleTreeNode, List[MerkleTreeNode]], Hash]
-    parent: Callable[[bytes], Hash]
+    parent: Callable[[MerkleTreeNode, List[MerkleTreeNode]], Hash]
-    blocks: int
+    roots: List[MerkleTreeNode] = attr.Factory(list)
-    roots: Optional[List[MerkleTreeNode]] = attr.Factory(list)
 
-    def next(self, data: bytes) -> List[MerkleTreeNode]:
+    def next(
-        """Further generate the treem based on the incoming data.
+        self, data: bytes, nodes: Optional[List[MerkleTreeNode]] = None
+    ) -> List[MerkleTreeNode]:
+        """Further generate the tree based on the incoming data.
 
         :param data: Incoming data
+        :param nodes: Pre-existing nodes
         """
-        pass
+        nodes = nodes or []
+
+        index = 2 * (self.blocks + 1)
+
+        leaf_node = MerkleTreeNode(
+            index=index,
+            parent=flat_tree.parent(index),
+            hash=None,
+            data=data,
+            size=len(data),
+        )
+
+        leaf_node.hash = self.leaf(leaf_node, self.roots)
+
+        self.roots.append(leaf_node)
+        nodes.append(leaf_node)
+
+        while len(self.roots) > 1:
+            left = self.roots[len(self.roots) - 2]
+            right = self.roots[len(self.roots) - 1]
+
+            if left.parent != right.parent:
+                break
+
+            self.roots.pop()
+
+            new_node = MerkleTreeNode(
+                index=left.parent,
+                parent=flat_tree.parent(left.parent),
+                hash=self.parent(left, [right]),
+                size=left.size + right.size,
+                data=None,
+            )
+
+            self.roots[len(self.roots) - 1] = new_node
+
+            nodes.append(new_node)
+
+        return nodes
 
     def __attrs_post_init__(self) -> Any:
         """Initialise parent and block defaults."""
+        index = self.roots[len(self.roots) - 1].index
+        right_span = flat_tree.right_span(index)
+        self.blocks = (1 + (right_span / 2)) if self.roots else 0
+
         for root in self.roots:
             if not root.parent:
                 root.parent = flat_tree.parent(root.index)
-
-        # https://github.com/mafintosh/merkle-tree-stream/blob/master/generator.js#L14
-        # self.roots[self.roots.length] ...
-        # self.blocks = (1 + (flat_tree.right_span(...) / 2)) if self.roots else 0

setup.cfg

@@ -45,6 +45,9 @@ package_dir =
   = .
 packages = find:
 zip_safe = False
+install_requires = 
+  attrs >= 19.1.0, < 20.0
+  flat-tree == 0.0.1a3 # TODO(decentral1se): use bounds when 0.0.1 lands
 
 [options.packages.find]
 where = .

test/conftest.py

@@ -0,0 +1,22 @@
+import hashlib
+
+import pytest
+
+
+@pytest.fixture
+def leaf():
+    def _leaf(node):
+        return hashlib.sha256(leaf.data).hexdigest()
+
+    return _leaf
+
+
+@pytest.fixture
+def parent():
+    def _parent(left, right):
+        sha256 = hashlib.sha256()
+        sha256.update(left)
+        sha256.update(right)
+        return sha256.hexdigest()
+
+    return _parent

test/test_generator.py

@@ -0,0 +1,66 @@
+"""Generator test module."""
+
+import hashlib
+
+from merkle_tree_stream import MerkleTreeGenerator, MerkleTreeNode
+
+
+def test_hashes(leaf, parent):
+    stream = MerkleTreeGenerator(leaf=leaf, parent=parent)
+
+    stream.next(b'a')
+
+    first_node = (
+        MerkleTreeNode(
+            index=0,
+            parent=1,
+            hash=hashlib.sha256(b'a').hexdigest(),
+            size=1,
+            data=b'a',
+        ),
+    )
+
+    stream.next(b'b')
+
+    second_node = (
+        MerkleTreeNode(
+            index=2,
+            parent=1,
+            hash=hashlib.sha256(b'b').hexdigest(),
+            size=1,
+            data=b'a',
+        ),
+    )
+
+    stream.next(b'c')
+
+    third = hashlib.sha256(b'a')
+    third.update(b'b')
+    third_hash = third.hexdigest()
+
+    third_node = (
+        MerkleTreeNode(index=1, parent=3, hash=third_hash, size=2, data=b'a'),
+    )
+
+    assert stream.nodes == [first_node, second_node, third_node]
+
+
+def test_single_root(leaf, parent):
+    stream = MerkleTreeGenerator(leaf=leaf, parent=parent)
+
+    stream.next('a')
+    stream.next('b')
+    stream.next('c')
+    stream.next('d')
+
+    assert stream.roots.length == 1
+
+
+def multiple_roots(leaf, parent):
+    stream = MerkleTreeGenerator(leaf=leaf, parent=parent)
+
+    stream.next('a')
+    stream.next('b')
+    stream.next('c')
+
+    assert stream.roots.length > 1