diff --git a/README.md b/README.md index 73cd41e..dfe9e78 100644 --- a/README.md +++ b/README.md @@ -1,16 +1,20 @@ -This Python package contains a few abstract base classes for tree data structures. -Trees are very common data structure that represents a hierarchy of common nodes. -There are many different ways to represent them. -This package tries to provide a uniform interface, mixin methods and some utility functions without settling on a concrete tree implementation. +Trees are a common data structure and there are many different ways to implement them. +This package provides a common interface to access and operate on these objects. -## Abstract base classes ## +## Installing ## + +Use [pip](https://pip.pypa.io/en/stable/getting-started/) to install abstracttree: -```python -from abstracttree import DownTree, to_mermaid - -to_mermaid(DownTree) +```sh +$ pip install --upgrade abstracttree ``` +## Usage ## + +You can start by implementing the mixins below. Otherwise, a lot of trees are supported out of the box. + +### Mixins ### + ```mermaid graph TD; Tree[Tree]; @@ -31,26 +35,21 @@ BinaryDownTree-->BinaryTree; Tree-->BinaryTree; ``` -A `Downtree` needs to have links to its direct children, but doesn't require a link to its parent. -A `Tree` needs to have links to both its `children` and its `parent`. +| ABC | Inherits from | Abstract Methods | Mixin Methods | +|-------------------|---------------------------|---------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `DownTree` | | `children` | `nodes`, `nodes.preorder()`, `nodes.postorder()`, `nodes.levelorder()`, `descendants`, `leaves`, `levels`, `levels.zigzag()`, `is_leaf`, `transform()`, `nid` | +| `Tree` | `DownTree` | `parent` | `root`, `is_root`, `ancestors`, `path`, `siblings` | +| `MutableDownTree` | `DownTree` | `add_child()`, `remove_child()` | `add_children()` | +| `MutableTree` | `MutableDownTree`, `Tree` | | `detach()` | +| `BinaryDownTree` | `DownTree` | `left_child`, `right_child` | `children`, `nodes.inorder()`, `descendants.inorder()` | +| `BinaryTree` | `BinaryDownTree`, `Tree` | | | -| ABC | Inherits from | Abstract Methods | Mixin Methods | -|-------------------|---------------------------|---------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `DownTree` | | `children` | `nodes`, `nodes.preorder()`, `nodes.postorder()`, `nodes.levelorder()`, `descendants`, `leaves`, `levels`, `levels.zigzag()`, `is_leaf`, `transform()`, `nid`, `eqv()` | -| `Tree` | `DownTree` | `parent` | `root`, `is_root`, `ancestors`, `path`, `siblings` | -| `MutableDownTree` | `DownTree` | `add_child()`, `remove_child()` | `add_children()` | -| `MutableTree` | `Tree`, `MutableDownTree` | | `detach()` | -| `BinaryDownTree` | `DownTree` | `left_child`, `right_child` | `children`, `nodes.inorder()`, `descendants.inorder()` | -| `BinaryTree` | `BinaryDownTree`, `Tree` | | | - -In your own code, you can inherit from these trees. -For example, if your tree only has links to children: +For example, to create a simple tree with children (but no parent): ```python -import abstracttree -from abstracttree import print_tree +from abstracttree import DownTree, print_tree -class MyTree(abstracttree.DownTree): +class MyTree(DownTree): def __init__(self, value, children=()): self.value = value self._children = children @@ -71,116 +70,118 @@ print_tree(tree) # └─ MyTree 3 ``` -## Adapter ## +### Generics ## -In practice, not all existing tree data structures implement one of these abstract classes. -As a bridge, you can use `AbstractTree.convert` to convert these trees to a `Tree` instance. -However, whenever possible, it's recommended to inherit from `Tree` directly for minimal overhead. +Unfortunately, not all trees inherit from the mixins above. Yet, some objects still have treelike behaviour. +Therefore, AbstractTree provides support for a slightly weaker protocol. -Examples: +The following objects are `TreeLike`: +- All objects that support `obj.children` and `obj.parent`. +- Builtins classes `pathlib.Path` and `zipfile.Path`. +- Third party tree classes from [anytree](https://github.com/c0fec0de/anytree), [bigtree](https://github.com/kayjan/bigtree), [itertree](https://github.com/BR1py/itertree) and [littletree](https://github.com/lverweijen/littletree). -```python -# Trees from built-ins and standard library -tree = Tree.convert(int) -tree = Tree.convert(ast.parse("1 + 1 == 2")) -tree = Tree.convert(pathlib.Path("abstracttree")) +The following objects are `DownTreeLike`: +- All objects that support `obj.children`. +- Anything implementing `DownTree`. +- Recursive collections like lists, tuples, sets, dicts. This can be useful when dealing with json-data. -# Anything that has parent and children attributes (anytree / bigtree / littletree) -tree = Tree.convert(anytree.Node('name')) +This can be tested using `isinstance`: -# Nested list -tree = Tree.convert([[1, 2, 3], [4, 5, 6]]) +```python +isinstance(Path(r"C:\\Windows\System"), TreeLike) # True +isinstance(range(100), DownTreeLike) # True +isinstance(range(100), TreeLike) # False +isinstance(5, DownTreeLike) # False +isinstance("some text", DownTreeLike) # False (even though it might be considered a collection by python). ``` -Or use `astree` if you need a custom function for `parent` or `children`: +### Basic functions +On downtreelikes: ```python -# Tree from json-data -data = {"name": "a", - "children": [ - {"name": "b", "children": []}, - {"name": "c", "children": []} -]} -astree(data, children=operator.itemgetter["children"]) +children(node) # Children of node +label(node) # String representation of node (similar to str, but output excludes parent and children) +nid(node) # Address of node (similar to id, but supports delegates). +eqv(node1, node2) # Check if 2 nodes have the same identity (similar to is, but supports delegates) +``` -# pyqt.QtWidget -astree(widget, children=lambda w: w.children(), parent = lambda w: w.parent()) +Additionally, on treelikes: +```python +parent(node) # Parent of node or None if node is root of its own tree. +root(node) # Find root of this tree. +``` -# Tree from treelib -astree(tree.root, children=lambda nid: tree.children(nid), parent=lambda nid: tree.parent(nid)) +Examples: +```python +>>> from abstracttree import * +>>> children([1, 2, 3]) +[1, 2, 3] +>>> children({"name": "Philip", "children": ["Pete", "Mariam"]}) +[MappingItem(key="name", value="Philip"), MappingItem(key="children", value=["Pete", "Miriam"])] +>>> parent(Path(r"C:\\Windows\System")) +Path(r"C:\\Windows") +>>> label(Path(r"C:\\Windows\System")) +"System" +>>> eqv(Path(r"C:\\Windows\System"), Path(r"C:\\Windows\System")) +True +>>> eqv([1, 2, 3], [1, 2, 3]) +False +``` -# itertree -astree(tree, children=iter, parent=lambda t: t.parent) +### Iterators -# Infinite binary tree -inf_binary = astree(0, children=lambda n: (2*n + 1, 2*n + 2)) +On downtreelikes: +```python +nodes(tree) # Iterate through all nodes in tree (in no particular order). +descendants(node) # Children and grand-(grand-*)-children of node. +leaves(root) # Leaves reachable from root ``` -## Utility functions +If you want to iterate though the nodes in a specific order, use: +```python +preorder(node) # Nodes in preorder (root comes first). +postorder(node) # Nodes in postorder (root comes last). +levelorder(node) # Nodes near the root come before later nodes. +``` +These will return tuples with (node, item). The item-object contains information about the depth of the node. -Pretty printing +Additionally, on treelikes: ```python -tree = astree(seq, children=lambda x: [x[:-2], x[1:]] if x else []) -print_tree(tree) -print(to_string(tree)) - -# ['a', 'b', 'c', 'd'] -# ├─ ['a', 'b'] -# │ └─ ['b'] -# └─ ['b', 'c', 'd'] -# ├─ ['b'] -# └─ ['c', 'd'] -# └─ ['d'] +ancestors(node) # Ancestors of node. +path(node) # Path from root to this node including this node. +siblings(node) # Siblings of node ``` -Plotting with matplotlib +### Adapters ### + +To upgrade a `TreeLike` to a full `Tree` use `as_tree`. + ```python -import matplotlib.pyplot as plt +path_tree = as_tree(pathlib.Path("my_documents")) # Optionally pass `children`, `parent`, `label`. -expr = ast.parse("y = x*x + 1") -plot_tree(expr) -plt.show() +# Iterate over all its descendants +for node in path_tree.descendants: + path_obj = node.value # Get back a Path-object from TreeAdapter ``` -![images/tree_calc_plot.png](images/tree_calc_plot.png) + +There is also `TreeAdapter` to help with classes that are very different. + +### Exporting ### Export to various formats ```python +print_tree(tree) + +# If matplotlib is installed +plot_tree(tree) + +# These may require graphviz or Pillow to be installed. to_dot(tree) to_mermaid(tree) to_latex(tree) to_reportlab(tree) to_image(Path('.'), "filetree.png", how="dot") -to_image(AbstractTree, "class_hierarchy.svg", how="mermaid") +to_image(DownTree, "tree_hierarchy.svg", how="mermaid") to_pillow(tree).show() ``` - -## Find distance between nodes - -```python -import heapq - -from abstracttree import HeapTree, Route - -tree = HeapTree([5, 4, 3, 2, 1]) -heapq.heapify(tree.heap) -route = Route(tree.left_child, tree.right_child) - -print(f"{route.lca = }") # => HeapTree([1, 2, 3, 5, 4], 0) -print(f"{route.nodes.count() = }") # => 3 -print(f"{route.edges.count() = }") # => 2 -``` - -## A few concrete tree implementations - -- [anytree](https://github.com/c0fec0de/anytree) -- [treelib](https://github.com/caesar0301/treelib) -- [bigtree](https://github.com/kayjan/bigtree) -- [itertree](https://github.com/BR1py/itertree) -- [dendropy](https://github.com/jeetsukumaran/DendroPy) -- [ete](https://github.com/etetoolkit/ete) -- [littletree](https://github.com/lverweijen/littletree) - also by me - -## Tree visualisation - -- [PrettyPrintTree](https://github.com/AharonSambol/PrettyPrintTree) - colored terminal output diff --git a/docs/source/CHANGES.md b/docs/source/CHANGES.md index 50088a7..0343a64 100644 --- a/docs/source/CHANGES.md +++ b/docs/source/CHANGES.md @@ -1,5 +1,20 @@ # Changelog +## Version 0.2.0 + +* Add generics `TreeLike` and `DownTreeLike`. +* Add many new functions that accept above generics as parameter. +* Change many existing functions to accept above generics above as parameter. +* Rename `astree(x)` to `as_tree(x)`. Also `as_tree` will now always return a `TreeView`. +* Remove `convert_tree.register`, but add granular methods `children.register`, `parent.register`. +* Iteration methods on `tree.nodes` and `tree.descendants` now use `None` as index for root (instead of 0). +* Change how `Mapping` is converted to Tree. `children(mapping)` is mostly similar to `mapping.items()`. This works well on jsonlike-data. +* Replace `x.eqv(y)` method by `eqv(x, y)` function. +* `TreeAdapter` remains, but some alternative adapters have been removed. +* `TreeAdapter` is hashable if the underlying object is hashable. +* `TreeAdapter.node` has been renamed to `TreeAdapter.value`. +* `HeapTree` has become immutable and hashable (by id(heap) and index). The heap itself may remain mutable without a problem. + ## Version 0.1.1 * Make it possible to pass options like `file_format` to `to_image`. diff --git a/docs/source/api.rst b/docs/source/api.rst index af8fcc0..d98b941 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -1,22 +1,18 @@ API ================== -Abstract base classes +Mixins --------------------- -.. automodule:: abstracttree.tree +.. automodule:: abstracttree.mixins :members: :show-inheritance: -.. automodule:: abstracttree.binarytree - :members: - :show-inheritance: - -Adapters ------------------- +Generics +--------------------- -.. automodule:: abstracttree.adapters - :members: astree, convert_tree +.. automodule:: abstracttree.generics + :members: :show-inheritance: Export @@ -32,10 +28,11 @@ Predicates :members: :show-inheritance: -HeapTree +Adapters ------------------ -.. automodule:: abstracttree.heaptree - :members: HeapTree + +.. automodule:: abstracttree.adapters + :members: :show-inheritance: Route @@ -43,3 +40,9 @@ Route .. automodule:: abstracttree.route :members: Route :show-inheritance: + +Utils +------------------ +.. automodule:: abstracttree.utils + :members: eqv + :show-inheritance: diff --git a/docs/source/conf.py b/docs/source/conf.py index f02d496..822cd38 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -42,5 +42,4 @@ import sys from pathlib import Path src_folder = Path(__file__).parent.parent.parent / "src" -print(src_folder) sys.path.insert(0, str(src_folder.resolve(strict=True))) diff --git a/docs/source/index.rst b/docs/source/index.rst index 779f5d7..3278753 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -6,17 +6,15 @@ Welcome to AbstractTree's documentation! ======================================== -Trees are very common data structure that represents a hierarchy of common nodes. -This package defines abstract base classes for these data structure in order to make code reusable. -It also provides an ``astree`` adapter in case it's not possible to inherit from any of these classes. -Finally, it provides many exports that even work on objects that don't inherit from any of the abstract base classes. +Trees are very common data structure to represent a hierarchy of nodes. +This package provides a common interface to access and operate on these objects. .. toctree:: :maxdepth: 2 :caption: Contents: installation - usage_bundled + usage api CHANGES diff --git a/docs/source/installation.rst b/docs/source/installation.rst index 47c29a3..8609834 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -3,7 +3,7 @@ Installation Use `pip `_ to install littletree:: - $ pip install --upgrade littletree + $ pip install --upgrade abstracttree In addition you may want to install the following for some export functions:: diff --git a/docs/source/usage_bundled.rst b/docs/source/usage.rst similarity index 55% rename from docs/source/usage_bundled.rst rename to docs/source/usage.rst index f5dbc75..56cfbc2 100644 --- a/docs/source/usage_bundled.rst +++ b/docs/source/usage.rst @@ -25,23 +25,23 @@ Abstract tree classes Tree-->BinaryTree; -A Downtrees is an object that has links to its direct children. -A Tree is has links to both its children and its parent. -A binary tree has exactly two children (left and right). +A downtree is an object that has links to its direct children. +A tree is similar to a downtree, but also has a link to its parent. +A binary tree may have at most two children. A mutable tree can change its structure once created. +---------------------+-------------------------------+-------------------------------------+------------------------------------------------------------------------------------+ | ABC | Inherits from | Abstract Methods | Mixin Methods | +=====================+===============================+=====================================+====================================================================================+ -| ``AbstractTree`` | | | ``nid``, ``eqv()`` | +| ``AbstractTree`` | | | ``nid`` | +---------------------+-------------------------------+-------------------------------------+------------------------------------------------------------------------------------+ | ``DownTree`` | ``AbstractTree`` | ``children`` | ``nodes``, ``descendants``, ``leaves``, ``levels``, ``is_leaf``, ``transform()`` | +---------------------+-------------------------------+-------------------------------------+------------------------------------------------------------------------------------+ -| ``Tree`` | ``DownTree`` | | ``siblings`` | +| ``Tree`` | ``DownTree`` | ``parent`` | ``ancestors``, ``siblings``, ``path`` | +---------------------+-------------------------------+-------------------------------------+------------------------------------------------------------------------------------+ | ``MutableDownTree`` | ``DownTree`` | ``add_child()``, ``remove_child()`` | ``add_children()`` | +---------------------+-------------------------------+-------------------------------------+------------------------------------------------------------------------------------+ -| ``MutableTree`` | ``Tree``, ``MutableDownTree`` | | ``detach()`` | +| ``MutableTree`` | ``MutableDownTree``, ``Tree`` | | ``detach()`` | +---------------------+-------------------------------+-------------------------------------+------------------------------------------------------------------------------------+ | ``BinaryDownTree`` | ``DownTree`` | ``left_child``, ``right_child`` | ``children`` | +---------------------+-------------------------------+-------------------------------------+------------------------------------------------------------------------------------+ @@ -51,10 +51,9 @@ A mutable tree can change its structure once created. In your own code, you can inherit from these trees. For example, if your tree only has links to children:: - import abstracttree - from abstracttree import print_tree + from abstracttree import DownTree, print_tree - class MyTree(abstracttree.DownTree): + class MyTree(DownTree): def __init__(self, value, children=()): self.value = value self._children = children @@ -76,61 +75,78 @@ You can now use this class in the following way to generate output:: # ├─ MyTree 2 # └─ MyTree 3 -Adapter ------------------- +Generics +--------------------- -In practice, not all existing tree data structures implement one of the abstract classes specified in `Abstract classes `_. -As a bridge, you can use ``Tree.convert`` to convert these trees to a ``Tree`` instance. -However, whenever possible, it's recommended to inherit from ``Tree`` directly for minimal overhead. +Unfortunately, not all trees inherit from the mixins above. Yet, some objects still have treelike behaviour. +Therefore, AbstractTree provides support for a slightly weaker protocol. -``Tree.convert`` already does the right thing on many objects of the standard library:: +The following objects are ``TreeLike``: - # Inheritance hierarchy - Tree.convert(int) +- All objects that support ``obj.children`` and ``obj.parent``. +- Builtins classes ``pathlib.Path`` and ``zipfile.Path``. +- Third party tree classes from `anytree `_, `bigtree `_, `itertree `_ and `littletree `_. - # Abstract syntax tree - Tree.convert(ast.parse("1 + 1 == 2")) +The following objects are ``DownTreeLike``: - # Filesystem - Tree.convert(pathlib.Path("abstracttree")) +- All objects that support ``obj.children``. +- Recursive collections like lists, tuples, sets, dicts. This can be useful when dealing with json-data. - # Zipfile - Tree.convert(zipfile.ZipFile("eclipse.jar")) +This can be tested using `isinstance`:: - # Nested list - Tree.convert([[1, 2, 3], [4, 5, 6]]) + isinstance(Path(r"C:\\Windows\System"), TreeLike) # True + isinstance(range(100), DownTreeLike) # True + isinstance(range(100), TreeLike) # False + isinstance(5, DownTreeLike) # False + isinstance("some text", DownTreeLike) # False (even though it might be considered a collection by python). -It can also construct a tree by ducktyping on ``parent`` and ``children`` attributes:: +Basic functions +--------------- - # Works on objects by anytree, bigtree and littletree - Tree.convert(anytree.Node('node')) +On downtreelikes:: -Alternatively, you can use ``astree`` and explicitly specify how to find ``children`` and ``parent``:: + children(node) # Children of node + label(node) # String representation of node (similar to str, but output excludes parent and children) + nid(node) # Address of node (similar to id, but supports delegates). + eqv(node1, node2) # Check if 2 nodes have the same identity (similar to is, but supports delegates) - # Tree from json-data - data = {"name": "a", - "children": [ - {"name": "b", "children": []}, - {"name": "c", "children": []} - ]} - astree(data, children=operator.itemgetter["children"]) +Additionally, on treelikes:: - # pyqt.QtWidget - astree(widget, children=lambda w: w.children(), parent = lambda w: w.parent()) + parent(node) # Parent of node or None if node is root of its own tree. + root(node) # Find root of this tree. - # Tree from treelib - astree(tree.root, children=lambda nid: tree.children(nid), parent=lambda nid: tree.parent(nid)) +Examples:: - # itertree - astree(tree, children=iter, parent=lambda t: t.parent) + >>> from abstracttree import * + >>> children([1, 2, 3]) + [1, 2, 3] + >>> children({"name": "Philip", "children": ["Pete", "Mariam"]}) + [MappingItem(key="name", value="Philip"), MappingItem(key="children", value=["Pete", "Miriam"])] + >>> parent(Path(r"C:\\Windows\System")) + Path(r"C:\\Windows") + >>> label(Path(r"C:\\Windows\System")) + "System" + >>> eqv(Path(r"C:\\Windows\System"), Path(r"C:\\Windows\System")) + True + >>> eqv([1, 2, 3], [1, 2, 3]) + False - # Infinite binary tree - inf_binary = astree(0, children=lambda n: (2*n + 1, 2*n + 2)) +Iterators +--------- + +The following methods can iterate through nodes:: + + nodes(tree) # Iterate through all nodes in tree (in no particular order). + descendants(node) # Children and grand-(grand-*)-children of node. + leaves(root) # Leaves reachable from root + ancestors(node) # Ancestors of node. + path(node) # Path from root to this node including this node. + siblings(node) # Siblings of node Traversal ----------------------------------------- +~~~~~~~~~ -There are 3 common ways to traverse a tree: +The following methods also iterate, but in a very specific order. Pre-order The parent is iterated over before its children. @@ -143,15 +159,17 @@ Level-order All these are possible by writing one of:: - for node, item in tree.nodes.preorder(): + for node, item in preorder(tree): ... - for node, item in tree.nodes.postorder(): + for node, item in postorder(tree): ... - for node, item in tree.nodes.levelorder(): + for node, item in levelorder(tree): ... + # If Downtree is implemented, tree.nodes.preorder() also works. + These methods return an item in addition to a node. This item is a tuple of the following fields: @@ -165,19 +183,45 @@ index The first child of a parent gets index 0, the second gets index 1. The root of the (sub)tree always gets an index of ``0`` even if it has prior siblings. -To iterate over the descendants (the nodes without the root), similar methods are defined:: +To iterate over the descendants without the root, use the following:: - for descendant, item in tree.descendants.preorder(): + for descendant, item in preorder(tree, include_root=False): ... + # If Downtree is implemented, tree.descendants.preorder() also works. If the order of iteration doesn't matter an alternative way to iterate is as follows:: - for node in tree.nodes: + for node in nodes(tree): ... - for descendant in tree.descendants: + for descendant in descendants(tree): ... +Adapters +------------------ + +If you want a ``Tree``-object, you can use ``as_tree`` to convert these treelikes to a full ``Tree``. +Alternatively, you can explicitly specify how to find ``children`` and ``parent``:: + + # Tree from json-data + data = {"name": "a", + "children": [ + {"name": "b", "children": []}, + {"name": "c", "children": []} + ]} + as_tree(data, children=operator.itemgetter["children"]) + + # pyqt.QtWidget + as_tree(widget, children=lambda w: w.children(), parent = lambda w: w.parent()) + + # Tree from treelib + as_tree(tree.root, children=lambda nid: tree.children(nid), parent=lambda nid: tree.parent(nid)) + + # itertree + as_tree(tree, children=iter, parent=lambda t: t.parent) + + # Infinite binary tree + inf_binary = as_tree(0, children=lambda n: (2*n + 1, 2*n + 2)) Export ---------------------------------------- @@ -185,27 +229,22 @@ Export Pretty printing:: print_tree(Path()) - # . - # ├─ abstracttree - # │ ├─ abstracttree\conversions.py - # │ ├─ abstracttree\export.py - # │ ├─ abstracttree\predicates.py - # │ ├─ abstracttree\treeclasses.py - # │ └─ abstracttree\__init__.py - # ├─ LICENSE - # ├─ Makefile - # ├─ manual.txt - # ├─ pyproject.toml - # ├─ README.md - # └─ tests - # ├─ tests\test_downtree.py - # ├─ tests\test_export.py - # ├─ tests\test_mutabletree.py - # ├─ tests\test_tree.py - # ├─ tests\test_uptree.py - # └─ tests\tree_instances.py - + # ├─ adapters + # │ ├─ adapters.py + # │ ├─ heaptree.py + # │ └─ __init__.py + # ├─ export.py + # ├─ generics.py + # ├─ iterators.py + # ├─ mixins + # │ ├─ trees.py + # │ ├─ views.py + # │ └─ __init__.py + # ├─ predicates.py + # ├─ route.py + # ├─ utils.py + # └─ __init__.py Plotting with matplotlib:: @@ -218,7 +257,7 @@ Plotting with matplotlib:: Export to graphviz:: - tree = astree(seq, children=lambda x: [x[:-2], x[1:]] if x else []) + tree = as_tree(seq, children=lambda x: [x[:-2], x[1:]] if x else []) to_graphviz(tree) diff --git a/examples/anytree_example.py b/examples/anytree_example.py index accf6a3..806b27e 100644 --- a/examples/anytree_example.py +++ b/examples/anytree_example.py @@ -2,7 +2,9 @@ import anytree -# This works for combining both +# If you want to use abstracttree as a mixin, it can be done like this. +# Usually the mixin would come second, but in this case anytree.Node has many similarly named methods and properties, +# which are already provided by abstracttree in a more generic way. class MyTree(abstracttree.Tree, anytree.Node): children = anytree.NodeMixin.children parent = anytree.NodeMixin.parent diff --git a/examples/itertree_example.py b/examples/itertree_example.py index 62cd621..18530e4 100644 --- a/examples/itertree_example.py +++ b/examples/itertree_example.py @@ -21,14 +21,9 @@ class MyTree(itertree.iTree, abstracttree.Tree): def children(self): return self - def __str__(self): - # Don't print subtrees - only_self = lambda n: n is self - return self.renders(filter_method=only_self) - - tree = MyTree(1) +# Add some descendants node = tree for i in range(7): node.append(child := MyTree(i)) @@ -36,7 +31,7 @@ def __str__(self): # AbstractTree print("Height of tree:", tree.levels.count() - 1) -print_tree(tree) +print_tree(tree, "MyTree({.tag})") # iTree print("Height of tree:", tree.max_depth) diff --git a/examples/networkx_example.py b/examples/networkx_example.py index 056b3b1..1d40e7b 100644 --- a/examples/networkx_example.py +++ b/examples/networkx_example.py @@ -13,13 +13,19 @@ class NetworkXTree(TreeAdapter): @property - def nid(self): - return id(self.identifier) + def graph(self): + (graph, _) = self.value + return graph - def eqv(self, other): - if not isinstance(other, type(self)): - return False - return self.graph is other.graph and self.identifier == other.identifier + @property + def identifier(self): + (_, identifier) = self.value + return identifier + + @property + def data(self): + (graph, identifier) = self.value + return graph.nodes[identifier] def __str__(self): return f"{self.identifier}" @@ -42,19 +48,8 @@ def parent_func(node): return graph, parents[0] @property - def graph(self): - (graph, _) = self.node - return graph - - @property - def identifier(self): - (_, identifier) = self.node - return identifier - - @property - def data(self): - (graph, identifier) = self.node - return graph.nodes[identifier] + def nid(self): + return id(self.graph) << 32 | id(self.identifier) def graph_to_tree(graph: nx.Graph, node=None, check=True): diff --git a/examples/sklearn_example.py b/examples/sklearn_example.py index e2b12b9..b4e9089 100644 --- a/examples/sklearn_example.py +++ b/examples/sklearn_example.py @@ -1,11 +1,28 @@ -# Attempt to contribute this directly to sklearn: -# https://github.com/scikit-learn/scikit-learn/pull/28364 - from abstracttree import print_tree, to_latex, to_image -from abstracttree.adapters import StoredParent +from abstracttree.adapters import TreeAdapter, as_tree + +def main(): + # Train a model + from sklearn.datasets import load_iris + from sklearn.model_selection import train_test_split + from sklearn.tree import DecisionTreeClassifier + iris = load_iris() + X = iris.data + y = iris.target + X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0) + + clf = DecisionTreeClassifier(max_leaf_nodes=3, random_state=0) + model = clf.fit(X_train, y_train) + + # Now treat the model as a tree + tree = convert_decisiontree(model) + print_tree(tree) + print(to_latex(tree, node_shape="rectangle", leaf_distance='8em', level_distance="10em")) + to_image(tree, 'sklearn_example.png') -class DecisionTreeAdapter(StoredParent): + +class DecisionTreeAdapter(TreeAdapter): @staticmethod def child_func(value): tree, idx = value @@ -17,32 +34,32 @@ def child_func(value): @property def feature(self): - tree, idx = self.node + tree, idx = self.value return tree.feature[idx] @property def threshold(self): - tree, idx = self.node + tree, idx = self.value return tree.threshold[idx] @property - def value(self): - tree, idx = self.node + def result(self): + tree, idx = self.value return tree.value[idx] @property def n_samples(self): - tree, idx = self.node + tree, idx = self.value return tree.n_node_samples[idx] @property def impurity(self): - tree, idx = self.node + tree, idx = self.value return tree.impurity[idx] def __str__(self): if self.children: - text = f"if X[:, {self.feature}] ≤ {self.threshold:.4g}: # {self.value}" + text = f"if X[:, {self.feature}] ≤ {self.threshold:.4g}: # {self.result}" else: text = f"return {self.value}" comment = f"# n_samples = {self.n_samples}, impurity = {self.impurity:.4g}" @@ -55,20 +72,4 @@ def convert_decisiontree(model): if __name__ == '__main__': - # Train a model - from sklearn.datasets import load_iris - from sklearn.model_selection import train_test_split - from sklearn.tree import DecisionTreeClassifier - iris = load_iris() - X = iris.data - y = iris.target - X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0) - - clf = DecisionTreeClassifier(max_leaf_nodes=3, random_state=0) - model = clf.fit(X_train, y_train) - - # Now treat the model as a tree - tree = convert_decisiontree(model) - print_tree(tree) - print(to_latex(tree, node_shape="rectangle", leaf_distance='8em', level_distance="10em")) - to_image(tree, 'sklearn_example.png') + main() diff --git a/pyproject.toml b/pyproject.toml index 61313d6..c922910 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ authors = [ ] description = "Abstract base classes for tree data structures" readme = "README.md" -requires-python = ">= 3.9" +requires-python = ">= 3.10" keywords = [ "tree", "datastructure", @@ -22,7 +22,6 @@ classifiers = [ "Operating System :: OS Independent", "License :: OSI Approved :: Apache Software License", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", diff --git a/src/abstracttree/__init__.py b/src/abstracttree/__init__.py index a398c9a..8783394 100644 --- a/src/abstracttree/__init__.py +++ b/src/abstracttree/__init__.py @@ -1,29 +1,4 @@ -__all__ = [ - "Tree", - "DownTree", - "MutableTree", - "MutableDownTree", - "BinaryTree", - "BinaryDownTree", - "astree", - "print_tree", - "plot_tree", - "to_string", - "to_image", - "to_dot", - "to_mermaid", - "to_pillow", - "to_reportlab", - "to_latex", - "RemoveDuplicates", - "PreventCycles", - "MaxDepth", - "HeapTree", - "Route", -] - -from .adapters import astree -from .binarytree import BinaryTree, BinaryDownTree +from .adapters import HeapTree, as_tree, convert_tree, TreeAdapter from .export import ( print_tree, plot_tree, @@ -35,7 +10,30 @@ to_latex, to_reportlab, ) -from .heaptree import HeapTree +from .generics import ( + TreeLike, + DownTreeLike, + children, + parent, + root, + nid, + label, + parents, +) +from .iterators import ( + nodes, + descendants, + preorder, + postorder, + levels, + levelorder, + levels_zigzag, + leaves, + siblings, + ancestors, + path, +) +from .mixins import Tree, DownTree, MutableDownTree, MutableTree, BinaryDownTree, BinaryTree from .predicates import RemoveDuplicates, PreventCycles, MaxDepth from .route import Route -from .tree import Tree, DownTree, MutableDownTree, MutableTree +from .utils import eqv diff --git a/src/abstracttree/adapters.py b/src/abstracttree/adapters.py deleted file mode 100644 index 3547c12..0000000 --- a/src/abstracttree/adapters.py +++ /dev/null @@ -1,404 +0,0 @@ -import ast -import functools -import operator -import pathlib -import xml.etree.ElementTree as ET -import zipfile -from collections.abc import Sequence, Mapping -from typing import TypeVar, Callable, Iterable, overload, Collection, Union - -from .tree import Tree, DownTree - -TWrap = TypeVar("TWrap") -BaseString = Union[str, bytes, bytearray] # Why did they ever remove this type? - - -@overload -def astree(obj) -> Tree: ... - - -@overload -def astree(obj: TWrap, children: Callable[[TWrap], Collection[TWrap]]) -> Tree: ... - - -@overload -def astree( - obj: TWrap, - children: Callable[[TWrap], Collection[TWrap]], - parent: Callable[[TWrap], TWrap], -) -> Tree: ... - - -def astree( - obj: TWrap, - children: Callable[[TWrap], Iterable[TWrap]] = None, - parent: Callable[[TWrap], TWrap] = None, -) -> Tree: - """Convert an arbitrary object into a tree. - - If no children or parents are given it will try a standard conversion. - If children and parents are given, they will be used as a function. - If children, but no parent is given, it will create a tree that has obj as its root. - """ - if not children: - return convert_tree(obj) - else: - if parent: - - class CustomTree(TreeAdapter): - child_func = staticmethod(children) - parent_func = staticmethod(parent) - - else: - - class CustomTree(StoredParent): - child_func = staticmethod(children) - - return CustomTree(obj) - - -@functools.singledispatch -def convert_tree(tree): - """Low level conversion of tree object to an abstracttree.Tree. - - The default implementation ducktypes on `tree.parent` and `tree.children`. - Classes can also define a method _abstracttree_ to override their conversion. - """ - explicit_conversion = getattr(tree, "_abstracttree_", None) - if explicit_conversion: - return explicit_conversion() - if hasattr(tree, "children"): - if hasattr(tree, "parent"): - return TreeAdapter(tree) - else: - return StoredParent(tree) - else: - raise NotImplementedError - - -@convert_tree.register -def _(tree: Tree): - return tree - - -@convert_tree.register -def _(tree: DownTree): - return UpgradedTree(tree) - - -@convert_tree.register -def _(tree: pathlib.PurePath): - return PathTree(tree) - - -@convert_tree.register -def _(zf: zipfile.ZipFile): - return PathTree(zipfile.Path(zf)) - - -@convert_tree.register -def _(path: zipfile.Path): - return PathTree(path) - - -@convert_tree.register -def _(tree: Sequence): - return SequenceTree(tree) - - -@convert_tree.register(str) -@convert_tree.register(bytes) -@convert_tree.register(bytearray) -def _(_: BaseString): - raise NotImplementedError( - "astree(x: str | bytes | bytearray) is unsafe, " - "because x is infinitely recursively iterable." - ) - - -@convert_tree.register -def _(tree: Mapping): - return MappingTree((None, tree)) - - -@convert_tree.register -def _(cls: type, invert=False): - if invert: - return InvertedTypeTree(cls) - return TypeTree(cls) - - -@convert_tree.register -def _(node: ast.AST): - return AstTree(node) - - -@convert_tree.register -def _(element: ET.Element): - return XmlTree(element) - - -@convert_tree.register -def _(tree: ET.ElementTree): - return XmlTree(tree.getroot()) - - -class TreeAdapter(Tree): - __slots__ = "node" - child_func: Callable[[TWrap], Iterable[TWrap]] = operator.attrgetter("children") - parent_func: Callable[[TWrap], TWrap] = operator.attrgetter("parent") - - def __init__(self, node): - self.node = node - - def __repr__(self): - return f"{type(self).__name__}({self.node})" - - def __str__(self): - return str(self.node) - - def __eq__(self, other): - return isinstance(other, type(self)) and self.node == other.node - - @property - def nid(self): - return id(self.node) - - def eqv(self, other): - return isinstance(other, type(self)) and self.node is other.node - - @property - def children(self): - cls = self.__class__ - child_func = self.child_func - return [cls(c) for c in child_func(self.node)] - - @property - def parent(self): - parent = self.parent_func(self.node) - if parent is not None: - return self.__class__(parent) - else: - return None - - -class PathTree(TreeAdapter): - __slots__ = () - _custom_nids = {} - - @property - def nid(self): - try: # Doesn't work on zipfile - st = self.node.lstat() - except (FileNotFoundError, AttributeError): - return self._custom_nids.setdefault(str(self.node), len(self._custom_nids)) - else: - return -st.st_ino - - def eqv(self, other): - return self.node == other.node - - @staticmethod - def parent_func(path): - # pathlib makes parent an infinite, but we want None - parent = path.parent - if path != parent: - return parent - else: - return None - - @property - def children(self): - try: - return list(map(type(self), self.child_func(self.node))) - except PermissionError: - return [] - - @staticmethod - def child_func(p): - return p.iterdir() if p.is_dir() else () - - @property - def root(self): - return type(self)(type(self.node)(self.node.anchor)) - - -class StoredParent(Tree): - __slots__ = "node", "_parent" - child_func: Callable[[TWrap], Iterable[TWrap]] = operator.attrgetter("children") - - def __init__(self, node, parent=None): - self.node = node - self._parent = parent - - def __repr__(self): - return f"{type(self).__name__}({self.node})" - - def __str__(self): - return str(self.node) - - def __eq__(self, other): - return isinstance(other, type(self)) and self.node == other.node - - @property - def nid(self): - return id(self.node) - - def eqv(self, other): - return isinstance(other, type(self)) and self.node is other.node - - @property - def children(self): - cls = type(self) - child_func = self.child_func - return [cls(c, self) for c in child_func(self.node)] - - @property - def parent(self): - return self._parent - - -class UpgradedTree(StoredParent): - __slots__ = () - - @property - def nid(self): - return self.node.nid - - def eqv(self, other): - return isinstance(other, type(self)) and self.node.eqv(other.node) - - -class SequenceTree(StoredParent): - __slots__ = () - - @staticmethod - def child_func(seq): - if isinstance(seq, Sequence) and not isinstance(seq, BaseString): - return seq - else: - return () - - def __str__(self): - if self.is_leaf: - return str(self.node) - else: - cls_name = type(self.node).__name__ - return f"{cls_name}[{len(self.node)}]" - - -class MappingTree(StoredParent): - __slots__ = () - - @staticmethod - def child_func(item): - _, value = item - if isinstance(value, Mapping): - return value.items() - elif value is not None: - return [(value, None)] - else: - return [] - - @property - def key(self): - key, _ = self.node - return key - - @property - def mapping(self): - _, mapping = self.node - return mapping - - def __str__(self): - key, mapping = self.node - if self.is_root: - cls_name = type(mapping).__name__ - return f"{cls_name}[{len(mapping)}]" - else: - return str(key) - - -class TypeTree(StoredParent): - __slots__ = () - - @property - def nid(self): - return id(self.node) - - @staticmethod - def child_func(cls): - return cls.__subclasses__() - - @property - def parents(self): - return list(map(type(self), self.node.__bases__)) - - def __str__(self): - return self.node.__qualname__ - - -class InvertedTypeTree(TypeTree): - __slots__ = () - - @staticmethod - def child_func(cls): - return cls.__bases__ - - -class AstTree(StoredParent): - __slots__ = () - CONT = "↓" - SINGLETON = Union[ast.expr_context, ast.boolop, ast.operator, ast.unaryop, ast.cmpop] - - @classmethod - def child_func(cls, node): - return tuple(node for node in ast.iter_child_nodes(node) if cls.is_child(node)) - - @classmethod - def is_child(cls, node): - return isinstance(node, ast.AST) and not isinstance(node, cls.SINGLETON) - - def __str__(self): - if self.is_leaf: - return ast.dump(self.node) - else: - format_value = self.format_value - args = [f"{name}={format_value(field)}" for name, field in ast.iter_fields(self.node)] - joined_args = ", ".join(args) - return f"{type(self.node).__name__}({joined_args})" - - @classmethod - def format_value(cls, field): - if cls.is_child(field): - field_str = cls.CONT - elif isinstance(field, ast.AST): - field_str = ast.dump(field) - elif isinstance(field, Sequence) and not isinstance(field, BaseString): - field = [cls.format_value(f) for f in field] - field_str = "[" + ", ".join(field) + "]" - else: - field_str = repr(field) - return field_str - - -class XmlTree(StoredParent): - @staticmethod - def child_func(element): - return element - - def __str__(self): - element = self.node - output = [f"<{element.tag}"] - - for k, v in element.items(): - output.append(f" {k}={v!r}") - output.append(">") - - if text := element.text and str(element.text).strip(): - output.append(text) - output.append(f"") - if tail := element.tail and str(element.tail).strip(): - output.append(tail) - - return "".join(output) diff --git a/src/abstracttree/adapters/__init__.py b/src/abstracttree/adapters/__init__.py new file mode 100644 index 0000000..e1990d4 --- /dev/null +++ b/src/abstracttree/adapters/__init__.py @@ -0,0 +1,9 @@ +__all__ = [ + "as_tree", + "convert_tree", + "HeapTree", + "TreeAdapter", +] + +from .heaptree import HeapTree +from .adapters import as_tree, convert_tree, TreeAdapter diff --git a/src/abstracttree/adapters/adapters.py b/src/abstracttree/adapters/adapters.py new file mode 100644 index 0000000..06a7e3f --- /dev/null +++ b/src/abstracttree/adapters/adapters.py @@ -0,0 +1,117 @@ +from collections.abc import Sequence, Callable, Iterable +from functools import lru_cache +from typing import Optional, TypeVar, Type + +import abstracttree.generics as generics +from abstracttree.generics import TreeLike, DownTreeLike +from abstracttree.mixins import Tree +from abstracttree.utils import eqv + +T = TypeVar("T") + + +def convert_tree(tree: DownTreeLike, required_type=Type[T]) -> T: + """Convert a TreeLike to a powerful Tree. + + If needed, it uses a TreeAdapter. + """ + if isinstance(tree, required_type): + return tree + elif hasattr(tree, '_abstracttree_'): + tree = tree._abstracttree_() + else: + tree = as_tree(tree) + + if isinstance(tree, required_type): + return tree + else: + raise TypeError(f"Unable to convert {type(tree)} to {required_type}") + + +def as_tree( + obj: T, + children: Callable[[T], Iterable[T]] = None, + parent: Callable[[T], Optional[T]] = None, + label: Callable[[T], str] = None, +) -> "TreeAdapter": + """Convert any object to a tree. + + Functions can be passed to control how the conversion should be done. + The original object can be accessed by using the value attribute. + """ + cls = type(obj) + adapter = compile_adapter(cls, children, parent, label) + tree = adapter(obj) + return tree + + +@lru_cache(maxsize=None) +def compile_adapter( + cls, + children: Callable[[T], Iterable[T]] = None, + parent: Callable[[T], Optional[T]] = None, + label: Callable[[T], str] = None, +): + if not parent and issubclass(cls, TreeLike): + parent = generics.parent.dispatch(cls) + + class CustomTreeAdapter(TreeAdapter): + child_func = staticmethod(children or generics.children.dispatch(cls)) + parent_func = staticmethod(parent) + label_func = staticmethod(label or generics.label.dispatch(cls)) + + return CustomTreeAdapter + +# Alias for backwards compatibility +astree = as_tree + + +class TreeAdapter(Tree): + child_func = staticmethod(generics.children) + parent_func = staticmethod(generics.parent) + label_func = staticmethod(generics.label) + + def __init__(self, value: TreeLike, _parent=None): + self._value = value + self._parent = _parent + + def __repr__(self) -> str: + return f"{type(self).__qualname__}({self.value!r})" + + def __str__(self) -> str: + return self.label_func(self._value) + + @property + def nid(self) -> int: + return generics.nid(self._value) + + def __eq__(self, other) -> bool: + """Check if the same node is wrapped. Similar to eqv(self.value, other.value).""" + return eqv(self, other) + + def __hash__(self) -> int: + """An adapter is hashable iff the underlying object is hashable.""" + return hash(self._value) + + @property + def value(self): + return self._value + + @property + def parent(self: T) -> Optional[T]: + if self._parent is not None: + return self._parent + + cls = type(self) + if pf := cls.parent_func: + par = pf(self._value) + if par is not None: + return cls(par) + return None + + @property + def children(self: T) -> Sequence[T]: + cls = type(self) + _child_func = cls.child_func + child_nodes = _child_func(self._value) + return [cls(c, self) for c in child_nodes] diff --git a/src/abstracttree/adapters/heaptree.py b/src/abstracttree/adapters/heaptree.py new file mode 100644 index 0000000..b6286a9 --- /dev/null +++ b/src/abstracttree/adapters/heaptree.py @@ -0,0 +1,95 @@ +from collections.abc import MutableSequence +from typing import Collection, Optional, TypeVar + +from ..mixins import BinaryTree + +D = TypeVar("D") + + +class HeapTree(BinaryTree): + """Provides a tree interface to a heap. + + Mainly useful for visualisation purposes. + >>> from abstracttree import print_tree + >>> import heapq + >>> tree = HeapTree() + >>> for n in range(5, 0, -1): + >>> heapq.heappush(tree.heap, n) + >>> print_tree(tree) + 0 → 1 + ├─ 1 → 2 + │ ├─ 3 → 5 + │ └─ 4 → 4 + └─ 2 → 3 + """ + + __slots__ = "_heap", "_index" + + def __init__(self, heap: MutableSequence[D] = None, index: int = 0): + if heap is None: + heap = [] + self._heap = heap + self._index = index + + def __repr__(self): + return f"{type(self).__qualname__}{self._heap, self._index}" + + def __str__(self): + try: + return f"{self._index} → {self.value}" + except IndexError: + return repr(self) + + @property + def nid(self): + return (id(self._heap) << 32) | self._index + + @property + def heap(self) -> MutableSequence[D]: + return self._heap + + @property + def index(self): + return self._index + + @property + def value(self) -> D: + return self._heap[self._index] + + def __eq__(self, other): + """Nodes should refer to the same _heap (identity) with the same _index.""" + return (isinstance(other, HeapTree) + and self._heap is other._heap + and self._index == other._index) + + def __hash__(self): + """Hashes by id(list) and _index.""" + return hash((id(self._heap), self._index)) + + @property + def children(self) -> Collection["HeapTree"]: + return [ + HeapTree(self._heap, i) + for i in range(2 * self._index + 1, 2 * self._index + 3) + if i < len(self._heap) + ] + + @property + def left_child(self) -> Optional["HeapTree"]: + i = 2 * self._index + 1 + if i < len(self._heap): + return HeapTree(self._heap, i) + + @property + def right_child(self) -> Optional["HeapTree"]: + i = 2 * self._index + 2 + if i < len(self._heap): + return HeapTree(self._heap, i) + + @property + def parent(self) -> Optional["HeapTree"]: + n = self._index + if n != 0: + return HeapTree(self._heap, (n - 1) // 2) + else: + return None diff --git a/src/abstracttree/binarytree.py b/src/abstracttree/binarytree.py deleted file mode 100644 index 052ca78..0000000 --- a/src/abstracttree/binarytree.py +++ /dev/null @@ -1,92 +0,0 @@ -from abc import ABCMeta, abstractmethod -from collections import deque -from typing import Optional, Sequence - -from abstracttree import tree - -from .tree import DownTree, Tree, TNode, NodeItem - - -class BinaryDownTree(DownTree, metaclass=ABCMeta): - """Binary-tree with links to children.""" - - __slots__ = () - - @property - @abstractmethod - def left_child(self) -> Optional[TNode]: - raise None - - @property - @abstractmethod - def right_child(self) -> Optional[TNode]: - raise None - - @property - def children(self) -> Sequence[TNode]: - children = list() - if self.left_child is not None: - children.append(self.left_child) - if self.right_child is not None: - children.append(self.right_child) - return children - - @property - def nodes(self): - return NodesView([self], 0) - - @property - def descendants(self): - return NodesView(self.children, 1) - - -class BinaryTree(BinaryDownTree, Tree, metaclass=ABCMeta): - """Binary-tree with links to children and to parent.""" - - __slots__ = () - - -class NodesView(tree.NodesView): - """Extend NodesView to make it do inorder.""" - - __slots__ = () - - def inorder(self, keep=None): - """ - Iterate through nodes in inorder (traverse left, yield root, traverse right). - - Note: - - `item.index` will be 0 for every left child - - `item.index` will be 1 for every right child (even if node.left_child is None) - This is a bit different from how `preorder()`, `postorder()` and `levelorder()` work, - because those functions always give index 0 to the first child, - regardless of whether it's a left or right child. - Like the other iterators, the root of a subtree always gets item.index equal to 0, - even if it is actually a right child in a bigger tree. - """ - stack = deque() - - for index, node in enumerate(self.nodes): - depth = self.level - item = NodeItem(index, depth) - - while node is not None or stack: - # Traverse down/left - left_child, left_item = node.left_child, NodeItem(0, depth + 1) - while left_child is not None and (not keep or keep(left_child, left_item)): - stack.append((node, item)) - node, item, depth = left_child, left_item, depth + 1 - left_child, left_item = node.left_child, NodeItem(0, depth + 1) - - yield node, item - - # Traverse right/up - right_child, right_item = node.right_child, NodeItem(1, depth + 1) - while stack and ( - right_child is None or (keep and not keep(right_child, right_item)) - ): - node, item = stack.pop() - yield node, item - depth -= 1 - right_child, right_item = node.right_child, NodeItem(1, depth + 1) - node, item, depth = right_child, right_item, depth + 1 diff --git a/src/abstracttree/export.py b/src/abstracttree/export.py index 09257a2..d072a02 100644 --- a/src/abstracttree/export.py +++ b/src/abstracttree/export.py @@ -4,12 +4,15 @@ import operator import subprocess import sys -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Callable from pathlib import Path -from typing import Union, Callable, TypedDict, Tuple, Any, TypeVar, Optional +from typing import Union, TypedDict, Tuple, Any, TypeVar, Optional +from . import generics +from .iterators import preorder, levels, levelorder +from .adapters import convert_tree +from .generics import TreeLike, DownTreeLike, label, nid from .predicates import PreventCycles, MaxDepth -from .tree import DownTree, Tree __all__ = [ "print_tree", @@ -20,6 +23,7 @@ "to_image", "to_pillow", "to_latex", + "to_reportlab", "LiteralText", ] @@ -69,7 +73,7 @@ def new_f(tree, file=None, *args, **kwargs): return new_f -def print_tree(tree, formatter=str, style=None, keep=None): +def print_tree(tree, formatter=generics.label, style=None, keep=None): """Print this tree. Shortcut for print(to_string(tree)).""" if sys.stdout: if not style: @@ -80,17 +84,21 @@ def print_tree(tree, formatter=str, style=None, keep=None): @_wrap_file def to_string( - tree: DownTree, - formatter=str, + tree: DownTreeLike, + formatter: Callable[[object], str] | str | None = generics.label, *, file=None, style: Union[str, Style] = "square", keep=None, ): """Converts tree to a string in a pretty format.""" - tree = Tree.convert(tree) + tree = convert_tree(tree, TreeLike) if isinstance(style, str): style = DEFAULT_STYLES[style] + if formatter is None: + formatter = generics.label + elif isinstance(formatter, str): + formatter = formatter.format empty_style = len(style["last"]) * " " lookup1 = [empty_style, style["vertical"]] lookup2 = [style["last"], style["branch"]] @@ -110,11 +118,14 @@ def to_string( def _iterate_patterns(tree, keep): # Yield for each node a list of continuation indicators. # The continuation indicator tells us whether the branch at a certain level is continued. + _children = generics.children.dispatch(type(tree)) + _parent = generics.parent.dispatch(type(tree)) + pattern = [] yield pattern, tree - for node, item in tree.descendants.preorder(keep=keep): + for node, item in preorder(tree, keep=keep, include_root=False): del pattern[item.depth - 1 :] - is_continued = item.index < len(node.parent.children) - 1 + is_continued = item.index < len(_children(_parent(node))) - 1 pattern.append(is_continued) yield pattern, node @@ -128,12 +139,22 @@ def _write_indent(file, pattern, lookup1, lookup2): file.write(lookup2[pattern[-1]]) -def plot_tree(tree: Tree, ax=None, formatter=str, keep=DEFAULT_PREDICATE, annotate_args=None): +def plot_tree( + tree: DownTreeLike, + ax=None, + formatter: Callable[[DownTreeLike], str] | str = label, + keep=DEFAULT_PREDICATE, + annotate_args=None +): """Plot the tree using matplotlib (if installed).""" # Roughly based on sklearn.tree.plot_tree() import matplotlib.pyplot as plt - tree = Tree.convert(tree) + tree = convert_tree(tree, TreeLike) + if isinstance(formatter, str): + formatter = formatter.format + + _parent = generics.parent.dispatch(type(tree)) if ax is None: ax = plt.gca() @@ -153,8 +174,8 @@ def plot_tree(tree: Tree, ax=None, formatter=str, keep=DEFAULT_PREDICATE, annota nodes_xy = {} - tree_height = max(it.depth for _, it in tree.descendants.preorder(keep=keep)) - for depth, level in zip(range(tree_height + 1), tree.levels): + tree_height = max(it.depth for _, it in preorder(tree, keep=keep, include_root=False)) + for depth, level in zip(range(tree_height + 1), levels(tree)): level = list(level) for i, node in enumerate(level): x = (i + 1) / (len(level) + 1) @@ -163,13 +184,13 @@ def plot_tree(tree: Tree, ax=None, formatter=str, keep=DEFAULT_PREDICATE, annota if node is tree: ax.annotate(formatter(node), (x, y), **kwargs) else: - parent = nodes_xy[node.parent.nid] - ax.annotate(formatter(node), parent, (x, y), **kwargs) - nodes_xy[node.nid] = x, y + p_node = nodes_xy[nid(_parent(node))] + ax.annotate(formatter(node), p_node, (x, y), **kwargs) + nodes_xy[nid(node)] = x, y return ax -TNode = TypeVar("TNode", bound=DownTree) +TNode = TypeVar("TNode", bound=DownTreeLike) TShape = Union[Tuple[str, str], str, Callable[[TNode], Union[Tuple[str, str], str]]] NodeAttributes = Union[Mapping[str, Any], Callable[[TNode], Mapping[str, Any]]] EdgeAttributes = Union[Mapping[str, Any], Callable[[TNode, TNode], Mapping[str, Any]]] @@ -177,7 +198,7 @@ def plot_tree(tree: Tree, ax=None, formatter=str, keep=DEFAULT_PREDICATE, annota def to_image( - tree: Tree, + tree: DownTreeLike, file=None, how="dot", *args, @@ -205,7 +226,7 @@ def to_image( _image_mermaid(tree, Path(file), *args, **kwargs) -def to_pillow(tree: Tree, **kwargs): +def to_pillow(tree: DownTreeLike, **kwargs): """Convert tree to pillow-format (uses graphviz on the background).""" from PIL import Image @@ -215,7 +236,7 @@ def to_pillow(tree: Tree, **kwargs): return Image.open(io.BytesIO(to_image(tree, file=None, how="dot", **kwargs))) -def to_reportlab(tree: Tree, **kwargs): +def to_reportlab(tree: DownTreeLike, **kwargs): """Convert tree to drawing for use with reportlab package.""" from svglib.svglib import svg2rlg @@ -228,7 +249,7 @@ def to_reportlab(tree: Tree, **kwargs): def _image_dot( - tree: Tree, + tree: DownTreeLike, file=None, file_format="png", program_path="dot", @@ -245,7 +266,7 @@ def _image_dot( def _image_mermaid( - tree: Tree, + tree: DownTreeLike, filename, program_path="mmdc", **kwargs, @@ -262,21 +283,23 @@ def _image_mermaid( @_wrap_file def to_dot( - tree: Tree, + tree: DownTreeLike, file=None, keep=DEFAULT_PREDICATE, node_name: Union[str, Callable[[TNode], str], None] = None, - node_label: Union[str, Callable[[TNode], str], None] = str, + node_label: Union[str, Callable[[TNode], str], None] = generics.label, node_shape: TShape = None, node_attributes: NodeAttributes = None, edge_attributes: EdgeAttributes = None, graph_attributes: GraphAttributes = None, ): """Export to `graphviz `_.""" - tree = Tree.convert(tree) + tree = convert_tree(tree, TreeLike) if node_name is None: node_name = _node_name_default + _parent = generics.parent.dispatch(type(tree)) + if node_attributes is None: node_attributes = dict() else: @@ -307,7 +330,7 @@ def to_dot( file.write(f"edge{attrs};\n") nodes = [] - for node, _ in tree.nodes.levelorder(keep=PreventCycles() & keep): + for node, _ in levelorder(tree, keep=PreventCycles() & keep): nodes.append(node) name = _escape_string(node_name(node), "dot") attrs = _handle_attributes(node_dynamic, node) @@ -315,9 +338,9 @@ def to_dot( nodes = iter(nodes) next(nodes) for node in nodes: - parent_name = _escape_string(node_name(node.parent), "dot") + parent_name = _escape_string(node_name(_parent(node)), "dot") child_name = _escape_string(node_name(node), "dot") - attrs = _handle_attributes(edge_dynamic, node.parent, node) + attrs = _handle_attributes(edge_dynamic, _parent(node), node) file.write(f"{parent_name}{arrow}{child_name}{attrs};\n") file.write("}\n") @@ -362,20 +385,22 @@ def _handle_attributes(attributes, *args): @_wrap_file def to_mermaid( - tree: Tree, + tree: DownTreeLike, file=None, keep=DEFAULT_PREDICATE, node_name: Union[str, Callable[[TNode], str], None] = None, - node_label: Union[str, Callable[[TNode], str], None] = str, + node_label: Union[str, Callable[[TNode], str], None] = generics.label, node_shape: TShape = "box", edge_arrow: Union[str, Callable[[TNode, TNode], str]] = "-->", graph_direction: str = "TD", ): """Export to `mermaid `_.""" - tree = Tree.convert(tree) + tree = convert_tree(tree, TreeLike) if node_name is None: node_name = _node_name_default + _parent = generics.parent.dispatch(type(tree)) + if isinstance(node_shape, str): node_shape = DEFAULT_SHAPES[node_shape] @@ -384,7 +409,7 @@ def to_mermaid( # Output nodes nodes = [] # Stop automatic garbage collecting - for node, _ in tree.nodes.levelorder(keep=PreventCycles() & keep): + for node, _ in levelorder(tree, keep=PreventCycles() & keep): left, right = _get_shape(node_shape, node) name = node_name(node) if node_label: @@ -398,18 +423,18 @@ def to_mermaid( nodes = iter(nodes) next(nodes) for node in nodes: - arrow = edge_arrow(node.parent, node) if callable(edge_arrow) else edge_arrow - parent = node_name(node.parent) + arrow = edge_arrow(_parent(node), node) if callable(edge_arrow) else edge_arrow + par = node_name(_parent(node)) child = node_name(node) - file.write(f"{parent}{arrow}{child};\n") + file.write(f"{par}{arrow}{child};\n") @_wrap_file def to_latex( - tree, + tree: DownTreeLike, file=None, keep=DEFAULT_PREDICATE, - node_label: Union[str, Callable[[TNode], str], None] = str, + node_label: Union[str, Callable[[TNode], str], None] = generics.label, node_shape: TShape = None, leaf_distance: Optional[str] = "2em", level_distance: Optional[str] = None, @@ -424,7 +449,6 @@ def to_latex( Make sure to put ``\\usepackage{tikz}`` in your preamble. Does not wrap output in a figure environment. """ - tree = DownTree.convert(tree) if isinstance(indent, int): indent = "\t" if indent == -1 else indent * " " @@ -445,12 +469,12 @@ def to_latex( node_options.append("draw") depth = 0 - label = _escape_string(node_label(tree), "latex") + my_label = _escape_string(node_label(tree), "latex") file.write(rf"\begin{{tikzpicture}}{_latex_options(tree, picture_options)}") file.write("\n") options = _latex_options(tree, node_options) - file.write(rf"\node{options}{{{label}}} [grow={graph_direction}]") - for node, item in tree.descendants.preorder(keep=keep): + file.write(rf"\node{options}{{{my_label}}} [grow={graph_direction}]") + for node, item in preorder(tree, keep=keep, include_root=False): if item.depth > depth: file.write("\n") else: @@ -460,9 +484,9 @@ def to_latex( depth = item.depth file.write(depth * indent) - label = _escape_string(node_label(node), "latex") + text = _escape_string(node_label(node), "latex") options = _latex_options(tree, node_options) - file.write(f"child {{node{options} {{{label}}}") + file.write(f"child {{node{options} {{{text}}}") # Close final leaf node on same line if depth: @@ -498,10 +522,12 @@ def _sibling_distances(tree, stop=100): It assumes sibling distance is constant for each level. Parameter stop is used to prevent infinite recursion. """ + _children = generics.children.dispatch(type(tree)) + level_ranks = [] - for level, _ in zip(tree.levels, range(stop)): + for level, _ in zip(levels(tree), range(stop)): cousins = itertools.pairwise(level) - mid_ranks = [(len(n1.children) + len(n2.children)) / 2 for n1, n2 in cousins] + mid_ranks = [(len(_children(n1)) + len(_children(n2))) / 2 for n1, n2 in cousins] level_ranks.append(max(max(mid_ranks), 1) if mid_ranks else 1) distances = len(level_ranks) * [1] @@ -510,8 +536,8 @@ def _sibling_distances(tree, stop=100): return distances -def _node_name_default(node: DownTree): - return hex(node.nid) +def _node_name_default(node: Any): + return hex(nid(node)) def _get_shape(shape_factory, node): diff --git a/src/abstracttree/generics.py b/src/abstracttree/generics.py new file mode 100644 index 0000000..15c668c --- /dev/null +++ b/src/abstracttree/generics.py @@ -0,0 +1,278 @@ +import ast +import os +import xml.etree.ElementTree as ET +import zipfile +from abc import ABCMeta +from collections import namedtuple +from collections.abc import Sequence, Mapping, Collection +from functools import singledispatch +from pathlib import Path +from typing import TypeVar, Optional, Union, Any + +BaseString = Union[str, bytes, bytearray] +_BaseString = [str, bytes, bytearray] # Support py3.10 +BasePath = Union[Path, zipfile.Path] +MappingItem = namedtuple("MappingItem", ["key", "value"]) + + + +class DownTreeLike(metaclass=ABCMeta): + """Any object that has an identifiable parent and/or identifiable children.""" + @classmethod + def __subclasshook__(cls, subclass): + has_children = (hasattr(subclass, "children") + or children.dispatch(object) is not children.dispatch(subclass)) + return has_children + + +class TreeLike(metaclass=ABCMeta): + """Any object that has an identifiable parent.""" + @classmethod + def __subclasshook__(cls, subclass): + has_parent = (hasattr(subclass, "parent") + or parent.dispatch(object) is not parent.dispatch(subclass)) + return has_parent and issubclass(subclass, DownTreeLike) + + +T = TypeVar("T", bound=TreeLike) +DT = TypeVar("DT", bound=DownTreeLike) + + +# Base cases +@singledispatch +def children(tree: DT) -> Collection[DT]: + """Returns children of any downtreelike-object.""" + try: + return tree.children + except AttributeError: + raise TypeError(f"{type(tree)} is not DownTreeLike. children(x) is not defined.") from None + +@singledispatch +def parent(tree: T) -> Optional[T]: + """Returns parent of any treelike-object.""" + try: + return tree.parent + except AttributeError: + raise TypeError(f"{type(tree)} is not TreeLike. parent(x) is not defined.") from None + +@singledispatch +def parents(tree: T) -> Sequence[T]: + """Like parent(tree) but return value as a sequence.""" + tree_parent = parent(tree) + if tree_parent is not None: + return (tree_parent,) + else: + return () + +@singledispatch +def root(node: T) -> T: + """Find the root of a node in a tree.""" + parent_ = parent.dispatch(type(node)) + maybe_parent = parent_(node) + while maybe_parent is not None: + node, maybe_parent = maybe_parent, parent_(maybe_parent) + return node + +@singledispatch +def label(node: object) -> str: + """Return a string representation of this node. + + This representation should always represent just the Node. + If the node has parents or children these should be omitted. + """ + ... +label.register(object, str) + +@singledispatch +def nid(node: Any): + """Unique idenitifier for node. + + Usually the same as id, but can be overwritten for classes that act as delegates. + """ + try: + return node.nid + except AttributeError: + return id(node) + + +# Collections (Handle Mapping, Sequence and BaseString together to allow specialisation). +@children.register +def _(coll: Collection): + match coll: + case Mapping(): + return [MappingItem(k, v) for k, v in coll.items()] + case MappingItem(value=value): + if isinstance(value, Collection) and not isinstance(value, BaseString): + return children(value) + else: + return [value] + case Collection() if not isinstance(coll, BaseString): + return coll + case _: + return () + +@label.register +def _(coll: Collection): + """In python a type can have multiple parent classes.""" + match coll: + case MappingItem(key=key): + return str(key) + case Collection() if not isinstance(coll, BaseString): + cls_name = type(coll).__name__ + return f"{cls_name}[{len(coll)}]" + case _: + return str(coll) + + +# BaseString (should not be treated as a collection). +# children.register(BaseString, children.dispatch(object)) # if py >= 3.11 +for cls in _BaseString: + children.register(cls, children.dispatch(object)) + + +# Types +@children.register +def _(cls: type): + # We need this static way of calling it, to make it work on type itself. + return type(cls).__subclasses__(cls) + +@parents.register +def _(cls: type): + """In python a type can have multiple parent classes. Therefore, parent is not defined.""" + return cls.__bases__ + +@label.register +def _(cls: type): + return cls.__qualname__ + +@root.register +def _(_: type) -> type: + return object + + +# BasePath +@children.register(Path) +@children.register(zipfile.Path) +def _(pth: BasePath): + if pth.is_dir(): + try: + return tuple(pth.iterdir()) + except PermissionError: + # Print error and continue + import traceback + traceback.print_exc() + return () + else: + return () + +@parent.register(Path) +@parent.register(zipfile.Path) +def _(pth: BasePath): + parent_path = pth.parent + if pth != parent_path: + return parent_path + else: + return None + +@label.register(Path) +@label.register(zipfile.Path) +def _(pth: BasePath): + return pth.name + +@root.register(Path) +@root.register(zipfile.Path) +def _(pth: BasePath): + return pth.anchor + + +# PathLike (not fully supported, except for Path) +@nid.register +def _(pth: os.PathLike): + # Some paths are circular. nid can be used to stop recursion in those cases. + try: + st = os.lstat(pth) + except (FileNotFoundError, AttributeError): + return id(pth) # Fall-back + else: + return -st.st_ino + +@label.register +def _(pth: os.PathLike): + return os.path.basename(pth) + + +# AST +AST_SINGLETON = Union[ast.expr_context, ast.boolop, ast.operator, ast.unaryop, ast.cmpop] +AST_CONT = "↓" + +@children.register +def _(node: ast.AST): + return tuple(child for child in ast.iter_child_nodes(node) + if isinstance(child, ast.AST) and not isinstance(child, AST_SINGLETON)) + +@label.register +def _(node: ast.AST): + def format_value(field): + if isinstance(field, ast.AST): + if isinstance(field, AST_SINGLETON): + field_str = ast.dump(field) + else: + field_str = AST_CONT + elif isinstance(field, Sequence) and not isinstance(field, BaseString): + field = [format_value(f) for f in field] + field_str = "[" + ", ".join(field) + "]" + else: + field_str = repr(field) + return field_str + + node_children = children(node) + + if not node_children: + return ast.dump(node) + else: + # format_value = format_value(node) + args = [f"{name}={format_value(field)}" for name, field in ast.iter_fields(node)] + joined_args = ", ".join(args) + return f"{type(node).__name__}({joined_args})" + +# Exception group (python 3.11 or higher) +try: + ExceptionGroup +except NameError: + pass +else: + @children.register(BaseExceptionGroup) + def _(group: Exception): + if isinstance(group, BaseExceptionGroup): + return group.exceptions + else: + return () + + @label.register(BaseExceptionGroup) + def _(group: Exception): + if isinstance(group, BaseExceptionGroup): + return f"{type(group).__qualname__}({group.message})" + else: + return repr(group) + + +# XML / ElementTree.Element +@children.register +def _(element: ET.Element): + return element + +@label.register +def _(element: ET.Element): + output = [f"<{element.tag}"] + + for k, v in element.items(): + output.append(f" {k}={v!r}") + output.append(">") + + if text := element.text and str(element.text).strip(): + output.append(text) + output.append(f"") + if tail := element.tail and str(element.tail).strip(): + output.append(tail) + + return "".join(output) diff --git a/src/abstracttree/generics.pyi b/src/abstracttree/generics.pyi new file mode 100644 index 0000000..dcff5b9 --- /dev/null +++ b/src/abstracttree/generics.pyi @@ -0,0 +1,60 @@ +""" +Provide static typing for generics.py + +Unfortunately, static type checkers can't handle __subclasshook__. +Therefore, type stubs are provided for TreeLike and DownTreeLike. +""" + +import ast +import zipfile +from collections.abc import Collection, Sequence +from functools import singledispatch +from pathlib import Path +from typing import Self, TypeVar, Optional, Any, Protocol, overload, Never + + +class _HasChildren(Protocol): + @property + def children(self) -> Collection[Self]: + ... + + +class _HasChildrenAndParent(_HasChildren): + @property + def parent(self) -> Self: + ... + +TreeLike = _HasChildrenAndParent | Path | zipfile.Path | type +DownTreeLike = TreeLike | _HasChildren | Collection | ast.AST + +try: + DownTreeLike |= ExceptionGroup +except AttributeError: + pass + +T = TypeVar("T", bound=TreeLike) +DT = TypeVar("DT", bound=DownTreeLike) + + +@singledispatch +@overload +def children(tree: str | bytes | bytearray) -> Never: ... +@overload +def children(tree: Collection) -> Any: ... +@overload +def children(tree: DT) -> Collection[DT]: ... + +@singledispatch +def parent(tree: T) -> Optional[T]: ... + +@singledispatch +def parents(tree: T) -> Sequence[T]: ... + +@singledispatch +def root(node: T) -> T: ... + +@singledispatch +def label(node: object) -> str: ... + +@singledispatch +def nid(node: Any): ... diff --git a/src/abstracttree/heaptree.py b/src/abstracttree/heaptree.py deleted file mode 100644 index d133ebf..0000000 --- a/src/abstracttree/heaptree.py +++ /dev/null @@ -1,78 +0,0 @@ -from typing import Collection, Optional - -from .binarytree import BinaryTree -from .tree import TNode - - -class HeapTree(BinaryTree): - """Provides a tree interface to a heap. - - Mainly useful for visualisation purposes. - >>> from abstracttree import print_tree - >>> import heapq - >>> tree = HeapTree() - >>> for n in range(5, 0, -1): - >>> heapq.heappush(tree.heap, n) - >>> print_tree(tree) - 0 → 1 - ├─ 1 → 2 - │ ├─ 3 → 5 - │ └─ 4 → 4 - └─ 2 → 3 - """ - - __slots__ = "heap", "index" - - def __init__(self, heap=None, index=0): - if heap is None: - heap = [] - self.heap = heap - self.index = index - - def __repr__(self): - return f"{type(self).__qualname__}{self.heap, self.index}" - - def __str__(self): - try: - return f"{self.index} → {self.value}" - except IndexError: - return repr(self) - - @property - def nid(self): - return self.index - - def eqv(self, other): - return type(self) is type(other) and self.index == other.index - - @property - def children(self: TNode) -> Collection[TNode]: - return [ - HeapTree(self.heap, i) - for i in range(2 * self.index + 1, 2 * self.index + 3) - if i < len(self.heap) - ] - - @property - def left_child(self) -> Optional[TNode]: - i = 2 * self.index + 1 - if i < len(self.heap): - return HeapTree(self.heap, i) - - @property - def right_child(self) -> Optional[TNode]: - i = 2 * self.index + 2 - if i < len(self.heap): - return HeapTree(self.heap, i) - - @property - def parent(self: TNode) -> Optional[TNode]: - n = self.index - if n != 0: - return HeapTree(self.heap, (n - 1) // 2) - else: - return None - - @property - def value(self): - return self.heap[self.index] diff --git a/src/abstracttree/iterators.py b/src/abstracttree/iterators.py new file mode 100644 index 0000000..df4a090 --- /dev/null +++ b/src/abstracttree/iterators.py @@ -0,0 +1,192 @@ +import itertools +from collections import namedtuple, deque +from collections.abc import Iterator, Sequence +from typing import TypeVar + +import abstracttree.generics as generics + +NodeItem = namedtuple("NodeItem", ["index", "depth"]) + +T = TypeVar("T", bound=generics.TreeLike) +DT = TypeVar("DT", bound=generics.DownTreeLike) + + +def ancestors(node: T) -> Iterator[T]: + """Iterate through ancestors of node.""" + parent = generics.parent.dispatch(type(node)) + while (node := parent(node)) is not None: + yield node + + +def path(node: T, reverse=False) -> Iterator[T]: + """Iterate through path of node.""" + if not reverse: + return itertools.chain(reversed(list(ancestors(node))), [node]) + else: + return itertools.chain([node], reversed(list(ancestors(node)))) + + +def nodes(tree: DT, include_root=True) -> Iterator[DT]: + """Iterate through all nodes of this tree. + + The order of nodes might change between versions. + Use methods like preorder, postorder, levelorder if the order is important. + """ + children = generics.children.dispatch(type(tree)) + coll = deque([tree] if include_root else children(tree)) + while coll: + yield (node := coll.pop()) + coll.extend(children(node)) + + +def descendants(node: DT) -> Iterator[DT]: + """Iterate through descendants of node in no particular order.""" + return nodes(node, include_root=False) + + +def preorder(tree: DT, keep=None, include_root=True) -> Iterator[tuple[DT, NodeItem]]: + """Iterate through nodes in pre-order. + + Only descend where keep(node). + Returns tuples (node, item) + Item denotes depth of iteration and index of child. + """ + children = generics.children.dispatch(type(tree)) + if include_root: + coll = deque([(tree, NodeItem(None, 0))]) + else: + coll = deque(reversed([(c, NodeItem(i, 1)) for i, c in enumerate(children(tree))])) + + while coll: + node, item = coll.pop() + if not keep or keep(node, item): + yield node, item + next_nodes = [(c, NodeItem(i, item.depth + 1)) for i, c in enumerate(children(node))] + coll.extend(reversed(next_nodes)) + + +def postorder(tree: DT, keep=None, include_root=True) -> Iterator[tuple[DT, NodeItem]]: + """Iterate through nodes in post-order. + + Only descend where keep(node). + Returns tuples (node, item) + Item denotes depth of iteration and index of child. + """ + children = generics.children.dispatch(type(tree)) + + if include_root: + coll = iter([(tree, NodeItem(None, 0))]) + else: + coll = iter([(c, NodeItem(i, 1)) for i, c in enumerate(children(tree))]) + + node, item = next(coll, (None, None)) + stack = [] + + while node or stack: + # Go down + keep_node = keep is None or keep(node, item) + while keep_node and (cc := children(node)): + stack.append((node, item, coll)) + coll = iter([ + (c, NodeItem(i, item.depth + 1)) for (i, c) in enumerate(cc) + ]) + node, item = next(coll) + keep_node = keep is None or keep(node, item) + if keep_node: + yield node, item + + # Go right or go up + node, item = next(coll, (None, None)) + while node is None and stack: + node, item, coll = stack.pop() + yield node, item + node, item = next(coll, (None, None)) + + +def levelorder(tree: DT, keep=None, include_root=True) -> Iterator[tuple[DT, NodeItem]]: + """Iterate through nodes in level-order. + + Only descend where keep(node). + Returns tuples (node, item) + Item denotes depth of iteration and index of child. + """ + children = generics.children.dispatch(type(tree)) + if include_root: + coll = deque([(tree, NodeItem(None, 0))]) + else: + coll = deque([(c, NodeItem(i, 1)) for i, c in enumerate(children(tree))]) + + while coll: + node, item = coll.popleft() + if not keep or keep(node, item): + yield node, item + next_nodes = [(c, NodeItem(i, item.depth + 1)) for i, c in enumerate(children(node))] + coll.extend(next_nodes) + + +def leaves(tree: DT) -> Iterator[DT]: + """Iterate through leaves of node.""" + children = generics.children.dispatch(type(tree)) + for node in nodes(tree): + if not children(node): + yield node + + +def siblings(node: T) -> Iterator[T]: + """Iterate through siblings of node.""" + nid = generics.nid.dispatch(type(node)) + node_nid = nid(node) + if p := generics.parent(node): + return (child for child in generics.children(p) if node_nid != nid(child)) + else: + return iter(()) + + +def levels(tree: DT) -> Iterator[Sequence[DT]]: + """Iterate through descendants in levels.""" + children = generics.children.dispatch(type(tree)) + level = [tree] + while level: + yield iter(level) + level = [child for node in level for child in children(node)] + + +def levels_zigzag(tree: DT) -> Iterator[Sequence[DT]]: + """Iterate through descendants in levels in zigzag-order.""" + children = generics.children.dispatch(type(tree)) + level = [tree] + while level: + yield iter(level) + level = [child for node in reversed(level) for child in reversed(children(node))] + + +def _inorder(node, keep=None, index=None, depth=0): + """Iterate through nodes of BinaryTree. + + This requires node to be a BinaryTree. + It's less generic than the other methods. Therefore, it's not exported by default. + It can be accessed through mixins.views.BinaryNodes. + """ + stack = deque() + item = NodeItem(index, depth) + + while node is not None or stack: + # Traverse down/left + left_child, left_item = node.left_child, NodeItem(0, depth + 1) + while left_child is not None and (not keep or keep(left_child, left_item)): + stack.append((node, item)) + node, item, depth = left_child, left_item, depth + 1 + left_child, left_item = node.left_child, NodeItem(0, depth + 1) + + yield node, item + + # Traverse right/up + right_child, right_item = node.right_child, NodeItem(1, depth + 1) + while stack and ( + right_child is None or (keep and not keep(right_child, right_item)) + ): + node, item = stack.pop() + yield node, item + depth -= 1 + right_child, right_item = node.right_child, NodeItem(1, depth + 1) + node, item, depth = right_child, right_item, depth + 1 diff --git a/src/abstracttree/mixins/__init__.py b/src/abstracttree/mixins/__init__.py new file mode 100644 index 0000000..281e78e --- /dev/null +++ b/src/abstracttree/mixins/__init__.py @@ -0,0 +1,10 @@ +__all__ = [ + "Tree", + "DownTree", + "MutableTree", + "MutableDownTree", + "BinaryDownTree", + "BinaryTree", +] + +from .trees import Tree, DownTree, MutableDownTree, MutableTree, BinaryDownTree, BinaryTree diff --git a/src/abstracttree/mixins/trees.py b/src/abstracttree/mixins/trees.py new file mode 100644 index 0000000..fa598e5 --- /dev/null +++ b/src/abstracttree/mixins/trees.py @@ -0,0 +1,207 @@ +import operator +from abc import abstractmethod, ABCMeta +from typing import TypeVar, Callable, Optional, Collection, Literal, Iterable, Sequence + +from .views import AncestorsView, PathView, NodesView, LeavesView, LevelsView, SiblingsView, BinaryNodesView +from .. import generics + +TNode = TypeVar("TNode") +TMutDownNode = TypeVar("TMutDownNode", bound="MutableDownTree") +Order = Literal["pre", "post", "level"] + + +class AbstractTree(metaclass=ABCMeta): + """Most abstract baseclass for everything.""" + + __slots__ = () + + @classmethod + def convert(cls, obj): + """Convert obj to tree-type or raise TypeError if that doesn't work. + + Maybe just use `as_tree(x) -> Tree` instead. + """ + from ..adapters import convert_tree + return convert_tree(obj, cls) + + @property + def nid(self) -> int: + """Unique number that represents this node.""" + return id(self) + + +class UpTree(AbstractTree, metaclass=ABCMeta): + """Abstract class for tree classes with parent but no children.""" + + __slots__ = () + + @property + @abstractmethod + def parent(self: TNode) -> Optional[TNode]: + """Parent of this node or None if root.""" + return None + + @property + def is_root(self) -> bool: + """Whether this node is a root (has no parent).""" + return self.parent is None + + @property + def root(self) -> TNode: + """Root of tree.""" + p, p2 = self, self.parent + while p2: + p, p2 = p2, p2.parent + return p + + @property + def ancestors(self): + """View of ancestors of node.""" + return AncestorsView(self) + + @property + def path(self): + """View of path from root to node.""" + return PathView(self) + + +class DownTree(AbstractTree, metaclass=ABCMeta): + """Abstract class for tree classes with children but no parent.""" + + __slots__ = () + + @property + @abstractmethod + def children(self: TNode) -> Collection[TNode]: + """Children of this node.""" + return () + + @property + def leaves(self): + """View of leaves from this node.""" + return LeavesView(self) + + @property + def nodes(self): + """View of this node and its descendants.""" + return NodesView(self) + + @property + def descendants(self): + """View of descendants of this node.""" + return NodesView(self, include_root=False) + + @property + def levels(self): + """View of this node and descendants by level.""" + return LevelsView(self) + + @property + def is_leaf(self) -> bool: + """Whether this node is a leaf (does not have children).""" + return not self.children + + def transform(self: TNode, f: Callable[[TNode], TMutDownNode], keep=None) -> TMutDownNode: + """Create new tree where each node of self is transformed by f.""" + stack = [] + for node, item in self.descendants.postorder(keep=keep): + depth = item.depth + while len(stack) < depth: + stack.append(list()) + stack[depth - 1].append(new := f(node)) + if len(stack) > depth: + new.add_children(stack.pop(-1)) + new = f(self) + if stack: + new.add_children(stack.pop()) + return new + + +class MutableDownTree(DownTree, metaclass=ABCMeta): + """Abstract class for mutable tree with children.""" + + __slots__ = () + + @abstractmethod + def add_child(self, node: TNode): + """Add node to children.""" + raise NotImplementedError + + @abstractmethod + def remove_child(self, node: TNode): + """Remove node from children.""" + raise NotImplementedError + + def add_children(self, children: Iterable[TNode]): + """Add multiple nodes to children.""" + for child in children: + self.add_child(child) + + +class Tree(UpTree, DownTree, metaclass=ABCMeta): + """Abstract class for tree classes with access to children and parents.""" + + __slots__ = () + + @property + def siblings(self): + """View of siblings of this node.""" + return SiblingsView(self) + + +class MutableTree(Tree, MutableDownTree, metaclass=ABCMeta): + """Abstract class for mutable tree with children and parent.""" + + __slots__ = () + + def detach(self) -> TNode: + """Remove parent if any and return self.""" + if p := self.parent: + p.remove_child(self) + return self + + +class BinaryDownTree(DownTree, metaclass=ABCMeta): + """Binary-tree with links to children.""" + + __slots__ = () + + @property + @abstractmethod + def left_child(self) -> Optional[TNode]: + return None + + @property + @abstractmethod + def right_child(self) -> Optional[TNode]: + return None + + @property + def children(self) -> Sequence[TNode]: + nodes = list() + if self.left_child is not None: + nodes.append(self.left_child) + if self.right_child is not None: + nodes.append(self.right_child) + return nodes + + @property + def nodes(self): + return BinaryNodesView(self) + + @property + def descendants(self): + return BinaryNodesView(self, include_root=False) + + +class BinaryTree(BinaryDownTree, Tree, metaclass=ABCMeta): + """Binary-tree with links to children and to parent.""" + + __slots__ = () + + +# Some optimizations +generics.children.register(DownTree, operator.attrgetter("children")) +generics.parent.register(UpTree, operator.attrgetter("parent")) +generics.nid.register(AbstractTree, operator.attrgetter("nid")) +generics.label.register(AbstractTree, str) diff --git a/src/abstracttree/mixins/views.py b/src/abstracttree/mixins/views.py new file mode 100644 index 0000000..d571b89 --- /dev/null +++ b/src/abstracttree/mixins/views.py @@ -0,0 +1,144 @@ +import itertools +from abc import ABCMeta +from collections.abc import Iterator +from typing import Iterable, TypeVar + +from .. import iterators as _iterators + +T = TypeVar("T", bound="Tree") + + +class TreeView(Iterable[T], metaclass=ABCMeta): + __slots__ = "_node" + itr_method = None + + def __init__(self, node: T): + self._node: T = node + + def __iter__(self) -> Iterator[T]: + return type(self).itr_method(self._node) + + def __bool__(self) -> bool: + try: + next(iter(self)) + except StopIteration: + return False + else: + return True + + def count(self) -> int: + """Count number of nodes in this view.""" + return _ilen(self) + + +class AncestorsView(TreeView): + """View over ancestors.""" + __slots__ = () + itr_method = _iterators.ancestors + + +class PathView(TreeView): + """View over path from root to self.""" + __slots__ = () + itr_method = _iterators.path + + def __bool__(self): + # A path always contains at least one node. No need to check. + return True + + def __contains__(self, item): + return item in reversed(self) + + def __reversed__(self): + return _iterators.path(self._node, reverse=True) + + def count(self): + return _ilen(reversed(self)) + + +class NodesView(TreeView): + """View over nodes.""" + __slots__ = "_include_root" + itr_method = _iterators.nodes + + def __init__(self, node, include_root: bool = True): + super().__init__(node) + self._include_root = include_root + + def __iter__(self): + return _iterators.nodes(self._node, include_root=self._include_root) + + def preorder(self, keep=None): + return _iterators.preorder(self._node, keep, include_root=self._include_root) + + def postorder(self, keep=None): + return _iterators.postorder(self._node, keep, include_root=self._include_root) + + def levelorder(self, keep=None): + return _iterators.levelorder(self._node, keep, include_root=self._include_root) + + +class LeavesView(TreeView): + """View over leaves.""" + __slots__ = () + itr_method = _iterators.leaves + + +class LevelsView(TreeView): + """View over levels.""" + __slots__ = () + itr_method = _iterators.levels + + def __bool__(self): + return True + + def zigzag(self): + """Zig-zag through levels.""" + return _iterators.levels_zigzag(self._node) + + +class SiblingsView(TreeView): + """View over siblings.""" + __slots__ = () + itr_method = _iterators.siblings + + def __contains__(self, other): + try: + other_parent = other.parent + except AttributeError: + return False # not a Tree + else: + return other_parent is self._node.parent and other is not self._node + + def __len__(self): + if p := self._node.parent: + return len(p.children) - 1 + return 0 + + count = __len__ + + +class BinaryNodesView(NodesView): + __slots__ = () + def inorder(self, keep=None): + """ + Iterate through nodes in inorder (traverse left, yield root, traverse right). + + Note: + - `item.index` will be 0 for every left child + - `item.index` will be 1 for every right child (even if node.left_child is None) + - `item.index` will be None for the root (even if root is a subtree) + This is a bit different from how `preorder()`, `postorder()` and `levelorder()` work, + because those functions always give index 0 to the first child, + regardless of whether it's a left or right child. + """ + if self._include_root: + yield from _iterators._inorder(self._node, keep, index=None, depth=0) + else: + yield from _iterators._inorder(self._node.left_child, keep, index=0, depth=1) + yield from _iterators._inorder(self._node.right_child, keep, index=1, depth=1) + + +def _ilen(itr): + """Recipe from https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.ilen""" + return sum(itertools.compress(itertools.repeat(1), zip(itr))) diff --git a/src/abstracttree/predicates.py b/src/abstracttree/predicates.py index ba924ab..e5f3194 100644 --- a/src/abstracttree/predicates.py +++ b/src/abstracttree/predicates.py @@ -1,10 +1,14 @@ from dataclasses import dataclass -from typing import Callable +from typing import Callable, TypeVar -from .tree import AbstractTree, NodeItem +from .iterators import NodeItem +from .generics import TreeLike, DownTreeLike, nid, parent +T = TypeVar("T", bound=TreeLike) +DT = TypeVar("DT", bound=DownTreeLike) -class Predicate(Callable[[AbstractTree, NodeItem], bool]): + +class Predicate(Callable[[T, NodeItem], bool]): __slots__ = () def __or__(self, other): @@ -26,7 +30,7 @@ class PredicateUnion(Predicate): def __init__(self, *preds): self.preds = preds - def __call__(self, node: AbstractTree, item: NodeItem): + def __call__(self, node: DT, item: NodeItem): return any(pred(node, item) for pred in self.preds) @@ -36,7 +40,7 @@ class PredicateIntersection(Predicate): def __init__(self, *preds): self.preds = preds - def __call__(self, node: AbstractTree, item: NodeItem): + def __call__(self, node: DT, item: NodeItem): return all(pred(node, item) for pred in self.preds) @@ -48,11 +52,12 @@ class RemoveDuplicates(Predicate): def __init__(self): self.seen = set() - def __call__(self, node: AbstractTree, item): - if node.nid in self.seen: + def __call__(self, node: DT, item): + node_nid = nid(node) + if node_nid in self.seen: return False else: - self.seen.add(node.nid) + self.seen.add(node_nid) return True @@ -69,13 +74,16 @@ def __init__(self): self.seen = set() self.duplicates = set() - def __call__(self, node: AbstractTree, item): - if node.parent is not None and node.parent.nid in self.duplicates and node.nid in self.seen: + def __call__(self, node: DT, item): + p = parent(node) + node_nid = nid(node) + + if p is not None and nid(p) in self.duplicates and node_nid in self.seen: return False - if node.nid in self.seen: - self.duplicates.add(node.nid) + if node_nid in self.seen: + self.duplicates.add(node_nid) else: - self.seen.add(node.nid) + self.seen.add(node_nid) return True diff --git a/src/abstracttree/route.py b/src/abstracttree/route.py index f9246bc..7842b5c 100644 --- a/src/abstracttree/route.py +++ b/src/abstracttree/route.py @@ -5,9 +5,10 @@ from functools import lru_cache from typing import TypeVar, Optional -from .tree import UpTree +from . import iterators as _iterators +from .generics import TreeLike, nid -TNode = TypeVar("TNode", bound=UpTree) +TNode = TypeVar("TNode", bound=TreeLike) class Route: @@ -19,7 +20,7 @@ class Route: __slots__ = "_apaths", "_lca" - def __init__(self, *anchors: TNode): + def __init__(self, *anchors: TreeLike): """Create a route through a few nodes. All nodes should belong to the same tree. @@ -31,21 +32,22 @@ def __init__(self, *anchors: TNode): self.add_anchor(anchor) def __repr__(self): - nodes_str = ", ".join([str(path[-1].identifier) for path in self._apaths]) + nodes_str = ", ".join([repr(p[-1]) for p in self._apaths]) return f"{self.__class__.__name__}({nodes_str})" - def add_anchor(self, anchor: TNode): + def add_anchor(self, anchor: TreeLike): """Add a node to the route. The node should belong to the same tree as any existing anchor nodes. """ self._lca = None - path = tuple(anchor.path) + anchor_path = list(_iterators.path(anchor)) apaths = self._apaths - if apaths and not apaths[0][0].eqv(path[0]): - raise ValueError("Different tree!") + + if not apaths or nid(apaths[0][0]) == nid(anchor_path[0]): + apaths.append(anchor_path) else: - apaths.append(path) + raise ValueError("Different tree!") @property def anchors(self): @@ -71,7 +73,7 @@ def lca(self) -> Optional[TNode]: if i := bisect( indices, False, - key=lambda ind: any(not path0[ind].eqv(p[ind]) for p in paths), + key=lambda ind: any(nid(path0[ind]) != nid(p[ind]) for p in paths), ): lca = self._lca = path0[i - 1] return lca @@ -82,7 +84,7 @@ def lca(self) -> Optional[TNode]: def _common2(self, i, j) -> int: path_i, path_j = self._apaths[i], self._apaths[j] indices = range(min(len(path_i), len(path_j))) - return bisect(indices, False, key=lambda ind: not path_i[ind].eqv(path_j[ind])) - 1 + return bisect(indices, False, key=lambda ind: nid(path_i[ind]) != nid(path_j[ind])) - 1 class RouteView(Sized, metaclass=ABCMeta): diff --git a/src/abstracttree/tree.py b/src/abstracttree/tree.py deleted file mode 100644 index 7cee4a0..0000000 --- a/src/abstracttree/tree.py +++ /dev/null @@ -1,372 +0,0 @@ -import itertools -from abc import abstractmethod, ABCMeta -from collections import deque, namedtuple -from typing import TypeVar, Callable, Optional, Collection, Literal, Iterable - -TNode = TypeVar("TNode") -TMutDownNode = TypeVar("TMutDownNode", bound="MutableDownTree") -Order = Literal["pre", "post", "level"] -NodeItem = namedtuple("NodeItem", ["index", "depth"]) -NodePredicate = Callable[[TNode, NodeItem], bool] - - -class AbstractTree(metaclass=ABCMeta): - """Most abstract baseclass for everything.""" - - __slots__ = () - - @classmethod - def convert(cls, obj): - """Convert obj to tree-type or raise TypeError if that doesn't work.""" - from .adapters import convert_tree - - if isinstance(obj, cls): - return obj - tree = convert_tree(obj) - if isinstance(tree, cls): - return tree - raise TypeError(f"{obj!r} cannot be converted to {cls.__name__}") - - @property - def nid(self) -> int: - """Unique number that represents this node.""" - return id(self) - - def eqv(self, other) -> bool: - """Check if both objects represent the same node. - - Should normally be operator.is, but can be overridden by delegates. - """ - return self is other - - -class UpTree(AbstractTree, metaclass=ABCMeta): - """Abstract class for tree classes with parent but no children.""" - - __slots__ = () - - @property - @abstractmethod - def parent(self: TNode) -> Optional[TNode]: - """Parent of this node or None if root.""" - return None - - @property - def is_root(self) -> bool: - """Whether this node is a root (has no parent).""" - return self.parent is None - - @property - def root(self) -> TNode: - """Root of tree.""" - p, p2 = self, self.parent - while p2: - p, p2 = p2, p2.parent - return p - - @property - def ancestors(self): - """View of ancestors of node.""" - return AncestorsView(self.parent) - - @property - def path(self): - """View of path from root to node.""" - return PathView(self) - - -class DownTree(AbstractTree, metaclass=ABCMeta): - """Abstract class for tree classes with children but no parent.""" - - __slots__ = () - - @property - @abstractmethod - def children(self: TNode) -> Collection[TNode]: - """Children of this node.""" - return () - - @property - def leaves(self): - """View of leaves from this node.""" - return LeavesView(self) - - @property - def nodes(self): - """View of this node and its descendants.""" - return NodesView([self], 0) - - @property - def descendants(self): - """View of descendants of this node.""" - return NodesView(self.children, 1) - - @property - def levels(self): - """View of this node and descendants by level.""" - return LevelsView(self) - - @property - def is_leaf(self) -> bool: - """Whether this node is a leaf (does not have children).""" - return not self.children - - def transform(self: TNode, f: Callable[[TNode], TMutDownNode], keep=None) -> TMutDownNode: - """Create new tree where each node of self is transformed by f.""" - stack = [] - for node, item in self.descendants.postorder(keep=keep): - depth = item.depth - while len(stack) < depth: - stack.append(list()) - stack[depth - 1].append(new := f(node)) - if len(stack) > depth: - new.add_children(stack.pop(-1)) - new = f(self) - if stack: - new.add_children(stack.pop()) - return new - - -class MutableDownTree(DownTree, metaclass=ABCMeta): - """Abstract class for mutable tree with children.""" - - __slots__ = () - - @abstractmethod - def add_child(self, node: TNode): - """Add node to children.""" - raise NotImplementedError - - @abstractmethod - def remove_child(self, node: TNode): - """Remove node from children.""" - raise NotImplementedError - - def add_children(self, children: Iterable[TNode]): - """Add multiple nodes to children.""" - for child in children: - self.add_child(child) - - -class Tree(UpTree, DownTree, metaclass=ABCMeta): - """Abstract class for tree classes with access to children and parents.""" - - __slots__ = () - - @property - def siblings(self): - """View of siblings of this node.""" - return SiblingsView(self) - - -class MutableTree(Tree, MutableDownTree, metaclass=ABCMeta): - """Abstract class for mutable tree with children and parent.""" - - __slots__ = () - - def detach(self) -> TNode: - """Remove parent if any and return self.""" - if p := self.parent: - p.remove_child(self) - return self - - -class TreeView(Iterable[TNode], metaclass=ABCMeta): - __slots__ = () - - def count(self) -> int: - """Count number of nodes in this view.""" - counter = itertools.count() - deque(zip(self, counter), maxlen=0) - return next(counter) - - -class AncestorsView(TreeView): - __slots__ = "parent" - - def __init__(self, parent): - self.parent = parent - - def __iter__(self): - p = self.parent - while p: - yield p - p = p.parent - - def __bool__(self): - return bool(self.parent) - - -class PathView(TreeView): - __slots__ = "view" - - def __init__(self, node): - self.view = AncestorsView(node) - - def __iter__(self): - return reversed(list(self.view)) - - def __reversed__(self): - return iter(self.view) - - def __contains__(self, node): - return node in self.view - - def __bool__(self): - return True - - def count(self): - return self.view.count() - - -class NodesView(TreeView): - __slots__ = "nodes", "level" - - def __init__(self, nodes, level): - self.nodes, self.level = nodes, level - - def __bool__(self): - return bool(self.nodes) - - def __iter__(self): - nodes = deque(self.nodes) - while nodes: - yield (node := nodes.pop()) - nodes.extend(node.children) - - def preorder(self, keep=None): - """Iterate through nodes in pre-order. - - Only descend where keep(node). - Returns tuples (node, item) - Item denotes depth of iteration and index of child. - """ - nodes = deque((c, NodeItem(i, self.level)) for (i, c) in enumerate(self.nodes)) - while nodes: - node, item = nodes.popleft() - if not keep or keep(node, item): - yield node, item - next_nodes = [(c, NodeItem(i, item.depth + 1)) for i, c in enumerate(node.children)] - nodes.extendleft(reversed(next_nodes)) - - def postorder(self, keep=None): - """Iterate through nodes in post-order. - - Only descend where keep(node). - Returns tuples (node, item) - Item denotes depth of iteration and index of child. - """ - children = iter([(c, NodeItem(i, self.level)) for (i, c) in enumerate(self.nodes)]) - node, item = next(children, (None, None)) - stack = [] - - while node or stack: - # Go down - keep_node = keep is None or keep(node, item) - while keep_node and node.children: - stack.append((node, item, children)) - children = iter([ - (c, NodeItem(i, item.depth + 1)) for (i, c) in enumerate(node.children) - ]) - node, item = next(children) - keep_node = keep is None or keep(node, item) - if keep_node: - yield node, item - - # Go right or go up - node, item = next(children, (None, None)) - while node is None and stack: - node, item, children = stack.pop() - yield node, item - node, item = next(children, (None, None)) - - def levelorder(self, keep=None): - """Iterate through nodes in level-order. - - Only descend where keep(node). - Returns tuples (node, item) - Item denotes depth of iteration and index of child. - """ - nodes = deque((c, NodeItem(i, self.level)) for (i, c) in enumerate(self.nodes)) - while nodes: - node, item = nodes.popleft() - if not keep or keep(node, item): - yield node, item - next_nodes = [(c, NodeItem(i, item.depth + 1)) for i, c in enumerate(node.children)] - nodes.extend(next_nodes) - - -class LeavesView(TreeView): - __slots__ = "root" - - def __init__(self, root): - self.root = root - - def __bool__(self): - return True - - def __iter__(self): - for node in self.root.nodes: - if node.is_leaf: - yield node - - def __contains__(self, node): - if not node.is_leaf: - return False - try: - ancestors = node.ancestors - except AttributeError: - return node in super() - else: - return self.root in ancestors - - -class LevelsView: - __slots__ = "tree" - - def __init__(self, tree): - self.tree = tree - - def __bool__(self): - return True - - def __iter__(self): - level = [self.tree] - while level: - yield iter(level) - level = [child for node in level for child in node.children] - - def zigzag(self): - """Traverse the levels in zigzag-order.""" - level = [self.tree] - while level: - yield iter(level) - level = [child for node in reversed(level) for child in reversed(node.children)] - - def count(self): - return 1 + max(it.depth for (_, it) in self.tree.nodes.preorder()) - - -class SiblingsView(TreeView): - __slots__ = "node" - - def __init__(self, node): - self.node = node - - def __iter__(self): - node = self.node - parent = node.parent - if parent is not None: - return (child for child in parent.children if not node.eqv(child)) - else: - return iter(()) - - def __contains__(self, node): - return not self.node.eqv(node) and self.node.parent.eqv(node.parent) - - def __len__(self): - if p := self.node.parent: - return len(p.children) - 1 - return 0 - - count = __len__ diff --git a/src/abstracttree/utils.py b/src/abstracttree/utils.py new file mode 100644 index 0000000..baa3b63 --- /dev/null +++ b/src/abstracttree/utils.py @@ -0,0 +1,10 @@ +from abstracttree.generics import DT, nid + + +def eqv(n1: DT, n2: DT) -> bool: + """Whether two nodes are equivalent. + + For nodes to be equivalent, they need to have the same nid and be of the same type. + The result is almost the same as ``n1 is n2``, but can be overridden for adapters, symlinks etc. + """ + return nid(n1) == nid(n2) and type(n1) == type(n2) diff --git a/tests/test_downtree.py b/tests/test_downtree.py index a728332..c0af717 100644 --- a/tests/test_downtree.py +++ b/tests/test_downtree.py @@ -1,6 +1,9 @@ from unittest import TestCase -import tree_instances as trees +try: + from . import tree_instances as trees +except ImportError: + import tree_instances as trees class TestDownTree(TestCase): @@ -102,7 +105,7 @@ def test_preorder_nodes(self): self.assertEqual(expected, values) items = [item for _, item in tree.nodes.preorder()] - expected = [(0, 0), (0, 1), (1, 1), (0, 2), (0, 3)] + expected = [(None, 0), (0, 1), (1, 1), (0, 2), (0, 3)] self.assertEqual(expected, items) def test_preorder_descendants(self): @@ -122,7 +125,7 @@ def test_postorder_nodes(self): self.assertEqual(expected, values) items = [item for _, item in tree.nodes.postorder()] - expected = [(0, 1), (0, 3), (0, 2), (1, 1), (0, 0)] + expected = [(0, 1), (0, 3), (0, 2), (1, 1), (None, 0)] self.assertEqual(expected, items) def test_postorder_descendants(self): @@ -142,7 +145,7 @@ def test_levelorder_nodes(self): self.assertEqual(expected, values) items = [item for _, item in tree.nodes.levelorder()] - expected = [(0, 0), (0, 1), (1, 1), (0, 2), (0, 3)] + expected = [(None, 0), (0, 1), (1, 1), (0, 2), (0, 3)] self.assertEqual(expected, items) def test_levelorder_descendants(self): @@ -162,7 +165,7 @@ def test_inorder_nodes(self): self.assertEqual(expected, values) items = [item for _, item in tree.nodes.inorder()] - expected = [(0, 1), (0, 0), (0, 2), (1, 3), (1, 1)] + expected = [(0, 1), (None, 0), (0, 2), (1, 3), (1, 1)] self.assertEqual(expected, items) def test_inorder_descendants(self): diff --git a/tests/test_export.py b/tests/test_export.py index f095ee9..492ac72 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -2,10 +2,16 @@ from unittest import TestCase from abstracttree import to_string, to_mermaid, to_dot, MaxDepth, to_latex, to_image -from tree_instances import INFINITE_TREE, BINARY_TREE + +try: + from .tree_instances import INFINITE_TREE, BINARY_TREE +except ImportError: + from tree_instances import INFINITE_TREE, BINARY_TREE class TestExport(TestCase): + maxDiff = None + def test_to_string(self): result = to_string(INFINITE_TREE, keep=MaxDepth(2)) expected = 3 * "Infinite singleton!\n" diff --git a/tests/test_generics.py b/tests/test_generics.py new file mode 100644 index 0000000..0dc7f7e --- /dev/null +++ b/tests/test_generics.py @@ -0,0 +1,83 @@ +import ast +from pathlib import Path +from unittest import TestCase + +from abstracttree.generics import children, MappingItem, parent, label + +SKIPTOKEN = "" +ERRTOKEN = "" + + +class A: pass + + +class B(A): pass + + +class TestGenerics(TestCase): + def setUp(self): + self.cases = [ + [1, 2, 3], + {"a": 1, "b": "hello", ("c", "d"): [4, 5, 6]}, + MappingItem("b", "hello"), + MappingItem(("c", "d"), [4, 5, 6]), + A, + ast.parse("1 + a", mode="eval").body, + Path(__file__).parent.parent / ".github" + ] + + def test_children(self): + + expectations = [ + [1, 2, 3], + list({"a": 1, "b": "hello", ("c", "d"): [4, 5, 6]}.items()), + ["hello"], + [4, 5, 6], + [B], + SKIPTOKEN, #[ast.parse("1", mode="eval").body, ast.parse("a", mode="eval").body], + (Path(__file__).parent.parent / ".github/workflows",) + ] + + for case, expected in zip(self.cases, expectations, strict=True): + if expected is SKIPTOKEN: + continue + + with self.subTest(f"Test for {case}"): + self.assertEqual(expected, children(case)) + + def test_parent(self): + expectations = [ + ERRTOKEN, + ERRTOKEN, + ERRTOKEN, + ERRTOKEN, + ERRTOKEN, + ERRTOKEN, + Path(__file__).parent.parent, + ] + + for case, expected in zip(self.cases, expectations, strict=True): + with self.subTest(f"Test for {case}"): + if expected is not ERRTOKEN: + self.assertEqual(expected, parent(case)) + else: + with self.assertRaises(TypeError): + parent(case) + + + def test_label(self): + expectations = [ + "list[3]", + "dict[3]", + "b", + "('c', 'd')", + "A", + "BinOp(left=↓, op=Add(), right=↓)", + ".github", + ] + + for case, expected in zip(self.cases, expectations, strict=True): + if expected is SKIPTOKEN: + continue + with self.subTest(f"Test for {case}"): + self.assertEqual(expected, label(case)) diff --git a/tests/test_mutabletree.py b/tests/test_mutabletree.py index fabaaf7..f0bc5e0 100644 --- a/tests/test_mutabletree.py +++ b/tests/test_mutabletree.py @@ -1,7 +1,11 @@ from unittest import TestCase from abstracttree import MaxDepth -from tree_instances import BinaryNode, INFINITE_BINARY_TREE + +try: + from .tree_instances import BinaryNode, INFINITE_BINARY_TREE +except ImportError: + from tree_instances import BinaryNode, INFINITE_BINARY_TREE class TestMutableDownTree(TestCase): @@ -43,7 +47,7 @@ def test_transform_II(self): """ def double(node): - return BinaryNode(value=2 * node.node) + return BinaryNode(value=2 * node.value) double_tree = INFINITE_BINARY_TREE.transform(double, keep=MaxDepth(2)) values = [node.value for node, _ in double_tree.nodes.preorder()] diff --git a/tests/test_route.py b/tests/test_route.py index d18a174..91c2ded 100644 --- a/tests/test_route.py +++ b/tests/test_route.py @@ -22,19 +22,19 @@ def setUp(self): self.heap_route = Route(HEAP_TREE.children[0], HEAP_TREE.children[1]) def test_anchors(self): - result = [node.node for node in self.route1.anchors] + result = [node.value for node in self.route1.anchors] self.assertEqual([0, 18, 0], result) self.assertEqual(3, len(self.route1.anchors)) - result = [node.node for node in self.route2.anchors] + result = [node.value for node in self.route2.anchors] self.assertEqual([18, 2, 18, 2], result) self.assertEqual(4, len(self.route2.anchors)) - result = [node.node for node in self.route3.anchors] + result = [node.value for node in self.route3.anchors] self.assertEqual([2, 2, 0, 0], result) self.assertEqual(4, len(self.route3.anchors)) - result = [node.node for node in self.route4.anchors] + result = [node.value for node in self.route4.anchors] self.assertEqual([7, 18, 8], result) self.assertEqual(3, len(self.route4.anchors)) @@ -43,19 +43,19 @@ def test_anchors(self): self.assertEqual(3, len(self.route4.anchors)) def test_nodes(self): - result = [node.node for node in self.route1.nodes] + result = [node.value for node in self.route1.nodes] self.assertEqual([0, 1, 3, 8, 18, 8, 3, 1, 0], result) self.assertEqual(9, len(self.route1.nodes)) - result = [node.node for node in self.route2.nodes] + result = [node.value for node in self.route2.nodes] self.assertEqual([18, 8, 3, 1, 0, 2, 0, 1, 3, 8, 18, 8, 3, 1, 0, 2], result) self.assertEqual(16, len(self.route2.nodes)) - result = [node.node for node in self.route3.nodes] + result = [node.value for node in self.route3.nodes] self.assertEqual([2, 0], result) self.assertEqual(2, len(self.route3.nodes)) - result = [node.node for node in self.route4.nodes] + result = [node.value for node in self.route4.nodes] self.assertEqual([7, 3, 8, 18, 8], result) self.assertEqual(5, len(self.route4.nodes)) @@ -64,14 +64,14 @@ def test_nodes(self): self.assertEqual(3, len(self.heap_route.nodes)) def test_edges(self): - result = [(v1.node, v2.node) for (v1, v2) in self.route1.edges] + result = [(v1.value, v2.value) for (v1, v2) in self.route1.edges] expected = [(0, 1), (1, 3), (3, 8), (8, 18), (18, 8), (8, 3), (3, 1), (1, 0)] self.assertEqual(expected, result) self.assertEqual(8, len(self.route1.edges)) def test_reversed(self): - node_result = [node.node for node in reversed(self.route2.nodes)] - edge_result = [(v1.node, v2.node) for (v1, v2) in reversed(self.route2.edges)] + node_result = [node.value for node in reversed(self.route2.nodes)] + edge_result = [(v1.value, v2.value) for (v1, v2) in reversed(self.route2.edges)] node_expected = [2, 0, 1, 3, 8, 18, 8, 3, 1, 0, 2, 0, 1, 3, 8, 18] edge_expected = [ (2, 0), @@ -94,10 +94,10 @@ def test_reversed(self): self.assertEqual(edge_expected, edge_result) def test_lca(self): - self.assertEqual(0, self.route1.lca.node) - self.assertEqual(0, self.route2.lca.node) - self.assertEqual(0, self.route3.lca.node) - self.assertEqual(3, self.route4.lca.node) + self.assertEqual(0, self.route1.lca.value) + self.assertEqual(0, self.route2.lca.value) + self.assertEqual(0, self.route3.lca.value) + self.assertEqual(3, self.route4.lca.value) self.assertEqual(0, self.heap_route.lca.index) def test_add_anchor(self): @@ -108,6 +108,14 @@ def test_add_anchor(self): with self.assertRaises(ValueError): self.route1.add_anchor(SEQTREE) + def test_paths(self): + from pathlib import Path + p1 = Path(__file__) + p2 = Path(__file__).parent.parent / "src" + + res = [node.name for node in Route(p1, p2).nodes] + self.assertEqual(['test_route.py', 'tests', 'abstracttree', 'src'], res) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_tree.py b/tests/test_tree.py index b96b7be..860cefa 100644 --- a/tests/test_tree.py +++ b/tests/test_tree.py @@ -1,15 +1,18 @@ from unittest import TestCase -import tree_instances as trees +try: + from . import tree_instances as trees +except ImportError: + import tree_instances as trees class TestTree(TestCase): def test_siblings_iter(self): tree = trees.INFINITE_BINARY_TREE - values = [node.node for node in tree.siblings] - values_left = [node.node for node in tree.children[0].siblings] - values_right = [node.node for node in tree.children[1].siblings] - values_lr = [node.node for node in tree.children[0].children[1].siblings] + values = [node.value for node in tree.siblings] + values_left = [node.value for node in tree.children[0].siblings] + values_right = [node.value for node in tree.children[1].siblings] + values_lr = [node.value for node in tree.children[0].children[1].siblings] self.assertEqual([], values) self.assertEqual([2], values_left) self.assertEqual([1], values_right) diff --git a/tests/test_treelike.py b/tests/test_treelike.py new file mode 100644 index 0000000..a70541e --- /dev/null +++ b/tests/test_treelike.py @@ -0,0 +1,56 @@ +import ast +from unittest import TestCase + +from abstracttree import HeapTree, as_tree +from abstracttree.generics import TreeLike, DownTreeLike +from pathlib import Path, PurePath + +try: + ExceptionGroup +except NameError: + # Patch this test for lower python versions + class ExceptionGroup(): + def __init__(self, *args): ... + + @property + def children(self): return [] + + +class TreeLikeTest(TestCase): + def setUp(self): + self.treelike_instances = [ + Path("/HasParent/HasChildren"), + HeapTree(list(range(40)), 5), + as_tree(ast.parse("f(x, y) + b")), + ] + self.pure_downtreelike_instances = [ + [], + ast.parse("f(x, y) + b"), + ExceptionGroup("A lot of things went wrong today", + [ValueError("value"), IndexError("index")]), + ] + self.untreelike_instances = [ + PurePath("/HasParent/HasNoChildren"), + 3, + "hello", + False, + None, + ] + + def test_treelike(self): + for inst in self.treelike_instances: + with self.subTest(msg=str(inst)): + self.assertIsInstance(inst, TreeLike) + self.assertIsInstance(inst, DownTreeLike) + + def test_downtreelike(self): + for inst in self.pure_downtreelike_instances: + with self.subTest(msg=str(inst)): + self.assertIsInstance(inst, DownTreeLike) + self.assertNotIsInstance(inst, TreeLike) + + def test_untreelike(self): + for inst in self.untreelike_instances: + with self.subTest(msg=str(inst)): + self.assertNotIsInstance(inst, DownTreeLike) + self.assertNotIsInstance(inst, TreeLike) diff --git a/tests/test_uptree.py b/tests/test_uptree.py index 063d9ff..79bd932 100644 --- a/tests/test_uptree.py +++ b/tests/test_uptree.py @@ -1,27 +1,31 @@ from pathlib import Path from unittest import TestCase -import tree_instances as trees -from abstracttree import astree +try: + from . import tree_instances as trees +except ImportError: + import tree_instances as trees + +from abstracttree import as_tree class TestUpTree(TestCase): def test_parent(self): - path_parent = astree(Path("this/path/should/not")) + path_parent = as_tree(Path("this/path/should/not")) self.assertEqual(None, trees.SINGLETON.parent) - self.assertEqual(path_parent, trees.NONEXISTENTPATH.parent) + self.assertEqual(str(path_parent), str(trees.NONEXISTENTPATH.parent)) # Or should there be an eqv generic that does samefile? self.assertEqual(None, trees.INFINITE_BINARY_TREE.parent) - self.assertEqual(2, trees.INFINITE_BINARY_TREE_SUBTREE.parent.node) - self.assertEqual(4, trees.COUNTDOWN.parent.node) + self.assertEqual(2, trees.INFINITE_BINARY_TREE_SUBTREE.parent.value) + self.assertEqual(4, trees.COUNTDOWN.parent.value) self.assertEqual(trees.INFINITE_TREE, trees.INFINITE_TREE) def test_root(self): - path_root = astree(Path("")) + path_root = as_tree(Path("")) self.assertEqual(trees.SINGLETON, trees.SINGLETON.root) - self.assertEqual(path_root, trees.NONEXISTENTPATH.root) - self.assertEqual(0, trees.INFINITE_BINARY_TREE.root.node) - self.assertEqual(0, trees.INFINITE_BINARY_TREE_SUBTREE.root.node) - self.assertEqual(0, trees.COUNTDOWN.root.node) + self.assertEqual(str(path_root), str(trees.NONEXISTENTPATH.root)) + self.assertEqual(0, trees.INFINITE_BINARY_TREE.root.value) + self.assertEqual(0, trees.INFINITE_BINARY_TREE_SUBTREE.root.value) + self.assertEqual(0, trees.COUNTDOWN.root.value) def test_is_root(self): self.assertTrue(trees.SINGLETON.is_root) diff --git a/tests/tree_instances.py b/tests/tree_instances.py index b03bf70..814aae3 100644 --- a/tests/tree_instances.py +++ b/tests/tree_instances.py @@ -1,8 +1,7 @@ import heapq from pathlib import Path -from abstracttree import MutableDownTree, Tree, astree, HeapTree -from abstracttree.binarytree import BinaryDownTree +from abstracttree import BinaryDownTree, MutableDownTree, Tree, as_tree, HeapTree class BinaryNode(MutableDownTree, BinaryDownTree): @@ -52,8 +51,8 @@ def __str__(self): return "Infinite singleton!" -SINGLETON = astree("Singleton", children=lambda n: ()) -NONEXISTENTPATH = Tree.convert(Path("this/path/should/not/exist")) +SINGLETON = as_tree("Singleton", children=lambda n: ()) +NONEXISTENTPATH = as_tree(Path("this/path/should/not/exist")) BINARY_TREE = BinaryNode(1) # 2 children BINARY_TREE.left = BinaryNode(2) # leaf @@ -63,17 +62,17 @@ def __str__(self): BINARY_TREE_SUBTREE = BINARY_TREE.right.left.right -INFINITE_BINARY_TREE = astree(0, children=lambda n: (2 * n + 1, 2 * n + 2)) +INFINITE_BINARY_TREE = as_tree(0, children=lambda n: (2 * n + 1, 2 * n + 2)) INFINITE_BINARY_TREE_SUBTREE = INFINITE_BINARY_TREE.children[1].children[0] COUNTDOWN_MAX = 5 -COUNTDOWN = astree( +COUNTDOWN = as_tree( COUNTDOWN_MAX, children=lambda n: [n + 1] if n < COUNTDOWN_MAX else (), parent=lambda n: n - 1 if n > 0 else None, ) -SEQTREE = Tree.convert([1, [2, 3], []]) +SEQTREE = as_tree([1, [2, 3], []]) INFINITE_TREE = InfiniteSingleton()