Add iterator(..., kind) to TypedTree (#18)

* Update take_the_tour.ipynb

* Add TypedTree.iterator(..., kind)` and `TypedNode.iterator(..., kind)`

* Add count_descendants(kind)

* Update ruff to 0.9

* More tests

Author: mar10

.pre-commit-config.yaml

@@ -1,7 +1,7 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.6.6
+    rev: v0.9.5
     hooks:
       # Run the linter.
       - id: ruff

CHANGELOG.md

@@ -1,7 +1,12 @@
 # Changelog
 
-## 1.0.1 (unreleased)
+## 1.1.0 (unreleased)
 
+- DEPRECATE: `TypedTree.iter_by_type()`. Use `iterator(.., kind)`instead.
+- New methods `TypedTree.iterator(..., kind=ANY_KIND)`, 
+  `TypedNode.iterator(..., kind=ANY_KIND)`,
+  and `TypedTree.count_descendants(leaves_only=False, kind=ANY_KIND)`
+ 
 ## 1.0.0 (2024-12-27)
 - Add benchmarks (using [Benchman](https://github.com/mar10/benchman)).
 - Drop support for Python 3.8

Pipfile

@@ -17,7 +17,7 @@ pytest = "*"
 pytest-cov = "*"
 PyYAML = "*"
 rdflib = "*"
-ruff = "*"
+ruff = "~=0.9"
 setuptools = ">=42.0"
 Sphinx = "*"
 sphinx_rtd_theme = "*"

Pipfile.lock

@@ -22,7 +22,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
     {
@@ -68,7 +68,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -103,7 +103,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -132,21 +132,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
       "Tree<'Organization'>\n",
-      "├── <__main__.Department object at 0x118c3fbc0>\n",
-      "│   ├── <__main__.Department object at 0x118646c30>\n",
-      "│   │   ╰── <__main__.Person object at 0x118c68170>\n",
-      "│   ╰── <__main__.Person object at 0x118c3f770>\n",
-      "├── <__main__.Department object at 0x118c3ffe0>\n",
-      "│   ╰── <__main__.Person object at 0x118c68bf0>\n",
-      "╰── <__main__.Person object at 0x118c3ffb0>\n"
+      "├── <__main__.Department object at 0x106a66e40>\n",
+      "│   ├── <__main__.Department object at 0x106a65610>\n",
+      "│   │   ╰── <__main__.Person object at 0x10c02e4e0>\n",
+      "│   ╰── <__main__.Person object at 0x10c02ede0>\n",
+      "├── <__main__.Department object at 0x106a64170>\n",
+      "│   ╰── <__main__.Person object at 0x10c02e5a0>\n",
+      "╰── <__main__.Person object at 0x10c02e150>\n"
      ]
     }
    ],
@@ -183,7 +183,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -215,16 +215,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 31,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "Node<'Person<Alice (25)>', data_id=294404091>"
+       "Node<'Person<Alice (25)>', data_id=281030165>"
       ]
      },
-     "execution_count": 31,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -243,16 +243,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<__main__.Person at 0x118c3ffb0>"
+       "<__main__.Person at 0x10c02e150>"
       ]
      },
-     "execution_count": 32,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -276,23 +276,42 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "[Node<'Department<Development>', data_id=294404028>,\n",
-       " Node<'Department<Marketing>', data_id=294404094>]"
+       "'Department<Test>'"
       ]
      },
-     "execution_count": 33,
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "str(tree[claire].parent.data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[Node<'Department<Development>', data_id=275408612>,\n",
+       " Node<'Department<Marketing>', data_id=275407895>]"
+      ]
+     },
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "# tree[alice].parent.data\n",
     "tree[alice].get_siblings()"
    ]
   },

docs/jupyter/take_the_tour.ipynb

@@ -356,7 +356,7 @@ Navigation methods are type-aware now::
     assert cause1.get_index() == 0
     assert cause1.get_index(any_kind=True) == 2
 
-    assert len(list(tree.iter_by_type("effect"))) == 3
+    assert len(list(tree.iterator(kind="effect"))) == 3
 
 Keep in mind that a tree node is unique within a tree, but may reference identical
 data objects, so these `clones` could exist at different locations of tree. 

docs/sphinx/ug_graphs.rst

@@ -1022,9 +1022,9 @@ def _self_check(self) -> Literal[True]:
             # assert node._data_id == self.calc_data_id(node.data), node
             assert node._data_id in self._nodes_by_data_id, node
             assert node._node_id == id(node), f"{node}: {node._node_id} != {id(node)}"
-            assert (
-                node._children is None or len(node._children) > 0
-            ), f"{node}: {node._children}"
+            assert node._children is None or len(node._children) > 0, (
+                f"{node}: {node._children}"
+            )
 
         assert len(self._node_by_id) == len(node_list)
 

nutree/tree.py

@@ -43,9 +43,9 @@ class Randomizer(ABC):
     """
 
     def __init__(self, *, probability: float = 1.0) -> None:
-        assert (
-            isinstance(probability, float) and 0.0 <= probability <= 1.0
-        ), f"probality must be in the range [0.0 .. 1.0]: {probability}"
+        assert isinstance(probability, float) and 0.0 <= probability <= 1.0, (
+            f"probality must be in the range [0.0 .. 1.0]: {probability}"
+        )
         self.probability = probability
 
     def _skip_value(self) -> bool:
@@ -84,9 +84,9 @@ def __init__(
         none_value: Any = None,
     ) -> None:
         super().__init__(probability=probability)
-        assert type(min_val) is type(
-            max_val
-        ), f"min_val and max_val must be of the same type: {min_val}, {max_val}"
+        assert type(min_val) is type(max_val), (
+            f"min_val and max_val must be of the same type: {min_val}, {max_val}"
+        )
         self.is_float = isinstance(min_val, float)
         self.min = min_val
         self.max = max_val
@@ -129,19 +129,19 @@ def __init__(
     ) -> None:
         super().__init__(probability=probability)
         assert isinstance(min_dt, date), f"min_dt must be a date: {min_dt}"
-        assert isinstance(
-            max_dt, (date, int)
-        ), f"max_dt must be a date or int: {max_dt}"
+        assert isinstance(max_dt, (date, int)), (
+            f"max_dt must be a date or int: {max_dt}"
+        )
 
         if isinstance(max_dt, int):
             self.delta_days = max_dt
             max_dt = min_dt + timedelta(days=self.delta_days)
         else:
             self.delta_days = (max_dt - min_dt).days
 
-        assert (
-            max_dt > min_dt
-        ), f"max_dt must be greater than min_dt: {min_dt}, {max_dt}"
+        assert max_dt > min_dt, (
+            f"max_dt must be greater than min_dt: {min_dt}, {max_dt}"
+        )
 
         self.min = min_dt
         self.max = max_dt

nutree/tree_generator.py

@@ -27,6 +27,7 @@
     ROOT_NODE_ID,
     DataIdType,
     DeserializeMapperType,
+    IterMethod,
     KeyMapType,
     MapperCallbackType,
     PredicateCallbackType,
@@ -137,12 +138,44 @@ def last_child(self, kind: str | type[ANY_KIND]) -> Self | None:
                 return n
         return None
 
+    def iterator(
+        self,
+        method: IterMethod = IterMethod.PRE_ORDER,
+        *,
+        add_self=False,
+        kind: str | type[ANY_KIND] = ANY_KIND,
+    ) -> Iterator[TypedNode[TData]]:
+        """Return an iterator that walks the tree in the specified order."""
+        if kind is ANY_KIND:
+            yield from super().iterator(method=method, add_self=add_self)
+            return
+
+        if add_self and self.kind == kind:
+            yield self
+        for n in super().iterator(method=method, add_self=False):
+            if n.kind == kind:
+                yield n
+        return
+
     def has_children(self, kind: str | type[ANY_KIND]) -> bool:
         """Return true if this node has one or more children."""
         if kind is ANY_KIND:
             return bool(self._children)
         return len(self.get_children(kind)) > 1
 
+    def count_descendants(
+        self, *, leaves_only=False, kind: str | type[ANY_KIND] = ANY_KIND
+    ) -> int:
+        """Return number of descendant nodes, not counting self."""
+        if kind is ANY_KIND:
+            return super().count_descendants(leaves_only=leaves_only)
+        all = not leaves_only
+        i = 0
+        for node in self.iterator():
+            if (all or not node._children) and node.kind == kind:
+                i += 1
+        return i
+
     def get_siblings(self, *, add_self=False, any_kind=False) -> list[Self]:
         """Return a list of all sibling entries of self (excluding self) if any."""
         if any_kind:
@@ -630,13 +663,30 @@ def last_child(self, kind: str | type[ANY_KIND]) -> TypedNode[TData] | None:
         return self.system_root.last_child(kind=kind)
 
     def iter_by_type(self, kind: str | type[ANY_KIND]) -> Iterator[TypedNode[TData]]:
+        """@deprecated: Use :meth:`iterator` with `kind` argument instead."""
+        yield from self.iterator(kind=kind)
+
+    def iterator(
+        self,
+        method: IterMethod = IterMethod.PRE_ORDER,
+        *,
+        kind: str | type[ANY_KIND] = ANY_KIND,
+    ) -> Iterator[TypedNode[TData]]:
         if kind == ANY_KIND:
-            yield from self.iterator()
-        for n in self.iterator():
+            yield from super().iterator(method=method)
+            return
+
+        for n in super().iterator(method=method):
             if n._kind == kind:
                 yield n
         return
 
+    def count_descendants(
+        self, *, leaves_only=False, kind: str | type[ANY_KIND] = ANY_KIND
+    ) -> int:
+        """Return number of nodes, optionally restricted to type."""
+        return self.system_root.count_descendants(leaves_only=leaves_only, kind=kind)
+
     def save(
         self,
         target: IO[str] | str | Path,

nutree/typed_tree.py

@@ -285,9 +285,9 @@ class FrozenItem:
 
         assert isinstance(dict_node.data, FrozenItem)
         assert dict_node.data is item, "dataclass should be stored as reference"
-        assert (
-            dict_node.price == 12.34
-        ), "should support attribute access via forwardinging"
+        assert dict_node.price == 12.34, (
+            "should support attribute access via forwardinging"
+        )
         with pytest.raises(AttributeError):
             _ = dict_node.foo
 

tests/test_objects.py

@@ -64,7 +64,7 @@ def test_typed_tree(self):
         assert cause1.get_index() == 0
         assert cause1.get_index(any_kind=True) == 2
 
-        assert len(list(tree.iter_by_type("effect"))) == 3
+        assert len(list(tree.iterator(kind="effect"))) == 3
 
         # tree.print()
         # print()

tests/test_rdf.py

@@ -95,7 +95,21 @@ def test_add_child(self):
 
         assert cause2.parent is fail1
 
+        assert tree.count == 8
+        assert tree.count_descendants() == 8
+        tree.print()
+        assert tree.count_descendants(leaves_only=True) == 6
+        assert tree.count_descendants(kind="cause") == 2
+        assert tree.count_descendants(leaves_only=True, kind="failure") == 1
+        assert tree.system_root.count_descendants(kind="failure") == 2
+
         assert len(list(tree.iter_by_type("cause"))) == 2
+        assert len(list(tree.iterator(kind="cause"))) == 2
+        assert len(list(tree.iterator())) == 8
+
+        assert len(list(tree.system_root.iterator(add_self=True))) == 9
+        assert len(list(tree.system_root.iterator(kind="cause"))) == 2
+        assert len(list(tree.system_root.iterator(kind="cause", add_self=True))) == 2
 
         assert cause2.get_children("undefined") == []