Focus on Iterator protocol - tests passing!

Too much naming copy/pasta will probably confuse Python programmers when
we have generator/iterators and I suppose we have streams but that is
just so generic that it won't help.

So, better make it explicit that this is an iterator.

It also seems to the fit the usage of the other implementors.

merkle_tree_stream/__init__.py

@@ -1,7 +1,7 @@
 """merkle-tree-stream module."""
 
-from merkle_tree_stream.generator import (  # noqa
-    MerkleTreeGenerator,
+from merkle_tree_stream.generate import (  # noqa
+    MerkleTreeIterator,
     MerkleTreeNode,
 )
 

merkle_tree_stream/generate.py

@@ -1,13 +1,13 @@
 """The merkle tree stream generator."""
 
-from typing import Any, Callable, List, Optional
+from typing import Any, Callable, Iterator, List, Optional
 
 import attr
 from flat_tree import FlatTreeAccessor
 
 Hash = str
 
-__all__ = ['MerkleTreeGenerator', 'MerkleTreeNode']
+__all__ = ['MerkleTreeIterator', 'MerkleTreeNode']
 
 flat_tree = FlatTreeAccessor()
 
@@ -26,7 +26,7 @@ class MerkleTreeNode:
     index: int
     parent: int
     size: int
-    data: Optional[bytes]
+    data: bytes
     hash: Optional[str] = None
 
     def __attrs_post_init__(self) -> Any:
@@ -35,29 +35,62 @@ class MerkleTreeNode:
 
 
 @attr.s(auto_attribs=True)
-class MerkleTreeGenerator:
-    """A stream that generates a merkle tree based on the incoming data.
+class MerkleTreeIterator:
+    """A merkle tree iterator based on incoming data.
 
     :param leaf: The leaf hash generation function
     :param parent: The parent hash generation function
     :param roots: The tree roots
     """
 
-    leaf: Callable[[MerkleTreeNode, List[MerkleTreeNode]], Hash]
-    parent: Callable[[MerkleTreeNode, List[MerkleTreeNode]], Hash]
+    leaf: Callable[[MerkleTreeNode], Hash]
+    parent: Callable[[MerkleTreeNode, MerkleTreeNode], Hash]
     roots: List[MerkleTreeNode] = attr.Factory(list)
 
-    def next(
-        self, data: bytes, nodes: Optional[List[MerkleTreeNode]] = None
-    ) -> List[MerkleTreeNode]:
-        """Further generate the tree based on the incoming data.
+    _position: int = 0
+    _nodes: List[MerkleTreeNode] = attr.Factory(list)
 
-        :param data: Incoming data
-        :param nodes: Pre-existing nodes
+    def __attrs_post_init__(self) -> Any:
+        """Initialise parent and block defaults."""
+        try:
+            index = self.roots[len(self.roots) - 1].index
+        except IndexError:
+            index = 0
+
+        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)
+
+    def __iter__(self) -> Iterator:
+        """The iterator initialisation."""
+        return self
+
+    def __next__(self) -> MerkleTreeNode:
+        """The following node."""
+        try:
+            node = self._nodes[self._position]
+        except IndexError:
+            raise StopIteration
+
+        self._position += 1
+
+        return node
+
+    def __len__(self) -> int:
+        """The number of nodes stored in the tree."""
+        return len(self._nodes)
+
+    def write(self, data: bytes):
+        """Write a new node to the tree.
+
+        :param data: The new tree data
         """
-        nodes = nodes or []
+        index = 2 * self.blocks
 
-        index = 2 * (self.blocks + 1)
+        self.blocks += 1
 
         leaf_node = MerkleTreeNode(
             index=index,
@@ -66,11 +99,10 @@ class MerkleTreeGenerator:
             data=data,
             size=len(data),
         )
-
-        leaf_node.hash = self.leaf(leaf_node, self.roots)
+        leaf_node.hash = self.leaf(leaf_node)
 
         self.roots.append(leaf_node)
-        nodes.append(leaf_node)
+        self._nodes.append(leaf_node)
 
         while len(self.roots) > 1:
             left = self.roots[len(self.roots) - 2]
@@ -84,23 +116,11 @@ class MerkleTreeGenerator:
             new_node = MerkleTreeNode(
                 index=left.parent,
                 parent=flat_tree.parent(left.parent),
-                hash=self.parent(left, [right]),
+                hash=self.parent(left, right),
                 size=left.size + right.size,
-                data=None,
+                data=b'',
             )
 
             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)
+            self._nodes.append(new_node)