From 3682ae498649091e582da88a209629e7ceec16e7 Mon Sep 17 00:00:00 2001 From: Martin Wendt Date: Sun, 3 Nov 2024 22:01:34 +0100 Subject: [PATCH] Generic data (#13) * Update fixture.py * Define a concept * Initial refactoring * Fix warnings --- CHANGELOG.md | 3 +- docs/jupyter/take_the_tour.ipynb | 141 ++++++++++++------ docs/sphinx/take_the_tour.md | 117 ++++++++++----- nutree/common.py | 2 +- nutree/diff.py | 4 +- nutree/dot.py | 2 +- nutree/node.py | 242 +++++++++++++++---------------- nutree/rdf.py | 2 +- nutree/tree.py | 143 ++++++++++++++---- nutree/tree_generator.py | 2 +- nutree/typed_tree.py | 97 ++++--------- tests/fixture.py | 16 +- tests/test_typing.py | 118 +++++++++++++++ tests/test_typing_concept.py | 146 ++++++++++++++++--- tests/tutorial_script.py | 87 ----------- 15 files changed, 699 insertions(+), 423 deletions(-) create mode 100644 tests/test_typing.py delete mode 100644 tests/tutorial_script.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 19b0c00..a2c987e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,8 @@ - tree.to_rdf() is now available for Tree (not only TypedTree). - New method `node.up()` allows method chaining when adding nodes. - Pass pyright 'typeCheckingMode = "standard"'. -- Use generic typing for improved type checking. +- Use generic typing for improved type checking, e.g. use `tree = Tree[Animals]()` + to create a type-aware container. ## 0.9.0 (2024-09-12) diff --git a/docs/jupyter/take_the_tour.ipynb b/docs/jupyter/take_the_tour.ipynb index 3e6ed04..532fc60 100644 --- a/docs/jupyter/take_the_tour.ipynb +++ b/docs/jupyter/take_the_tour.ipynb @@ -140,13 +140,13 @@ "output_type": "stream", "text": [ "Tree<'Organization'>\n", - "├── <__main__.Department object at 0x109bbf6b0>\n", - "│ ├── <__main__.Department object at 0x1078baa50>\n", - "│ │ ╰── <__main__.Person object at 0x109cb23c0>\n", - "│ ╰── <__main__.Person object at 0x109cb2300>\n", - "├── <__main__.Department object at 0x1078b8f20>\n", - "│ ╰── <__main__.Person object at 0x109cb2060>\n", - "╰── <__main__.Person object at 0x109cb1f70>\n" + "├── <__main__.Department object at 0x114bc4140>\n", + "│ ├── <__main__.Department object at 0x114bc4440>\n", + "│ │ ╰── <__main__.Person object at 0x114bc4740>\n", + "│ ╰── <__main__.Person object at 0x114bc46b0>\n", + "├── <__main__.Department object at 0x114bc45f0>\n", + "│ ╰── <__main__.Person object at 0x114bc47a0>\n", + "╰── <__main__.Person object at 0x114bc4590>\n" ] } ], @@ -221,7 +221,7 @@ { "data": { "text/plain": [ - "Node<'Person', data_id=278704631>" + "Node<'Person', data_id=290178137>" ] }, "execution_count": 6, @@ -267,9 +267,9 @@ { "data": { "text/plain": [ - "[Node<'Person', data_id=278704700>,\n", - " Node<'Department', data_id=276347122>,\n", - " Node<'Person', data_id=278704631>]" + "[Node<'Person', data_id=290178164>,\n", + " Node<'Department', data_id=290178143>,\n", + " Node<'Person', data_id=290178137>]" ] }, "execution_count": 8, @@ -294,7 +294,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 9, "metadata": {}, "outputs": [ { @@ -302,8 +302,8 @@ "output_type": "stream", "text": [ "Tree<'Organization'>\n", - "╰── Node<'Department', data_id=8e6ae0d7-b268-4c9d-8da8-bdc5cdcb4f8d>\n", - " ╰── Node<'Person', data_id=ad20de1f-34c4-40cc-a102-51c28fb6ec5b>\n" + "╰── Node<'Department', data_id=1145f7c6-9a12-4110-a9f7-f84c4d53df77>\n", + " ╰── Node<'Person', data_id=a5adfa10-ebbb-4949-9638-6a99dab6165b>\n" ] } ], @@ -331,7 +331,7 @@ { "data": { "text/plain": [ - "Node<'Person', data_id=ad20de1f-34c4-40cc-a102-51c28fb6ec5b>" + "Node<'Person', data_id=a5adfa10-ebbb-4949-9638-6a99dab6165b>" ] }, "execution_count": 10, @@ -354,7 +354,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 11, "metadata": {}, "outputs": [ { @@ -396,7 +396,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 13, "metadata": {}, "outputs": [ { @@ -459,9 +459,9 @@ "output_type": "stream", "text": [ "Tree<\"Copy of Tree<'Organization'>\">\n", - "├── Node<'Department', data_id=278642539>\n", - "│ ╰── Node<'Department', data_id=276347557>\n", - "╰── Node<'Department', data_id=276347122>\n" + "├── Node<'Department', data_id=290178068>\n", + "│ ╰── Node<'Department', data_id=290178116>\n", + "╰── Node<'Department', data_id=290178143>\n" ] } ], @@ -487,8 +487,8 @@ "output_type": "stream", "text": [ "Tree<\"Copy of Tree<'Organization'>\">\n", - "├── Node<'Department', data_id=278642539>\n", - "╰── Node<'Department', data_id=276347122>\n" + "├── Node<'Department', data_id=290178068>\n", + "╰── Node<'Department', data_id=290178143>\n" ] } ], @@ -567,14 +567,14 @@ "name": "stdout", "output_type": "stream", "text": [ - "Node<'Department', data_id=278642539>\n", - "├── Node<'Department', data_id=276347557>\n", - "│ ╰── Node<'Person', data_id=278704700>\n", - "╰── Node<'Person', data_id=278704631>\n", - "Node<'Department', data_id=276347122>\n", - "├── Node<'Person', data_id=278704646>\n", - "╰── Node<'Person', data_id=278704688>\n", - "Node<'Person', data_id=278704631>\n" + "Node<'Department', data_id=290178068>\n", + "├── Node<'Department', data_id=290178116>\n", + "│ ╰── Node<'Person', data_id=290178164>\n", + "╰── Node<'Person', data_id=290178137>\n", + "Node<'Department', data_id=290178143>\n", + "├── Node<'Person', data_id=290178170>\n", + "╰── Node<'Person', data_id=290178155>\n", + "Node<'Person', data_id=290178137>\n" ] } ], @@ -591,8 +591,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "Node<'Person', data_id=278704631>, parent=None\n", - "Node<'Person', data_id=278704631>, parent=Node<'Department', data_id=278642539>\n" + "Node<'Person', data_id=290178137>, parent=None\n", + "Node<'Person', data_id=290178137>, parent=Node<'Department', data_id=290178068>\n" ] } ], @@ -621,7 +621,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Tree<'4459280720'>\n", + "Tree<'4642856144'>\n", "├── 'A'\n", "│ ├── 'a1'\n", "│ ╰── 'a2'\n", @@ -661,11 +661,11 @@ "name": "stdout", "output_type": "stream", "text": [ - "Tree<'4459281584'>\n", - "├── Node<'A', data_id=646631547239422124>\n", - "│ ╰── Node<\"DictWrapper<{'title': 'foo', 'id': 1}>\", data_id=4459448768>\n", - "╰── Node<'B', data_id=-8358617541791855429>\n", - " ╰── Node<\"DictWrapper<{'title': 'foo', 'id': 1}>\", data_id=4459448768>\n" + "Tree<'4626598512'>\n", + "├── Node<'A', data_id=5473560910289120105>\n", + "│ ╰── Node<\"DictWrapper<{'title': 'foo', 'id': 1}>\", data_id=4628284480>\n", + "╰── Node<'B', data_id=5825221258629577999>\n", + " ╰── Node<\"DictWrapper<{'title': 'foo', 'id': 1}>\", data_id=4628284480>\n" ] } ], @@ -702,7 +702,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "TypedTree<'4459281248'>\n", + "TypedTree<'4626500208'>\n", "╰── friend → Mia\n", " ├── brother → Noah\n", " ╰── sister → Olivia\n" @@ -739,10 +739,69 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Typing\n", + "# Type Hints\n", "\n", "Nutree comes fully typed (passing [pyright](https://microsoft.github.io/pyright/#/) \n", - "standard checks). This improves type-safety and auto-complete features in your IDE." + "standard checks). This improves type-safety and auto-complete features inside \n", + "IDEs.\n", + "\n", + "Assuming we have a hierrchy of objects:" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [], + "source": [ + "class OrgaEntry:\n", + " def __init__(self, name: str):\n", + " self.name: str = name\n", + " self.guid = uuid.uuid4()\n", + "\n", + "\n", + "class Person(OrgaEntry):\n", + " def __init__(self, name: str, age: int):\n", + " super().__init__(name)\n", + " self.age: int = age\n", + "\n", + "\n", + "class Department(OrgaEntry):\n", + " def __init__(self, name: str):\n", + " super().__init__(name)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now create a type-safe tree to store these objects:" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "b871d538-8505-464c-992f-a1bf6a781711\n" + ] + } + ], + "source": [ + "tree = Tree[OrgaEntry]()\n", + "\n", + "dev = tree.add(Department(\"Development\"))\n", + "alice = dev.add(Person(\"Alice\", 25))\n", + "\n", + "tree.add(42) # Flags a typing error in the IDE (but not at runtime)\n", + "\n", + "# IDE infers that the type `OrgaEntry` is stored in the `data` attribute\n", + "# and provides code completion\n", + "print(alice.data.guid)" ] } ], diff --git a/docs/sphinx/take_the_tour.md b/docs/sphinx/take_the_tour.md index 288195d..12577d4 100644 --- a/docs/sphinx/take_the_tour.md +++ b/docs/sphinx/take_the_tour.md @@ -90,13 +90,13 @@ tree.print() ``` Tree<'Organization'> - ├── <__main__.Department object at 0x109bbf6b0> - │ ├── <__main__.Department object at 0x1078baa50> - │ │ ╰── <__main__.Person object at 0x109cb23c0> - │ ╰── <__main__.Person object at 0x109cb2300> - ├── <__main__.Department object at 0x1078b8f20> - │ ╰── <__main__.Person object at 0x109cb2060> - ╰── <__main__.Person object at 0x109cb1f70> + ├── <__main__.Department object at 0x114bc4140> + │ ├── <__main__.Department object at 0x114bc4440> + │ │ ╰── <__main__.Person object at 0x114bc4740> + │ ╰── <__main__.Person object at 0x114bc46b0> + ├── <__main__.Department object at 0x114bc45f0> + │ ╰── <__main__.Person object at 0x114bc47a0> + ╰── <__main__.Person object at 0x114bc4590> Tree nodes store a reference to the object in the `node.data` attribute. @@ -134,7 +134,7 @@ tree[alice] - Node<'Person', data_id=278704631> + Node<'Person', data_id=290178137> @@ -158,9 +158,9 @@ tree.find_all(match=lambda node: "i" in node.data.name) - [Node<'Person', data_id=278704700>, - Node<'Department', data_id=276347122>, - Node<'Person', data_id=278704631>] + [Node<'Person', data_id=290178164>, + Node<'Department', data_id=290178143>, + Node<'Person', data_id=290178137>] @@ -181,8 +181,8 @@ tree_2.print(repr="{node}") ``` Tree<'Organization'> - ╰── Node<'Department', data_id=8e6ae0d7-b268-4c9d-8da8-bdc5cdcb4f8d> - ╰── Node<'Person', data_id=ad20de1f-34c4-40cc-a102-51c28fb6ec5b> + ╰── Node<'Department', data_id=1145f7c6-9a12-4110-a9f7-f84c4d53df77> + ╰── Node<'Person', data_id=a5adfa10-ebbb-4949-9638-6a99dab6165b> now we could also search by the GUID, for example: @@ -195,7 +195,7 @@ tree_2.find(data_id=str(bob.guid)) - Node<'Person', data_id=ad20de1f-34c4-40cc-a102-51c28fb6ec5b> + Node<'Person', data_id=a5adfa10-ebbb-4949-9638-6a99dab6165b> @@ -260,9 +260,9 @@ tree_copy.print(repr="{node}") ``` Tree<"Copy of Tree<'Organization'>"> - ├── Node<'Department', data_id=278642539> - │ ╰── Node<'Department', data_id=276347557> - ╰── Node<'Department', data_id=276347122> + ├── Node<'Department', data_id=290178068> + │ ╰── Node<'Department', data_id=290178116> + ╰── Node<'Department', data_id=290178143> In-place filtering is also available: @@ -274,8 +274,8 @@ tree_copy.print(repr="{node}") ``` Tree<"Copy of Tree<'Organization'>"> - ├── Node<'Department', data_id=278642539> - ╰── Node<'Department', data_id=276347122> + ├── Node<'Department', data_id=290178068> + ╰── Node<'Department', data_id=290178143> ## Mutation @@ -320,14 +320,14 @@ identical data_id: tree.print(repr="{node}", title=False) ``` - Node<'Department', data_id=278642539> - ├── Node<'Department', data_id=276347557> - │ ╰── Node<'Person', data_id=278704700> - ╰── Node<'Person', data_id=278704631> - Node<'Department', data_id=276347122> - ├── Node<'Person', data_id=278704646> - ╰── Node<'Person', data_id=278704688> - Node<'Person', data_id=278704631> + Node<'Department', data_id=290178068> + ├── Node<'Department', data_id=290178116> + │ ╰── Node<'Person', data_id=290178164> + ╰── Node<'Person', data_id=290178137> + Node<'Department', data_id=290178143> + ├── Node<'Person', data_id=290178170> + ╰── Node<'Person', data_id=290178155> + Node<'Person', data_id=290178137> @@ -336,8 +336,8 @@ for clone in tree.find_all(alice): print(f"{clone}, parent={clone.parent}") ``` - Node<'Person', data_id=278704631>, parent=None - Node<'Person', data_id=278704631>, parent=Node<'Department', data_id=278642539> + Node<'Person', data_id=290178137>, parent=None + Node<'Person', data_id=290178137>, parent=Node<'Department', data_id=290178068> ## Special Data Types @@ -356,7 +356,7 @@ tree_str.add("B") tree_str.print() ``` - Tree<'4459280720'> + Tree<'4642856144'> ├── 'A' │ ├── 'a1' │ ╰── 'a2' @@ -386,11 +386,11 @@ tree.print(repr="{node}") # tree.find(d) ``` - Tree<'4459281584'> - ├── Node<'A', data_id=646631547239422124> - │ ╰── Node<"DictWrapper<{'title': 'foo', 'id': 1}>", data_id=4459448768> - ╰── Node<'B', data_id=-8358617541791855429> - ╰── Node<"DictWrapper<{'title': 'foo', 'id': 1}>", data_id=4459448768> + Tree<'4626598512'> + ├── Node<'A', data_id=5473560910289120105> + │ ╰── Node<"DictWrapper<{'title': 'foo', 'id': 1}>", data_id=4628284480> + ╰── Node<'B', data_id=5825221258629577999> + ╰── Node<"DictWrapper<{'title': 'foo', 'id': 1}>", data_id=4628284480> ## Typed Trees @@ -410,7 +410,7 @@ typed_tree.add("Mia", kind="friend").add("Noah", kind="brother").up().add( typed_tree.print() ``` - TypedTree<'4459281248'> + TypedTree<'4626500208'> ╰── friend → Mia ├── brother → Noah ╰── sister → Olivia @@ -427,7 +427,48 @@ Even conversion to and SVG, PNG is possible: Read the [User Guide](https://nutree.readthedocs.io/en/latest/ug_serialize.html) for different methods to save, load, or convert a tree to different output formats. -# Typing +# Type Hints Nutree comes fully typed (passing [pyright](https://microsoft.github.io/pyright/#/) -standard checks). This improves type-safety and auto-complete features in your IDE. +standard checks). This improves type-safety and auto-complete features inside +IDEs. + +Assuming we have a hierrchy of objects: + + +```python +class OrgaEntry: + def __init__(self, name: str): + self.name: str = name + self.guid = uuid.uuid4() + + +class Person(OrgaEntry): + def __init__(self, name: str, age: int): + super().__init__(name) + self.age: int = age + + +class Department(OrgaEntry): + def __init__(self, name: str): + super().__init__(name) +``` + +We can now create a type-safe tree to store these objects: + + +```python +tree = Tree[OrgaEntry]() + +dev = tree.add(Department("Development")) +alice = dev.add(Person("Alice", 25)) + +tree.add(42) # Flags a typing error in the IDE (but not at runtime) + +# IDE infers that the type `OrgaEntry` is stored in the `data` attribute +# and provides code completion +print(alice.data.guid) +``` + + b871d538-8505-464c-992f-a1bf6a781711 + diff --git a/nutree/common.py b/nutree/common.py index 2e23e98..1c66a4c 100644 --- a/nutree/common.py +++ b/nutree/common.py @@ -121,7 +121,7 @@ def __init__(self, value=None): DataIdType = Union[str, int] #: Type of ``Tree(..., calc_data_id)``` -CalcIdCallbackType = Callable[["Tree[Any]", Any], DataIdType] +CalcIdCallbackType = Callable[["Tree", Any], DataIdType] #: Type of ``format(..., repr=)``` ReprArgType = Union[str, Callable[["Node"], str]] diff --git a/nutree/diff.py b/nutree/diff.py index 7aa7184..ea0ed9e 100644 --- a/nutree/diff.py +++ b/nutree/diff.py @@ -82,7 +82,9 @@ def diff_node_formatter(node): return s -def diff_tree(t0: Tree[Any], t1: Tree[Any], *, ordered=False, reduce=False) -> Tree: +def diff_tree( + t0: Tree[Any, Any], t1: Tree[Any, Any], *, ordered=False, reduce=False +) -> Tree: from nutree import Tree t2 = Tree(f"diff({t0.name!r}, {t1.name!r})") diff --git a/nutree/dot.py b/nutree/dot.py index e0183b0..0388b2c 100644 --- a/nutree/dot.py +++ b/nutree/dot.py @@ -111,7 +111,7 @@ def _attr_str(attr_def: dict, mapper=None, node=None): def tree_to_dotfile( - tree: Tree[Any], + tree: Tree[Any, Any], target: IO[str] | str | Path, *, format=None, diff --git a/nutree/node.py b/nutree/node.py index eab983b..877a9bf 100644 --- a/nutree/node.py +++ b/nutree/node.py @@ -13,6 +13,7 @@ IO, TYPE_CHECKING, Any, + Generic, Iterable, Iterator, # TypeVar, @@ -37,7 +38,6 @@ from nutree.common import ( CONNECTORS, - AmbiguousMatchError, DataIdType, DeserializeMapperType, FlatJsonDictType, @@ -63,14 +63,14 @@ from nutree.dot import node_to_dot from nutree.rdf import RDFMapperCallbackType, node_to_rdf -TNode = TypeVar("TNode", bound="Node", default="Node") -# TNode = TypeVar("TNode", bound="Node", default="Node", covariant=True) +TData = TypeVar("TData", bound="Any", default="Any") +TNode = TypeVar("TNode", bound="Node", default="Node[TData]") # ------------------------------------------------------------------------------ # - Node # ------------------------------------------------------------------------------ -class Node: +class Node(Generic[TData]): """ A Node represents a single element in the tree. It is a shallow wrapper around a user data instance, that adds navigation, @@ -121,14 +121,14 @@ class Node: def __init__( self, - data, + data: TData, *, parent: Self, data_id: DataIdType | None = None, node_id: int | None = None, meta: dict | None = None, ): - self._data = data + self._data: TData = data self._parent: Self = parent tree = parent._tree @@ -147,7 +147,7 @@ def __init__( self._meta = meta - tree._register(self) + tree._register(self) # type: ignore def __repr__(self) -> str: return f"{self.__class__.__name__}<{self.name!r}, data_id={self.data_id}>" @@ -179,16 +179,6 @@ def __getattr__(self, name: str) -> Any: return getattr(self, name[:-2])() raise AttributeError(repr(name)) - # def __iadd__(self, other) -> None: - # """Add child node(s).""" - # if isinstance(other, (Node, str)): - # self.add_child(other) - # elif isinstance(other, (list, tuple)): - # for o in other: - # self += o - # else: - # raise NotImplementedError - # Do not define __len__: we don't want leaf nodes to evaluate as falsy # def __len__(self) -> int: # raise NotImplementedError("Use `len(node.data)` or `len(node._children)`.") @@ -244,7 +234,7 @@ def children(self) -> list[Self]: return [] if c is None else c @property - def data(self) -> Any: + def data(self) -> TData: """Return the wrapped data instance (use :meth:`~nutree.tree.Tree.set_data()` to modify).""" return self._data @@ -319,77 +309,92 @@ def set_data( with_clones: bool | None = None, ) -> None: """Change node's `data` and/or `data_id` and update bookkeeping.""" - if not data and not data_id: - raise ValueError("Missing data or data_id") - - tree = self._tree - - if data is None or data is self._data: - new_data = None - else: - new_data = data - if data_id is None: - data_id = tree.calc_data_id(data) - - if data_id is None or data_id == self._data_id: - new_data_id = None - else: - new_data_id = data_id - - node_map = tree._nodes_by_data_id - cur_nodes = node_map[self._data_id] - has_clones = len(cur_nodes) > 1 - - if has_clones and with_clones is None: - raise AmbiguousMatchError( - "set_data() for clones requires `with_clones` decision" - ) + return self.tree._set_data( + self, # type: ignore + data, + data_id=data_id, + with_clones=with_clones, + ) - if new_data_id: - # data_id (and possibly data) changes: we have to update the map - if has_clones: - if with_clones: - # Move the whole slot (but check if new id already exist) - prev_clones = node_map[self._data_id] - del node_map[self._data_id] - try: # are we adding to existing clones now? - node_map[new_data_id].extend(prev_clones) - except KeyError: # still a singleton, just a new data_id - node_map[new_data_id] = prev_clones - for n in prev_clones: - n._data_id = new_data_id - if new_data: - n._data = new_data - else: - # Move this one node to another slot in the map - node_map[self._data_id].remove(self) - try: # are we adding to existing clones again? - node_map[new_data_id].append(self) - except KeyError: # now a singleton with a new data_id - node_map[new_data_id] = [self] - self._data_id = new_data_id - if new_data: - self._data = new_data - else: - # data_id (and possibly data) changed for a *single* node - del node_map[self._data_id] - try: # are we creating a clone now? - node_map[new_data_id].append(self) - except KeyError: # still a singleton, just a new data_id - node_map[new_data_id] = [self] - self._data_id = new_data_id - if new_data: - self._data = new_data - elif new_data: - # `data` changed, but `data_id` remains the same: - # simply replace the reference - if with_clones: - for n in cur_nodes: - n._data = data - else: - self._data = new_data + # def set_data( + # self, + # data, + # *, + # data_id: DataIdType | None = None, + # with_clones: bool | None = None, + # ) -> None: + # """Change node's `data` and/or `data_id` and update bookkeeping.""" + # if not data and not data_id: + # raise ValueError("Missing data or data_id") + + # tree = self.tree + + # if data is None or data is self._data: + # new_data = None + # else: + # new_data = data + # if data_id is None: + # data_id = tree.calc_data_id(data) - return + # if data_id is None or data_id == self._data_id: + # new_data_id = None + # else: + # new_data_id = data_id + + # node_map = tree._nodes_by_data_id + # cur_nodes = node_map[self._data_id] + # has_clones = len(cur_nodes) > 1 + + # if has_clones and with_clones is None: + # raise AmbiguousMatchError( + # "set_data() for clones requires `with_clones` decision" + # ) + + # if new_data_id: + # # data_id (and possibly data) changes: we have to update the map + # if has_clones: + # if with_clones: + # # Move the whole slot (but check if new id already exist) + # prev_clones = node_map[self._data_id] + # del node_map[self._data_id] + # try: # are we adding to existing clones now? + # node_map[new_data_id].extend(prev_clones) + # except KeyError: # still a singleton, just a new data_id + # node_map[new_data_id] = prev_clones + # for n in prev_clones: + # n._data_id = new_data_id + # if new_data: + # n._data = new_data + # else: + # # Move this one node to another slot in the map + # node_map[self._data_id].remove(self) + # try: # are we adding to existing clones again? + # node_map[new_data_id].append(self) + # except KeyError: # now a singleton with a new data_id + # node_map[new_data_id] = [self] + # self._data_id = new_data_id + # if new_data: + # self._data = new_data + # else: + # # data_id (and possibly data) changed for a *single* node + # del node_map[self._data_id] + # try: # are we creating a clone now? + # node_map[new_data_id].append(self) + # except KeyError: # still a singleton, just a new data_id + # node_map[new_data_id] = [self] + # self._data_id = new_data_id + # if new_data: + # self._data = new_data + # elif new_data: + # # `data` changed, but `data_id` remains the same: + # # simply replace the reference + # if with_clones: + # for n in cur_nodes: + # n._data = data + # else: + # self._data = new_data + + # return def get_children(self) -> list[Self]: """Return list of direct child nodes (list may be empty).""" @@ -433,7 +438,7 @@ def last_sibling(self) -> Self: def get_clones(self, *, add_self=False) -> list[Self]: """Return a list of all nodes that reference the same data if any.""" - clones = self._tree._nodes_by_data_id[self._data_id] + clones = cast(list[Self], self._tree._nodes_by_data_id[self._data_id]) if add_self: return clones.copy() return [n for n in clones if n is not self] @@ -564,7 +569,7 @@ def get_path(self, *, add_self=True, separator="/", repr="{node.name}") -> str: def add_child( self, - child: Self | Tree | Any, + child: Self | Tree | TData, *, before: Self | bool | int | None = None, deep: bool | None = None, @@ -598,7 +603,7 @@ def add_child( - : prepend the new node before this child node Args: - child (Node|Tree|Any): + child (Node|Tree|TData): Either an existing Node or a data object. before (bool|int|Node|None): Optional position. @@ -622,7 +627,7 @@ def add_child( if isinstance(child, self._tree.__class__): if deep is None: deep = True - topnodes = child._root.children + topnodes = cast(list[Self], child.system_root.children) if isinstance(before, (int, Node)) or before is True: topnodes.reverse() n = None @@ -630,6 +635,7 @@ def add_child( self.add_child(n, before=before, deep=deep) return cast(Self, n) # need to return a node + factory = self.tree.node_factory source_node: Self = None # type: ignore new_node: Self = None # type: ignore @@ -654,22 +660,16 @@ def add_child( if data_id and data_id != source_node._data_id: raise UniqueConstraintError(f"data_id conflict: {source_node}") - # If creating an inherited node, use the parent class as constructor - # child_class = child.__class__ - - new_node = cast( - Self, - self.tree.node_factory( - source_node.data, parent=self, data_id=data_id, node_id=node_id - ), + new_node = factory( + source_node.data, # type: ignore + parent=self, # type: ignore + data_id=data_id, + node_id=node_id, ) else: - new_node = cast( - Self, - self.tree.node_factory( - child, parent=self, data_id=data_id, node_id=node_id - ), - ) + new_node = factory(child, parent=self, data_id=data_id, node_id=node_id) # type: ignore + + new_node = cast(Self, new_node) if before is True: before = 0 # prepend @@ -701,7 +701,7 @@ def add_child( def append_child( self, - child: Self | Tree | Any, + child: Self | Tree | TData, *, deep=None, data_id: DataIdType | None = None, @@ -717,7 +717,7 @@ def append_child( def prepend_child( self, - child: Self | Tree | Any, + child: Self | Tree | TData, *, deep=None, data_id: DataIdType | None = None, @@ -737,7 +737,7 @@ def prepend_child( def prepend_sibling( self, - child: Self | Tree | Any, + child: Self | Tree | TData, *, deep=None, data_id: DataIdType | None = None, @@ -753,7 +753,7 @@ def prepend_sibling( def append_sibling( self, - child: Self | Tree | Any, + child: Self | Tree | TData, *, deep=None, data_id: DataIdType | None = None, @@ -770,7 +770,7 @@ def append_sibling( def move_to( self, - new_parent: Self | Tree[Self], + new_parent: Self | Tree[TData, Self], *, before: Self | bool | int | None = None, ): @@ -783,9 +783,9 @@ def move_to( if not isinstance(new_parent, Node): # it's a Tree assert isinstance(new_parent, self.tree.__class__) - new_parent = new_parent._root + new_parent = new_parent.system_root - if new_parent._tree is not self._tree: + if new_parent.tree is not self.tree: raise NotImplementedError("Can only move nodes inside same tree") self._parent._children.remove(self) # type: ignore @@ -836,34 +836,34 @@ def remove(self, *, keep_children=False, with_clones=False) -> None: if not pc: # store None instead of `[]` pc = self._parent._children = None - self._tree._unregister(self) + self._tree._unregister(self) # type: ignore def remove_children(self) -> None: """Remove all children of this node, making it a leaf node.""" _unregister = self._tree._unregister for n in self._iter_post(): - _unregister(n) + _unregister(n) # type: ignore self._children = None return def copy( self, *, add_self=True, predicate: PredicateCallbackType | None = None - ) -> Tree[Self]: + ) -> Tree[TData, Self]: """Return a new :class:`~nutree.tree.Tree` instance from this branch. See also :ref:`iteration-callbacks`. """ - new_tree = self._tree.__class__() + new_tree = cast("Tree[TData, Self]", self._tree.__class__()) if add_self: root = new_tree.add(self) else: - root = new_tree._root + root = new_tree.system_root root._add_from(self, predicate=predicate) return new_tree def copy_to( self, - target: Self | Tree[Self], + target: Self | Tree[TData, Self], *, add_self=True, before: Self | bool | int | None = None, @@ -884,7 +884,7 @@ def copy_to( if add_self: res = target.add_child(self, before=before, deep=deep) return cast(Self, res) # if target is Tree, type is not inferred? - # return target.add_child(self, before=before, deep=deep) + assert before is None if not self._children: raise ValueError("Need child nodes when `add_self=False`") @@ -974,7 +974,7 @@ def _visit(other: Self) -> None: pass return - def filtered(self, predicate: PredicateCallbackType) -> Tree[Self]: + def filtered(self, predicate: PredicateCallbackType) -> Tree[TData, Self]: """Return a filtered copy of this node and descendants as tree. See also :ref:`iteration-callbacks`. @@ -1096,7 +1096,7 @@ def visit( *, add_self=False, method: IterMethod = IterMethod.PRE_ORDER, - memo: Any = None, + memo: Any | None = None, ) -> None | Any: """Call `callback(node, memo)` for all subnodes. diff --git a/nutree/rdf.py b/nutree/rdf.py index 26f95fe..a8aaf84 100644 --- a/nutree/rdf.py +++ b/nutree/rdf.py @@ -163,7 +163,7 @@ def node_to_rdf( def tree_to_rdf( - tree: Tree[Any], + tree: Tree[Any, Any], *, node_mapper: RDFMapperCallbackType | None = None, ) -> Graph: diff --git a/nutree/tree.py b/nutree/tree.py index c43cca7..7d5ff3e 100644 --- a/nutree/tree.py +++ b/nutree/tree.py @@ -59,7 +59,7 @@ MermaidFormatType, MermaidNodeMapperCallbackType, ) -from nutree.node import Node, TNode +from nutree.node import Node, TData, TNode from nutree.rdf import tree_to_rdf _DELETED_TAG = "" @@ -77,7 +77,7 @@ class _SystemRootNode(Node): """Invisible system root node.""" def __init__(self, tree: Tree) -> None: - self._tree: Tree = tree # type: ignore + self._tree: Tree = tree self._parent = None # type: ignore self._node_id = ROOT_NODE_ID self._data_id = ROOT_DATA_ID @@ -89,7 +89,7 @@ def __init__(self, tree: Tree) -> None: # ------------------------------------------------------------------------------ # - Tree # ------------------------------------------------------------------------------ -class Tree(Generic[TNode]): +class Tree(Generic[TData, TNode]): """ A Tree object is a shallow wrapper around a single, invisible system root node. All visible toplevel nodes are direct children of this root node. @@ -105,7 +105,7 @@ class Tree(Generic[TNode]): **Note:** Use with care, see also :ref:`forward-attributes`. """ - node_factory: Type[Node] = Node + node_factory: Type[TNode] = cast(Type[TNode], Node) root_node_factory = _SystemRootNode #: Default connector prefixes ``format(style=...)`` argument. @@ -129,7 +129,7 @@ def __init__( self._lock = threading.RLock() #: Tree name used for logging self.name: str = str(id(self) if name is None else name) - self._root: TNode = self.root_node_factory(self) # type: ignore + self._root: Node = self.root_node_factory(self) # type: ignore self._node_by_id: dict[int, TNode] = {} self._nodes_by_data_id: dict[DataIdType, list[TNode]] = {} # Optional callback that calculates data_ids from data objects @@ -197,10 +197,6 @@ def __getitem__(self, data: object) -> TNode: ) return res[0] - # def __iadd__(self, other) -> None: - # """Support `tree += node(s)` syntax""" - # self._root += other - def __len__(self): """Make ``len(tree)`` return the number of nodes (also makes empty trees falsy).""" @@ -219,7 +215,7 @@ def calc_data_id(self, data: Any) -> DataIdType: repeated invocations of Python. """ if self._calc_data_id_hook: - return self._calc_data_id_hook(self, data) + return self._calc_data_id_hook(self, data) # type: ignore return hash(data) def _register(self, node: TNode) -> None: @@ -278,16 +274,16 @@ def _unregister(self, node: TNode, *, clear: bool = True) -> None: def children(self) -> list[TNode]: """Return list of direct child nodes, i.e. toplevel nodes (list may be empty).""" - return self._root.children + return self.system_root.children def get_toplevel_nodes(self) -> list[TNode]: """Return list of direct child nodes, i.e. toplevel nodes (may be empty, alias for :meth:`children`).""" - return self._root.children + return self.system_root.children @property def system_root(self) -> TNode: - return self._root + return cast(TNode, self._root) @property def count(self) -> int: @@ -305,6 +301,85 @@ def count_unique(self) -> int: """ return len(self._nodes_by_data_id) + def _set_data( + self, + node: TNode, + data: TData, + *, + data_id: DataIdType | None = None, + with_clones: bool | None = None, + ) -> None: + """Change node's `data` and/or `data_id` and update bookkeeping.""" + if not data and not data_id: + raise ValueError("Missing data or data_id") + + if data is None or data is node._data: + new_data = None + else: + new_data = data + if data_id is None: + data_id = self.calc_data_id(data) + + if data_id is None or data_id == node._data_id: + new_data_id = None + else: + new_data_id = data_id + + node_map = self._nodes_by_data_id + cur_nodes = node_map[node._data_id] + has_clones = len(cur_nodes) > 1 + + if has_clones and with_clones is None: + raise AmbiguousMatchError( + "set_data() for clones requires `with_clones` decision" + ) + + if new_data_id: + # data_id (and possibly data) changes: we have to update the map + if has_clones: + if with_clones: + # Move the whole slot (but check if new id already exist) + prev_clones = node_map[node._data_id] + del node_map[node._data_id] + try: # are we adding to existing clones now? + node_map[new_data_id].extend(prev_clones) + except KeyError: # still a singleton, just a new data_id + node_map[new_data_id] = prev_clones + for n in prev_clones: + n._data_id = new_data_id + if new_data: + n._data = new_data + else: + # Move this one node to another slot in the map + node_map[node._data_id].remove(node) + try: # are we adding to existing clones again? + node_map[new_data_id].append(node) + except KeyError: # now a singleton with a new data_id + node_map[new_data_id] = [node] + node._data_id = new_data_id + if new_data: + node._data = new_data + else: + # data_id (and possibly data) changed for a *single* node + del node_map[node._data_id] + try: # are we creating a clone now? + node_map[new_data_id].append(node) + except KeyError: # still a singleton, just a new data_id + node_map[new_data_id] = [node] + node._data_id = new_data_id + if new_data: + node._data = new_data + elif new_data: + # `data` changed, but `data_id` remains the same: + # simply replace the reference + if with_clones: + for n in cur_nodes: + n._data = data + else: + node._data = new_data + + return + @classmethod def serialize_mapper(cls, node: Node, data: dict) -> dict | None: """Used as default `mapper` argument for :meth:`save`.""" @@ -319,11 +394,11 @@ def deserialize_mapper(cls, parent: Node, data: dict) -> str | object | None: def first_child(self) -> TNode | None: """Return the first toplevel node.""" - return self._root.first_child() + return self.system_root.first_child() def last_child(self) -> TNode | None: """Return the last toplevel node.""" - return self._root.last_child() + return self.system_root.last_child() def get_random_node(self) -> TNode: """Return a random node. @@ -335,7 +410,7 @@ def get_random_node(self) -> TNode: def calc_height(self) -> int: """Return the maximum depth of all nodes.""" - return self._root.calc_height() + return self.system_root.calc_height() def visit( self, @@ -349,7 +424,9 @@ def visit( See Node's :meth:`~nutree.node.Node.visit` method and :ref:`iteration-callbacks` for details. """ - return self._root.visit(callback, add_self=False, method=method, memo=memo) + return self.system_root.visit( + callback, add_self=False, method=method, memo=memo + ) def iterator(self, method: IterMethod = IterMethod.PRE_ORDER) -> Iterator[TNode]: """Traverse tree structure and yield nodes. @@ -362,7 +439,7 @@ def iterator(self, method: IterMethod = IterMethod.PRE_ORDER) -> Iterator[TNode] values = list(self._node_by_id.values()) random.shuffle(values) return (n for n in values) - return self._root.iterator(method=method) + return self.system_root.iterator(method=method) #: Implement ``for node in tree: ...`` syntax to iterate nodes depth-first. __iter__ = iterator @@ -376,7 +453,9 @@ def format_iter( if title: yield f"{self}" if title is True else f"{title}" has_title = title is not False - yield from self._root.format_iter(repr=repr, style=style, add_self=has_title) + yield from self.system_root.format_iter( + repr=repr, style=style, add_self=has_title + ) def format( self, *, repr: ReprArgType | None = None, style=None, title=None, join="\n" @@ -405,7 +484,7 @@ def print( def add_child( self, - child: TNode | Self | Any, + child: TNode | Self | TData, *, before: TNode | bool | int | None = None, deep: bool | None = None, @@ -416,7 +495,7 @@ def add_child( See Node's :meth:`~nutree.node.Node.add_child` method for details. """ - return self._root.add_child( + return self.system_root.add_child( child, before=before, deep=deep, @@ -448,7 +527,7 @@ def copy( name = f"Copy of {self}" new_tree = self.__class__(name) with self: - new_tree._root._add_from(self._root, predicate=predicate) + new_tree.system_root._add_from(self.system_root, predicate=predicate) return new_tree def copy_to(self, target: TNode | Self, *, deep=True) -> None: @@ -457,14 +536,14 @@ def copy_to(self, target: TNode | Self, *, deep=True) -> None: See Node's :meth:`~nutree.node.Node.copy_to` method for details. """ with self: - self._root.copy_to(target, add_self=False, before=None, deep=deep) + self.system_root.copy_to(target, add_self=False, before=None, deep=deep) def filter(self, predicate: PredicateCallbackType) -> None: """In-place removal of unmatching nodes. See also :ref:`iteration-callbacks`. """ - self._root.filter(predicate=predicate) + self.system_root.filter(predicate=predicate) def filtered(self, predicate: PredicateCallbackType) -> Self: """Return a filtered copy of this tree. @@ -477,7 +556,7 @@ def filtered(self, predicate: PredicateCallbackType) -> Self: def clear(self) -> None: """Remove all nodes from this tree.""" - self._root.remove_children() + self.system_root.remove_children() def find_all( self, @@ -504,7 +583,7 @@ def find_all( return [] elif match is not None: - return self._root.find_all(match=match, max_results=max_results) + return self.system_root.find_all(match=match, max_results=max_results) raise NotImplementedError @@ -534,7 +613,7 @@ def find_first( return res[0] if res else None elif match is not None: assert node_id is None - return self._root.find_first(match=match) + return self.system_root.find_first(match=match) elif node_id is not None: return self._node_by_id.get(node_id) raise NotImplementedError @@ -548,14 +627,14 @@ def sort(self, *, key: SortKeyType | None = None, reverse=False, deep=True) -> N `key` defaults to ``attrgetter("name")``, so children are sorted by their string representation. """ - self._root.sort_children(key=key, reverse=reverse, deep=deep) + self.system_root.sort_children(key=key, reverse=reverse, deep=deep) def to_dict_list(self, *, mapper: SerializeMapperType | None = None) -> list[dict]: """Call Node's :meth:`~nutree.node.Node.to_dict` method for all child nodes and return list of results.""" res = [] with self: - for n in self._root._children: # type: ignore[reportOptionalIterable] + for n in self.system_root._children: # type: ignore[reportOptionalIterable] res.append(n.to_dict(mapper=mapper)) return res @@ -568,7 +647,7 @@ def from_dict(cls, obj: list[dict], *, mapper=None) -> Self: :ref:`iteration-callbacks`. """ new_tree = cls() - new_tree._root.from_dict(obj, mapper=mapper) + new_tree.system_root.from_dict(obj, mapper=mapper) return new_tree def to_list_iter( @@ -579,7 +658,7 @@ def to_list_iter( value_map: ValueMapType | None = None, ) -> Iterator[tuple[int, Union[FlatJsonDictType, str]]]: """Yield a parent-referencing list of child nodes.""" - yield from self._root.to_list_iter( + yield from self.system_root.to_list_iter( mapper=mapper, key_map=key_map, value_map=value_map ) @@ -685,7 +764,7 @@ def _from_list( mapper: DeserializeMapperType | None = None, ) -> Self: tree = cls() # Tree or TypedTree - node_idx_map: dict[int, TNode] = {0: tree._root} + node_idx_map: dict[int, TNode] = {0: tree.system_root} if mapper is None: mapper = cls.deserialize_mapper diff --git a/nutree/tree_generator.py b/nutree/tree_generator.py index f57567d..73ec6b9 100644 --- a/nutree/tree_generator.py +++ b/nutree/tree_generator.py @@ -391,7 +391,7 @@ def _make_tree( return -def build_random_tree(*, tree_class: type[Tree[Any]], structure_def: dict) -> Tree: +def build_random_tree(*, tree_class: type[Tree[Any, Any]], structure_def: dict) -> Tree: """ Return a nutree.TypedTree with random data from a specification. See :ref:`randomize` for details. diff --git a/nutree/typed_tree.py b/nutree/typed_tree.py index f00a5e8..b3d9143 100644 --- a/nutree/typed_tree.py +++ b/nutree/typed_tree.py @@ -9,15 +9,14 @@ from collections import Counter from pathlib import Path -from typing import IO, Any, Iterator, Type, cast +from typing import IO, Iterator, Type, cast # typing.Self requires Python 3.11 -from typing_extensions import Self +from typing_extensions import Any, Self from nutree.common import ( ROOT_DATA_ID, ROOT_NODE_ID, - CalcIdCallbackType, DataIdType, DeserializeMapperType, KeyMapType, @@ -28,7 +27,7 @@ call_mapper, sentinel, ) -from nutree.node import Node +from nutree.node import Node, TData from nutree.tree import Tree @@ -43,7 +42,7 @@ class TAnyKind: # ------------------------------------------------------------------------------ # - TypedNode # ------------------------------------------------------------------------------ -class TypedNode(Node): +class TypedNode(Node[TData]): """ A special node variant, derived from :class:`~nutree.node.Node` and used by :class:`~nutree.typed_tree.TypedTree`. @@ -64,14 +63,14 @@ class TypedNode(Node): def __init__( self, kind: str, - data, + data: TData, *, parent: Self, data_id: DataIdType | None = None, node_id: int | None = None, meta: dict | None = None, ): - self._kind = kind # tree._register() checks for this attribute + self._kind: str = kind # tree._register() checks for this attribute super().__init__( data, parent=parent, data_id=data_id, node_id=node_id, meta=meta ) @@ -86,34 +85,10 @@ def __repr__(self) -> str: f"{self.name}, data_id={self.data_id!r}>" ) - # @property - # def name(self) -> str: - # """String representation of the embedded `data` object with kind.""" - # # return f"{self._kind} → {self.data}" - # # Inspired by clarc notation: http://www.jclark.com/xml/xmlns.htm - # # return f"{{{self._kind}}}:{self.data}" - # return f"{self.data}" - @property def kind(self) -> str: return self._kind - # @property - # def parent(self) -> TypedNode | None: - # """Return parent node or None for toplevel nodes.""" - # p = self._parent - # return p if p._parent else None - - # @property - # def children(self) -> list[TypedNode]: - # """Return list of direct child nodes (list may be empty). - - # Note that this property returns all children, independent of the kind. - # See also :meth:`get_children`. - # """ - # c = self._children - # return [] if c is None else c - def get_children(self, kind: str | ANY_KIND) -> list[Self]: """Return list of direct child nodes of a given type (list may be empty).""" all_children = self._children @@ -123,12 +98,6 @@ def get_children(self, kind: str | ANY_KIND) -> list[Self]: return all_children return list(filter(lambda n: n._kind == kind, all_children)) - # def set_data( - # self, kind: str, data, *, data_id=None, with_clones: bool = None - # ) -> None: - # """Change node's `data` and/or `data_id` and update bookkeeping.""" - # super().set_data(data, data_id=data_id, with_clones=with_clones) - def first_child(self, kind: str | ANY_KIND) -> Self | None: """First direct child node or None if no children exist.""" all_children = self._children @@ -238,7 +207,7 @@ def is_last_sibling(self, *, any_kind=False) -> bool: def add_child( self, - child: Self | TypedTree | Any, + child: Self | TypedTree | TData, *, kind: str | None = None, before: Self | bool | int | None = None, @@ -262,12 +231,12 @@ def add_child( if isinstance(child, TypedTree): if deep is None: deep = True - topnodes = child._root.children + topnodes = cast(list[Self], child.system_root.children) if isinstance(before, (int, Node)) or before is True: topnodes.reverse() for n in topnodes: self.add_child(n, kind=n.kind, before=before, deep=deep) - return child._root # type: ignore + return child.system_root # type: ignore source_node: Self = None # type: ignore new_node: Self = None # type: ignore @@ -333,7 +302,7 @@ def add_child( def append_child( self, - child: Self | TypedTree | Any, + child: Self | TypedTree | TData, *, kind: str | None = None, deep: bool | None = None, @@ -355,7 +324,7 @@ def append_child( def prepend_child( self, - child: Self | TypedTree | Any, + child: Self | TypedTree | TData, *, kind: str | None = None, deep: bool | None = None, @@ -377,7 +346,7 @@ def prepend_child( def prepend_sibling( self, - child: Self | TypedTree | Any, + child: Self | TypedTree | TData, *, kind: str | None = None, deep=None, @@ -394,7 +363,7 @@ def prepend_sibling( def append_sibling( self, - child: Self | TypedTree | Any, + child: Self | TypedTree | TData, *, kind: str | None = None, deep: bool | None = None, @@ -482,13 +451,13 @@ def __init__(self, tree: TypedTree) -> None: self._data = tree.name self._children = [] self._meta = None - self._kind = None + self._kind = None # type: ignore # ------------------------------------------------------------------------------ # - TypedTree # ------------------------------------------------------------------------------ -class TypedTree(Tree[TypedNode]): +class TypedTree(Tree[TData, TypedNode[TData]]): """ A special tree variant, derived from :class:`~nutree.tree.Tree`, that uses :class:`~nutree.typed_tree.TypedNode` objects, which maintain @@ -507,20 +476,6 @@ class TypedTree(Tree[TypedNode]): #: Default value for ``add_child`` when loading. DEFAULT_CHILD_TYPE = "child" - def __init__( - self, - name: str | None = None, - *, - calc_data_id: CalcIdCallbackType | None = None, - forward_attrs: bool = False, - ): - super().__init__( - name, - calc_data_id=calc_data_id, - forward_attrs=forward_attrs, - ) - # self._root = _SystemRootTypedNode(self) - @classmethod def deserialize_mapper(cls, parent: Node, data: dict) -> str | object | None: """Used as default `mapper` argument for :meth:`load`.""" @@ -534,19 +489,19 @@ def deserialize_mapper(cls, parent: Node, data: dict) -> str | object | None: def add_child( self, - child: TypedNode | Self | Any, + child: TypedNode[TData] | Self | TData, *, kind: str | None = None, - before: TypedNode | bool | int | None = None, + before: TypedNode[TData] | bool | int | None = None, deep: bool | None = None, data_id: DataIdType | None = None, node_id: int | None = None, - ) -> TypedNode: + ) -> TypedNode[TData]: """Add a toplevel node. See Node's :meth:`~nutree.node.Node.add_child` method for details. """ - return self._root.add_child( + return self.system_root.add_child( child, kind=kind, before=before, @@ -558,15 +513,15 @@ def add_child( #: Alias for :meth:`add_child` add = add_child # Must re-bind here - def first_child(self, kind: str | ANY_KIND) -> TypedNode | None: + def first_child(self, kind: str | ANY_KIND) -> TypedNode[TData] | None: """Return the first toplevel node.""" - return self._root.first_child(kind=kind) + return self.system_root.first_child(kind=kind) - def last_child(self, kind: str | ANY_KIND) -> TypedNode | None: + def last_child(self, kind: str | ANY_KIND) -> TypedNode[TData] | None: """Return the last toplevel node.""" - return self._root.last_child(kind=kind) + return self.system_root.last_child(kind=kind) - def iter_by_type(self, kind: str | ANY_KIND) -> Iterator[TypedNode]: + def iter_by_type(self, kind: str | ANY_KIND) -> Iterator[TypedNode[TData]]: if kind == ANY_KIND: return self.iterator() for n in self.iterator(): @@ -622,7 +577,7 @@ def _from_list( mapper = cls.deserialize_mapper # System root has index #0: - node_idx_map: dict[int, TypedNode] = {0: tree._root} + node_idx_map: dict[int, TypedNode[TData]] = {0: tree.system_root} # Start reading data lines starting at index #1: for idx, (parent_idx, data) in enumerate(obj, 1): @@ -630,7 +585,7 @@ def _from_list( # print(idx, parent_idx, data, parent) if isinstance(data, str): # This can only happen if the source was generated by a plain Tree - n = parent.add(data, kind=cls.DEFAULT_CHILD_TYPE) + n = parent.add(data, kind=cls.DEFAULT_CHILD_TYPE) # type: ignore elif isinstance(data, int): first_clone = node_idx_map[data] n = parent.add( diff --git a/tests/fixture.py b/tests/fixture.py index 63dfdad..d12130e 100644 --- a/tests/fixture.py +++ b/tests/fixture.py @@ -61,7 +61,12 @@ def __repr__(self) -> str: def create_tree( - *, style="simple", name="fixture", clones=False, tree=None, print=True + *, + style="simple", + name="fixture", + clones=False, + tree: Tree | None = None, + print=True, ) -> Tree: if tree is not None: assert not tree, "must be empty" @@ -123,7 +128,12 @@ def create_tree( def create_typed_tree( - *, style="simple", name="fixture", clones=False, tree=None, print=True + *, + style="simple", + name="fixture", + clones=False, + tree: TypedTree | None = None, + print=True, ) -> TypedTree: if tree is not None: assert not tree, "must be empty" @@ -282,7 +292,7 @@ def check_content( ignore_tree_name=True, repr: ReprArgType | None = None, style=None, -): +) -> bool: err = _check_content(tree, expect_ascii, msg, ignore_tree_name, repr, style) if err: raise AssertionError(err) from None diff --git a/tests/test_typing.py b/tests/test_typing.py new file mode 100644 index 0000000..99d7f0a --- /dev/null +++ b/tests/test_typing.py @@ -0,0 +1,118 @@ +# ruff: noqa: T201, T203 `print` found +# type: ignore +from __future__ import annotations + +from uuid import UUID, uuid4 + +from nutree.tree import Tree +from nutree.typed_tree import ANY_KIND, TypedTree +from typing_extensions import reveal_type + + +# --- Sample data ------------------------------------------------------------ +class OrgaEntry: + def __init__(self, name: str): + self.name: str = name + self.guid: UUID = uuid4() + + +class Person(OrgaEntry): + def __init__(self, name: str, age: int): + super().__init__(name) + self.age: int = age + + +class Department(OrgaEntry): + def __init__(self, name: str): + super().__init__(name) + + +# --- Test ------------------------------------------------------------------- + + +class TestTreeTyping: + def test_tree(self): + tree = Tree() + n = tree.add("top") + n.add("child") + tree.add(42) + n.add(42) + tree.first_child().add("child2") + + reveal_type(tree.system_root) + reveal_type(n) + reveal_type(tree.first_child()) + reveal_type(tree.first_child().data) + + def test_str_tree(self): + tree = Tree[str]() + + n = tree.add("child") + n.add("child2") + n.add(42) + tree.add(42) + + reveal_type(tree.system_root) + reveal_type(n) + reveal_type(tree.first_child()) + reveal_type(tree.first_child().data) + + def test_orga_tree(self): + tree = Tree[OrgaEntry]() + + dev = tree.add(Department("Development")) + alice = dev.add(Person("Alice", 42)) + tree.add(42) + alice.add(42) + + reveal_type(tree.system_root) + reveal_type(alice) + reveal_type(tree.first_child()) + reveal_type(tree.first_child().data) + reveal_type(alice) + reveal_type(alice.data) + + +class TestTypedTreeTyping: + def test_typed_tree(self): + tree = TypedTree() + + n = tree.add("child", kind="child") + n.add("child2", kind="child") + tree.add(42, kind="child") + + tree.first_child(kind=ANY_KIND).add("child3", kind="child") + + reveal_type(tree.system_root) + reveal_type(n) + reveal_type(tree.first_child(kind=ANY_KIND)) + reveal_type(tree.first_child(kind=ANY_KIND).data) + + def test_typed_tree_str(self): + tree = TypedTree[str]() + + n = tree.add("child", kind="child") + n.add("child2", kind="child") + n.add(42, kind="child") + tree.add(42, kind="child") + + tree.first_child(kind=ANY_KIND).add("child3", kind="child") + + reveal_type(tree.system_root) + reveal_type(n) + reveal_type(tree.first_child(kind=ANY_KIND)) + reveal_type(tree.first_child(kind=ANY_KIND).data) + + def test_typed_tree_orga(self): + tree = TypedTree[OrgaEntry]() + + dev = tree.add(Department("Development"), kind="department") + alice = dev.add(Person("Alice", 42), kind="member") + tree.add(42, kind="child") + alice.add(42, kind="child") + + reveal_type(alice) + reveal_type(tree.first_child(kind=ANY_KIND)) + reveal_type(tree.first_child(kind=ANY_KIND).data) + reveal_type(alice) + reveal_type(alice.data) diff --git a/tests/test_typing_concept.py b/tests/test_typing_concept.py index 5ed4063..231d3fb 100644 --- a/tests/test_typing_concept.py +++ b/tests/test_typing_concept.py @@ -1,37 +1,46 @@ # ruff: noqa: T201, T203 `print` found # pyright: reportIncompatibleMethodOverride=false +# mypy: ignore[override] + +# type: ignore from __future__ import annotations -from typing import Any, Generic, List, Type, cast +from typing import Generic, List, Type, cast +from uuid import UUID, uuid4 -from typing_extensions import Self, TypeVar +from typing_extensions import Any, Self, TypeVar, reveal_type -TNode = TypeVar("TNode", bound="Node", default="Node") +TData = TypeVar("TData", bound="Any", default="Any") +TNode = TypeVar("TNode", bound="Node", default="Node[TData]") -class Node: - def __init__(self, data: Any, parent: Self): - self.data: Any = data +class Node(Generic[TData]): + def __init__(self, data: TData, parent: Self): + self.data: TData = data self.parent: Self = parent self.children: List[Self] = [] - def add(self, data: Any) -> Self: + def add(self, data: TData) -> Self: node = self.__class__(data, self) self.children.append(node) return node -class Tree(Generic[TNode]): +class Tree(Generic[TData, TNode]): node_factory: Type[TNode] = cast(Type[TNode], Node) def __init__(self): - self.root: TNode = self.node_factory("__root__", cast(TNode, None)) + self._root: Node = self.node_factory("__root__", None) # type: ignore - def add(self, data: Any) -> TNode: + def add(self, data: TData) -> TNode: node = self.root.add(data) return node + @property + def root(self) -> TNode: + return cast(TNode, self._root) + def first(self) -> TNode: return self.root.children[0] @@ -39,44 +48,133 @@ def first(self) -> TNode: # ---------------------------- -class TypedNode(Node): - def __init__(self, data: Any, kind: str, parent: Self): +class TypedNode(Node[TData]): + def __init__(self, data: TData, kind: str, parent: Self): super().__init__(data, parent) self.kind: str = kind # self.children: List[TypedNode] = [] - def add(self, data: Any, kind: str) -> Self: + def add(self, data: TData, kind: str) -> Self: node = self.__class__(data, kind, self) self.children.append(node) return node -class TypedTree(Tree[TypedNode]): +class TypedTree(Tree[TData, TypedNode[TData]]): node_factory = TypedNode def __init__(self): - self.root: TypedNode = self.node_factory( - "__root__", "__root__", cast(TypedNode, None) - ) + self._root = TypedNode("__root__", "__root__", None) # type: ignore - def add(self, data: Any, kind: str) -> TypedNode: + def add(self, data: TData, kind: str) -> TypedNode[TData]: node = self.root.add(data, kind) return node -class TestTypingSelf: +# --- Sample data ------------------------------------------------------------ +class OrgaEntry: + def __init__(self, name: str): + self.name: str = name + self.guid: UUID = uuid4() + + +class Person(OrgaEntry): + def __init__(self, name: str, age: int): + super().__init__(name) + self.age: int = age + + +class Department(OrgaEntry): + def __init__(self, name: str): + super().__init__(name) + + +# --- Test ------------------------------------------------------------------- + + +class TestTreeTyping: def test_tree(self): - # tree0 = Tree() - tree = Tree[Node]() + tree = Tree() n = tree.add("top") n.add("child") + tree.add(42) + n.add(42) tree.first().add("child2") + reveal_type(tree.root) + reveal_type(n) + reveal_type(tree.first()) + reveal_type(tree.first().data) + + def test_str_tree(self): + tree = Tree[str]() + + n = tree.add("child") + n.add("child2") + n.add(42) + tree.add(42) + + reveal_type(tree.root) + reveal_type(n) + reveal_type(tree.first()) + reveal_type(tree.first().data) + + def test_orga_tree(self): + tree = Tree[OrgaEntry]() + + dev = tree.add(Department("Development")) + alice = dev.add(Person("Alice", 42)) + tree.add(42) + alice.add(42) + + reveal_type(tree.root) + reveal_type(alice) + reveal_type(tree.first()) + reveal_type(tree.first().data) + reveal_type(alice) + reveal_type(alice.data) + + +class TestTypedTreeTyping: def test_typed_tree(self): tree = TypedTree() - tree.add("child", kind="child") - tree.add("child2", kind="child") + n = tree.add("child", kind="child") + n.add("child2", kind="child") + tree.add(42, kind="child") - # tree.first().children tree.first().add("child3", kind="child") + + reveal_type(tree.root) + reveal_type(n) + reveal_type(tree.first()) + reveal_type(tree.first().data) + + def test_typed_tree_str(self): + tree = TypedTree[str]() + + n = tree.add("child", kind="child") + n.add("child2", kind="child") + n.add(42, kind="child") + tree.add(42, kind="child") + + tree.first().add("child3", kind="child") + + reveal_type(tree.root) + reveal_type(n) + reveal_type(tree.first()) + reveal_type(tree.first().data) + + def test_typed_tree_orga(self): + tree = TypedTree[OrgaEntry]() + + dev = tree.add(Department("Development"), kind="department") + alice = dev.add(Person("Alice", 42), kind="member") + tree.add(42, kind="child") + alice.add(42, kind="child") + + reveal_type(alice) + reveal_type(tree.first()) + reveal_type(tree.first().data) + reveal_type(alice) + reveal_type(alice.data) diff --git a/tests/tutorial_script.py b/tests/tutorial_script.py deleted file mode 100644 index 9c0d9bf..0000000 --- a/tests/tutorial_script.py +++ /dev/null @@ -1,87 +0,0 @@ -# (c) 2021-2024 Martin Wendt; see https://github.com/mar10/nutree -# Licensed under the MIT license: https://www.opensource.org/licenses/mit-license.php -""" """ -# ruff: noqa: T201, T203 `print` found -# pyright: reportAttributeAccessIssue=false -# pyright: reportOptionalMemberAccess=false - -import json - -from nutree import Node, Tree - - -class Item: - def __init__(self, name, price, count): - self.name = name - self.price = float(price) - self.count = int(count) - - def __repr__(self): - return f"Item<{self.name!r}, {self.price:.2f}$>" - - -if __name__ == "__main__": - tree = Tree("Store") - - n = tree.add("Records") - n.add("Let It Be") - n.add("Get Yer Ya-Ya's Out!") - n = tree.add("Books") - n.add("The Little Prince") - - tree.print() - - print(tree.format(repr="{node}", style="lines32", title="My Store")) - - assert tree.count == 5 - - record_node = tree["Records"] - assert isinstance(record_node, Node) - assert tree.first_child() is record_node - - assert len(record_node.children) == 2 - assert record_node.depth() == 1 - - assert tree.find("Records") is record_node - assert tree.find("records") is None - - n = record_node.first_child() - assert record_node.find("Let It Be") is n - - assert n.name == "Let It Be" - assert n.depth() == 2 - assert n.parent is record_node - assert n.prev_sibling() is None - assert n.next_sibling().name == "Get Yer Ya-Ya's Out!" - assert not n.children - - res = tree.find_all(match=r"[GL]et.*") - print(res) - assert len(res) == 2 - - res = tree.find_all(match=lambda n: "y" in n.name.lower()) - assert len(res) == 1 - - # Note that `find()` is an alias for `find_first()` - res = tree.find_first(match=r"[GL]et.*") - assert res.name == "Let It Be" - assert tree._self_check() - - tree = Tree("multi") - a = tree.add("A") - a.add("a1") - a.add("a2") - b = tree.add("B") - b.add("b1") - b.add("a2") - tree.print() - print(tree.find("a2")) - print(tree.find_all("a2")) - - res = tree.dumps("dict") - - res = json.dumps(tree.to_dict_list()) - print(res) - - res = list(tree.to_list_iter()) - print(res)