Enhance dagnode

CHANGELOG.md

@@ -5,6 +5,14 @@ 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).
 
 ## [Unreleased]
+### Added
+- DAGNode: Able to access and delete node children via name with square bracket accessor with `__getitem__` and `__delitem__` magic methods.
+- DAGNode: Able to delete all children for a node.
+- DAGNode: Able to check if node contains child node with `__contains__` magic method.
+- DAGNode: Able to iterate the node to access children with `__iter__` magic method.
+### Changed
+- Tree Search: Modify type hints to include DAGNode for `find_children`, `find_child`, and `find_child_by_name`.
+- Misc: Neater handling of strings for tests.
 
 ## [0.15.5] - 2023-01-17
 ### Changed

README.md

@@ -555,7 +555,7 @@ Below are the tables of attributes available to `BaseNode` and `Node` classes.
 | Get parent                        | `node_e.parent`            | Node(/a/b, )                                                                         |
 | Get siblings                      | `node_e.siblings`          | (Node(/a/b/d, ), Node(/a/b/f, ))                                                     |
 | Get left sibling                  | `node_e.left_sibling`      | Node(/a/b/d, )                                                                       |
-| Get right sibling                 | `node_e.left_sibling`      | Node(/a/b/f, )                                                                       |
+| Get right sibling                 | `node_e.right_sibling`     | Node(/a/b/f, )                                                                       |
 | Get ancestors (lazy evaluation)   | `list(node_e.ancestors)`   | [Node(/a/b, ), Node(/a, )]                                                           |
 | Get descendants (lazy evaluation) | `list(node_b.descendants)` | [Node(/a/b/d, ), Node(/a/b/e, ), Node(/a/b/f, ), Node(/a/b/f/h, ), Node(/a/b/f/i, )] |
 | Get leaves (lazy evaluation)      | `list(node_b.leaves)`      | [Node(/a/b/d, ), Node(/a/b/e, ), Node(/a/b/f/h, ), Node(/a/b/f/i, )]                 |
@@ -1151,8 +1151,10 @@ Compared to nodes in tree, nodes in DAG are able to have multiple parents.
 
 1. **From `DAGNode`**
 
-DAGNode can be linked to each other with `parents` and `children` setter methods,
-or using bitshift operator with the convention `parent_node >> child_node` or `child_node << parent_node`.
+DAGNodes can be linked to each other in the following ways:
+  - Using `parents` and `children` setter methods
+  - Directly passing `parents` or `children` argument
+  - Using bitshift operator with the convention `parent_node >> child_node` or `child_node << parent_node`
 
 {emphasize-lines="5-8,10"}
 ```python
@@ -1238,6 +1240,55 @@ print([(parent.node_name, child.node_name) for parent, child in dag_iterator(dag
 # [('a', 'd'), ('c', 'd'), ('d', 'e'), ('a', 'c'), ('b', 'c')]
 ```
 
+### DAG Attributes and Operations
+
+Note that using `DAGNode` as superclass inherits the default class attributes (properties) and operations (methods).
+
+```python
+from bigtree import list_to_dag
+
+relations_list = [
+   ("a", "c"),
+   ("a", "d"),
+   ("b", "c"),
+   ("c", "d"),
+   ("d", "e")
+]
+dag = list_to_dag(relations_list)
+dag
+# DAGNode(d, )
+
+# Accessing children
+node_e = dag["e"]
+node_a = dag.parents[0]
+```
+
+Below are the tables of attributes available to `DAGNode` class.
+
+| Attributes wrt self                  | Code             | Returns |
+|--------------------------------------|------------------|---------|
+| Check if root                        | `node_a.is_root` | True    |
+| Check if leaf node                   | `dag.is_leaf`    | False   |
+| Get node name (only for `Node`)      | `dag.node_name`  | 'd'     |
+
+| Attributes wrt structure     | Code                  | Returns                                                              |
+|------------------------------|-----------------------|----------------------------------------------------------------------|
+| Get child/children           | `node_a.children`     | (DAGNode(c, ), DAGNode(d, ))                                         |
+| Get parents                  | `dag.parents`         | (DAGNode(a, ), DAGNode(c, ))                                         |
+| Get siblings                 | `dag.siblings`        | (DAGNode(c, ),)                                                      |
+| Get ancestors                | `dag.ancestors`       | [DAGNode(a, ), DAGNode(b, ), DAGNode(c, )]                           |
+| Get descendants              | `dag.descendants`     | [DAGNode(e, )]                                                       |
+
+Below is the table of operations available to `DAGNode` class.
+
+| Operations                            | Code                                                       | Returns                                                                                                          |
+|---------------------------------------|------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------|
+| Get node information                  | `dag.describe(exclude_prefix="_")`                         | [('name', 'd')]                                                                                                  |
+| Find path(s) from one node to another | `node_a.go_to(dag)`                                        | [[DAGNode(a, ), DAGNode(c, ), DAGNode(d, description=dag-tag)], [DAGNode(a, ), DAGNode(d, description=dag-tag)]] |
+| Set attribute(s)                      | `dag.set_attrs({"description": "dag-tag"})`                | None                                                                                                             |
+| Get attribute                         | `dag.get_attr("description")`                              | 'dag-tag'                                                                                                        |
+| Copy DAG                              | `dag.copy()`                                               | None                                                                                                             |
+
 ----
 
 ## Demo Usage

bigtree/node/dagnode.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import copy
-from typing import Any, Dict, Iterable, List, Optional, Tuple, TypeVar
+from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple, TypeVar
 
 from bigtree.utils.exceptions import LoopError, TreeError
 from bigtree.utils.iterators import preorder_iter
@@ -319,6 +319,13 @@ def children(self: T, new_children: Iterable[T]) -> None:
                     self.__children.remove(new_child)
             raise TreeError(exc_info)
 
+    @children.deleter
+    def children(self) -> None:
+        """Delete child node(s)"""
+        for child in self.children:
+            self.__children.remove(child)  # type: ignore
+            child.__parents.remove(self)  # type: ignore
+
     def __pre_assign_children(self: T, new_children: Iterable[T]) -> None:
         """Custom method to check before attaching children
         Can be overridden with `_DAGNode__pre_assign_children()`
@@ -553,6 +560,32 @@ def __copy__(self: T) -> T:
         obj.__dict__.update(self.__dict__)
         return obj
 
+    def __getitem__(self, child_name: str) -> T:
+        """Get child by name identifier
+
+        Args:
+            child_name (str): name of child node
+
+        Returns:
+            (Self): child node
+        """
+        from bigtree.tree.search import find_child_by_name
+
+        return find_child_by_name(self, child_name)  # type: ignore
+
+    def __delitem__(self, child_name: str) -> None:
+        """Delete child by name identifier, will not throw error if child does not exist
+
+        Args:
+            child_name (str): name of child node
+        """
+        from bigtree.tree.search import find_child_by_name
+
+        child = find_child_by_name(self, child_name)
+        if child:
+            self.__children.remove(child)  # type: ignore
+            child.__parents.remove(self)  # type: ignore
+
     def __repr__(self) -> str:
         """Print format of DAGNode
 
@@ -567,20 +600,39 @@ def __repr__(self) -> str:
         return f"{class_name}({self.node_name}, {node_description})"
 
     def __rshift__(self: T, other: T) -> None:
-        """Set children using >> bitshift operator for self >> other
+        """Set children using >> bitshift operator for self >> children (other)
 
         Args:
             other (Self): other node, children
         """
         other.parents = [self]
 
     def __lshift__(self: T, other: T) -> None:
-        """Set parent using << bitshift operator for self << other
+        """Set parent using << bitshift operator for self << parent (other)
 
         Args:
             other (Self): other node, parent
         """
         self.parents = [other]
 
+    def __iter__(self) -> Generator[T, None, None]:
+        """Iterate through child nodes
+
+        Returns:
+            (Self): child node
+        """
+        yield from self.children  # type: ignore
+
+    def __contains__(self, other_node: T) -> bool:
+        """Check if child node exists
+
+        Args:
+            other_node (T): child node
+
+        Returns:
+            (bool)
+        """
+        return other_node in self.children
+
 
 T = TypeVar("T", bound=DAGNode)

bigtree/tree/search.py

@@ -1,6 +1,7 @@
-from typing import Any, Callable, Iterable, List, Tuple, TypeVar
+from typing import Any, Callable, Iterable, List, Tuple, TypeVar, Union
 
 from bigtree.node.basenode import BaseNode
+from bigtree.node.dagnode import DAGNode
 from bigtree.node.node import Node
 from bigtree.utils.exceptions import SearchError
 from bigtree.utils.iterators import preorder_iter
@@ -24,6 +25,7 @@
 
 T = TypeVar("T", bound=BaseNode)
 NodeT = TypeVar("NodeT", bound=Node)
+DAGNodeT = TypeVar("DAGNodeT", bound=DAGNode)
 
 
 def findall(
@@ -368,11 +370,11 @@ def find_attrs(
 
 
 def find_children(
-    tree: T,
-    condition: Callable[[T], bool],
+    tree: Union[T, DAGNodeT],
+    condition: Callable[[Union[T, DAGNodeT]], bool],
     min_count: int = 0,
     max_count: int = 0,
-) -> Tuple[T, ...]:
+) -> Tuple[Union[T, DAGNodeT], ...]:
     """
     Search children for nodes matching condition (callable function).
 
@@ -385,15 +387,15 @@ def find_children(
     (Node(/a/b, age=65), Node(/a/c, age=60))
 
     Args:
-        tree (BaseNode): tree to search for its children
+        tree (BaseNode/DAGNode): tree to search for its children
         condition (Callable): function that takes in node as argument, returns node if condition evaluates to `True`
         min_count (int): checks for minimum number of occurrences,
             raise SearchError if the number of results do not meet min_count, defaults to None
         max_count (int): checks for maximum number of occurrences,
             raise SearchError if the number of results do not meet min_count, defaults to None
 
     Returns:
-        (BaseNode)
+        (BaseNode/DAGNode)
     """
     result = tuple([node for node in tree.children if node and condition(node)])
     if min_count and len(result) < min_count:
@@ -408,9 +410,9 @@ def find_children(
 
 
 def find_child(
-    tree: T,
-    condition: Callable[[T], bool],
-) -> T:
+    tree: Union[T, DAGNodeT],
+    condition: Callable[[Union[T, DAGNodeT]], bool],
+) -> Union[T, DAGNodeT]:
     """
     Search children for *single node* matching condition (callable function).
 
@@ -423,18 +425,20 @@ def find_child(
     Node(/a/b, age=65)
 
     Args:
-        tree (BaseNode): tree to search for its child
+        tree (BaseNode/DAGNode): tree to search for its child
         condition (Callable): function that takes in node as argument, returns node if condition evaluates to `True`
 
     Returns:
-        (BaseNode)
+        (BaseNode/DAGNode)
     """
     result = find_children(tree, condition, max_count=1)
     if result:
         return result[0]
 
 
-def find_child_by_name(tree: NodeT, name: str) -> NodeT:
+def find_child_by_name(
+    tree: Union[NodeT, DAGNodeT], name: str
+) -> Union[NodeT, DAGNodeT]:
     """
     Search tree for single node matching name attribute.
 
@@ -449,10 +453,10 @@ def find_child_by_name(tree: NodeT, name: str) -> NodeT:
     Node(/a/c/d, age=40)
 
     Args:
-        tree (Node): tree to search, parent node
+        tree (Node/DAGNode): tree to search, parent node
         name (str): value to match for name attribute, child node
 
     Returns:
-        (Node)
+        (Node/DAGNode)
     """
     return find_child(tree, lambda node: node.node_name == name)

tests/binarytree/test_helper.py

@@ -9,7 +9,16 @@ class TestCloneTree:
     def test_clone_tree(binarytree_node):
         root_clone = clone_tree(binarytree_node, node_type=Node)
         assert isinstance(root_clone, Node), "Wrong type returned"
-        expected_str = """1\n├── 2\n│   ├── 4\n│   │   └── 8\n│   └── 5\n└── 3\n    ├── 6\n    └── 7\n"""
+        expected_str = (
+            "1\n"
+            "├── 2\n"
+            "│   ├── 4\n"
+            "│   │   └── 8\n"
+            "│   └── 5\n"
+            "└── 3\n"
+            "    ├── 6\n"
+            "    └── 7\n"
+        )
         assert_print_statement(print_tree, expected_str, tree=binarytree_node)
 
 
@@ -87,6 +96,10 @@ def test_tree_diff(binarytree_node):
         other_tree_node = prune_tree(binarytree_node, "1/3")
         tree_only_diff = get_tree_diff(binarytree_node, other_tree_node, only_diff=True)
         expected_str = (
-            """1\n└── 2 (-)\n    ├── 4 (-)\n    │   └── 8 (-)\n    └── 5 (-)\n"""
+            "1\n"
+            "└── 2 (-)\n"
+            "    ├── 4 (-)\n"
+            "    │   └── 8 (-)\n"
+            "    └── 5 (-)\n"
         )
         assert_print_statement(print_tree, expected_str, tree=tree_only_diff)

tests/node/test_dagnode.py

@@ -951,6 +951,56 @@ def assert_dag_structure_self(self):
             actual == expected
         ), f"Node attribute should be {expected}, but it is {actual}"
 
+    # Test accessing with square bracket (__getitem__)
+    assert (
+        self.a["c"] == self.c
+    ), f"Accessor method not returning correct for node {self.a} for {self.c}"
+    assert (
+        self.a["d"] == self.d
+    ), f"Accessor method not returning correct for node {self.a} for {self.d}"
+    assert (
+        self.a["c"]["d"] == self.d
+    ), f"Accessor method not returning correct for node {self.a} for {self.d}"
+
+    # Test deletion of children with square bracket (__delitem__)
+    assert len(self.a.children) == 2, f"Before deletion: error in {self.a.children}"
+    del self.a["d"]
+    assert len(self.a.children) == 1, f"After deletion: error in {self.a.children}"
+    self.a >> self.d
+    assert (
+        len(self.a.children) == 2
+    ), f"Revert after deletion: error in {self.a.children}"
+
+    # Test deletion of non-existent children with square bracket (__delitem__)
+    assert len(self.a.children) == 2, f"Before deletion: error in {self.a.children}"
+    del self.a["b"]
+    assert len(self.a.children) == 2, f"After deletion: error in {self.a.children}"
+
+    # Test deletion of all children
+    assert len(self.a.children) == 2, f"Before deletion: error in {self.a.children}"
+    assert self.a in self.c.parents, f"Before deletion: error in {self.c.parents}"
+    assert self.a in self.d.parents, f"Before deletion: error in {self.d.parents}"
+    del self.a.children
+    assert len(self.a.children) == 0, f"After deletion: error in {self.a.children}"
+    assert self.a not in self.c.parents, f"Before deletion: error in {self.c.parents}"
+    assert self.a not in self.d.parents, f"Before deletion: error in {self.d.parents}"
+    self.a >> self.c
+    self.a >> self.d
+    assert (
+        len(self.a.children) == 2
+    ), f"Revert after deletion: error in {self.a.children}"
+
+    # Test iteration (__iter__)
+    expected = [self.c, self.d]
+    actual = [child for child in self.a]
+    assert (
+        actual == expected
+    ), f"Node {self.a} should have {expected} children when iterated, but it has {actual}"
+
+    # Test contains (__contains__)
+    assert self.c in self.a, f"Check if {self.a} contains {self.c}, expected True"
+    assert self.b not in self.a, f"Check if {self.a} contains {self.b}, expected False"
+
 
 def assert_dag_child_attr(dag, parent_name, child_name, child_attr, child_value):
     """Test tree attributes"""

tests/node/test_node.py

@@ -422,7 +422,7 @@ def assert_tree_structure_node_self(self):
             actual == expected
         ), f"Node should have path {expected}, but path is {actual}"
 
-    # Test show
+    # Test show()
     expected_str = (
         "a\n"
         "├── b\n"
@@ -454,6 +454,7 @@ def assert_tree_structure_node_self(self):
         style="ansi",
     )
 
+    # Test hshow()
     expected_str = (
         "           ┌─ d\n"
         "     ┌─ b ─┤     ┌─ g\n"