Merge branch 'main' of https://github.com/mar10/nutree

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

@@ -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()