Merge pull request #81 from kayjan/v0.11.0

V0.11.0

CHANGELOG.md

@@ -4,6 +4,15 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.11.0] - 2023-09-08
+### Added
+- Tree Helper: Pruning tree to allow pruning by `prune_path` and `max_depth`.
+- Tree Plot: Implement Enhanced Reingold Tilford Algorithm to retrieve (x, y) coordinates for a tree structure.
+### Changed
+- BaseNode/DAGNode: `get_attr` method to allow default return value.
+### Fixed
+- Utility Iterator: Relax type hinting using TypeVar.
+
 ## [0.10.3] - 2023-08-12
 ### Added
 - Tree Constructor: `add_path_to_tree`, `dataframe_to_tree`, `dataframe_to_tree_by_relation` to allow custom node types that takes in constructor arguments.
@@ -157,10 +166,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [0.6.3] - 2022-11-15
 ### Added
 - DAGNode: Rollback functionality to original state when there is error setting parent and children (backwards-compatible).
-
 ### Changed
 - BaseNode, BNode, DAGNode: Refactor by abstracting checks.
-
 ### Fixed
 - BaseNode: Fix rollback logic to handle failure in pre-assign checks and reassigning same child / parent.
 - BNode: Fix issue of reassigning children shifting existing child from right to left.
@@ -181,15 +188,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - BTree Constructor: From list.
 - BNode Iterator: Level-Order Iterator.
 - Misc: Add Tips and Tricks to documentation (List Directory).
-
 ### Fixed
 - DAGNode: Fix issue of duplicate parent constructor creating duplicate children.
 
 ## [0.5.5] - 2022-11-12
 ### Added
 - Misc: More docstring examples.
 - Misc: More test cases.
-
 ### Fixed
 - Tree Modifier: Fix issue with `merge_children` argument not working as expected.
 
@@ -199,15 +204,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Node: Made more extendable with pre-/post-assign checks.
 - Misc: Add Tips and Tricks to documentation (Extending Nodes).
 - Misc: More test cases.
-
 ### Fixed
 - Tree Search: Type hints.
 
 ## [0.5.3] - 2022-11-11
 ### Added
 - DAG and Tree Exporter: More customizations allowed on edges.
 - Add Tips and Tricks to documentation (Weighted Trees, Merging Trees).
-
 ### Fixed
 - Tree Modifier: Fix issue with `merge_children` argument not working as expected.
 
@@ -222,15 +225,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [0.5.0] - 2022-11-09
 ### Added
 - Misc: Clean codes and documentation.
-
 ### Changed
 - Tree Exporter: Printing tree to group multiple arguments together.
 - DAG and Tree Exporter: Export to dot able to plot multiple disjointed trees/dags, rename `bgcolor` to `bg_colour`.
 
 ## [0.4.6] - 2022-11-09
 ### Added
 - Tree Constructor: From DataFrame of parent-child columns.
-
 ### Changed
 - Tree Exporter: Printing tree to define node name or path, and default to const style.
 - Tree Constructor: Rename `list_to_tree_tuples` to `list_to_tree_by_relation`.
@@ -263,7 +264,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [0.3.3] - 2022-11-07
 ### Added
 - DAG Exporter: To list, nested dictionary, pandas DataFrame.
-
 ### Changed
 - BaseNode and DAGNode: Modify docstring.
 - Tree Exporter: Support Nodes with same name.
@@ -303,6 +303,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Utility Iterator: Tree traversal methods.
 - Workflow To Do App: Tree use case with to-do list implementation.
 
+[0.11.0]: https://github.com/kayjan/bigtree/compare/0.10.3...0.11.0
 [0.10.3]: https://github.com/kayjan/bigtree/compare/0.10.2...0.10.3
 [0.10.2]: https://github.com/kayjan/bigtree/compare/0.10.1...0.10.2
 [0.10.1]: https://github.com/kayjan/bigtree/compare/0.10.0...0.10.1

README.md

@@ -24,7 +24,7 @@ Related Links:
 There are 3 segments to Big Tree consisting of Tree, Binary Tree, and Directed Acyclic Graph (DAG) implementation.
 
 
-For **Tree** implementation, there are 8 main components.
+For **Tree** implementation, there are 9 main components.
 
 1. [**Node**](https://bigtree.readthedocs.io/en/latest/node.html)
    1. ``BaseNode``, extendable class
@@ -57,12 +57,14 @@ For **Tree** implementation, there are 8 main components.
    1. Cloning tree to another `Node` type
    2. Prune tree
    3. Get difference between two trees
-7. [**Exporting Tree**](https://bigtree.readthedocs.io/en/latest/bigtree/tree/export.html)
+7. [**Plotting Tree**](https://bigtree.readthedocs.io/en/latest/bigtree/utils/plot.html)
+   1. Enhanced Reingold Tilford Algorithm to retrieve (x, y) coordinates for a tree structure
+8. [**Exporting Tree**](https://bigtree.readthedocs.io/en/latest/bigtree/tree/export.html)
    1. Print to console
    2. Export to *dictionary*, *nested dictionary*, or *pandas DataFrame*
    3. Export tree to dot (can save to .dot, .png, .svg, .jpeg files)
    4. Export tree to Pillow (can save to .png, .jpg)
-8. [**Workflows**](https://bigtree.readthedocs.io/en/latest/workflows.html)
+9. [**Workflows**](https://bigtree.readthedocs.io/en/latest/workflows.html)
    1. Sample workflows for tree demonstration!
 
 For **Binary Tree** implementation, there are 3 main components.

bigtree/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.10.3"
+__version__ = "0.11.0"
 
 from bigtree.binarytree.construct import list_to_binarytree
 from bigtree.dag.construct import dataframe_to_dag, dict_to_dag, list_to_dag
@@ -70,5 +70,6 @@
     zigzag_iter,
     zigzaggroup_iter,
 )
+from bigtree.utils.plot import reingold_tilford
 from bigtree.workflows.app_calendar import Calendar
 from bigtree.workflows.app_todo import AppToDo

bigtree/node/basenode.py

@@ -547,7 +547,7 @@ def describe(
             and (not len(exclude_prefix) or not item[0].startswith(exclude_prefix))
         ]
 
-    def get_attr(self, attr_name: str) -> Any:
+    def get_attr(self, attr_name: str, default_value: Any = None) -> Any:
         """Get value of node attribute
         Returns None if attribute name does not exist
 
@@ -558,14 +558,15 @@ def get_attr(self, attr_name: str) -> Any:
 
         Args:
             attr_name (str): attribute name
+            default_value (Any): default value if attribute does not exist, defaults to None
 
         Returns:
             (Any)
         """
         try:
             return getattr(self, attr_name)
         except AttributeError:
-            return None
+            return default_value
 
     def set_attrs(self, attrs: Dict[str, Any]) -> None:
         """Set node attributes

bigtree/node/dagnode.py

@@ -443,20 +443,21 @@ def describe(
             and (not len(exclude_prefix) or not item[0].startswith(exclude_prefix))
         ]
 
-    def get_attr(self, attr_name: str) -> Any:
+    def get_attr(self, attr_name: str, default_value: Any = None) -> Any:
         """Get value of node attribute
         Returns None if attribute name does not exist
 
         Args:
             attr_name (str): attribute name
+            default_value (Any): default value if attribute does not exist, defaults to None
 
         Returns:
             (Any)
         """
         try:
             return getattr(self, attr_name)
         except AttributeError:
-            return None
+            return default_value
 
     def set_attrs(self, attrs: Dict[str, Any]) -> None:
         """Set node attributes

bigtree/tree/helper.py

@@ -7,6 +7,7 @@
 from bigtree.tree.export import tree_to_dataframe
 from bigtree.tree.search import find_path
 from bigtree.utils.exceptions import NotFoundError
+from bigtree.utils.iterators import levelordergroup_iter
 
 __all__ = ["clone_tree", "prune_tree", "get_tree_diff"]
 
@@ -50,55 +51,83 @@ def recursive_add_child(_new_parent_node: BaseNode, _parent_node: BaseNode) -> N
 
 
 def prune_tree(
-    tree: Union[BinaryNode, Node], prune_path: str, sep: str = "/"
+    tree: Union[BinaryNode, Node],
+    prune_path: str = "",
+    sep: str = "/",
+    max_depth: int = 0,
 ) -> Union[BinaryNode, Node]:
-    """Prune tree to leave only the prune path, returns the root of a *copy* of the original tree.
+    """Prune tree by path or depth, returns the root of a *copy* of the original tree.
 
-    All siblings along the prune path will be removed.
-    Prune path name should be unique, can be full path or partial path (trailing part of path) or node name.
+    For pruning by `prune_path`,
+      All siblings along the prune path will be removed.
+      Prune path name should be unique, can be full path or partial path (trailing part of path) or node name.
+
+    For pruning by `max_depth`,
+      All nodes that are beyond `max_depth` will be removed.
 
     Path should contain `Node` name, separated by `sep`.
       - For example: Path string "a/b" refers to Node("b") with parent Node("a").
 
     >>> from bigtree import Node, prune_tree
     >>> root = Node("a")
     >>> b = Node("b", parent=root)
-    >>> c = Node("c", parent=root)
+    >>> c = Node("c", parent=b)
+    >>> d = Node("d", parent=b)
+    >>> e = Node("e", parent=root)
     >>> root.show()
     a
     ├── b
-    └── c
+    │   ├── c
+    │   └── d
+    └── e
 
     >>> root_pruned = prune_tree(root, "a/b")
     >>> root_pruned.show()
     a
     └── b
+        ├── c
+        └── d
+
+    >>> root_pruned = prune_tree(root, max_depth=2)
+    >>> root_pruned.show()
+    a
+    ├── b
+    └── e
 
     Args:
         tree (Union[BinaryNode, Node]): existing tree
         prune_path (str): prune path, all siblings along the prune path will be removed
-        sep (str): path separator
+        sep (str): path separator of `prune_path`
+        max_depth (int): maximum depth of pruned tree, based on `depth` attribute, defaults to None
 
     Returns:
         (Union[BinaryNode, Node])
     """
-    prune_path = prune_path.replace(sep, tree.sep)
+    if not prune_path and not max_depth:
+        raise ValueError("Please specify either `prune_path` or `max_depth` or both.")
+
     tree_copy = tree.copy()
-    child = find_path(tree_copy, prune_path)
-    if not child:
-        raise NotFoundError(
-            f"Cannot find any node matching path_name ending with {prune_path}"
-        )
 
-    if isinstance(child.parent, BinaryNode):
+    # Prune by path (prune bottom-up)
+    if prune_path:
+        prune_path = prune_path.replace(sep, tree.sep)
+        child = find_path(tree_copy, prune_path)
+        if not child:
+            raise NotFoundError(
+                f"Cannot find any node matching path_name ending with {prune_path}"
+            )
         while child.parent:
-            child.parent.children = [child, None]  # type: ignore
+            for other_children in child.parent.children:
+                if other_children != child:
+                    other_children.parent = None
             child = child.parent
-        return tree_copy
 
-    while child.parent:
-        child.parent.children = [child]  # type: ignore
-        child = child.parent
+    # Prune by depth (prune top-down)
+    if max_depth:
+        for depth, level_nodes in enumerate(levelordergroup_iter(tree_copy), 1):
+            if depth == max_depth:
+                for level_node in level_nodes:
+                    del level_node.children
     return tree_copy