from __future__ import annotations import copy from typing import Any, Dict, Generator, Iterable, List, Optional, Set, Tuple, TypeVar from bigtree.globals import ASSERTIONS from bigtree.utils.exceptions import CorruptedTreeError, LoopError, TreeError from bigtree.utils.iterators import preorder_iter class BaseNode: """ BaseNode extends any Python class to a tree node. Nodes can have attributes if they are initialized from `Node`, *dictionary*, or *pandas DataFrame*. Nodes can be linked to each other with `parent` and `children` setter methods, or using bitshift operator with the convention `parent_node >> child_node` or `child_node << parent_node`. Examples: >>> from bigtree import Node, print_tree >>> root = Node("a", age=90) >>> b = Node("b", age=65) >>> c = Node("c", age=60) >>> d = Node("d", age=40) >>> root.children = [b, c] >>> d.parent = b >>> print_tree(root, attr_list=["age"]) a [age=90] ├── b [age=65] │ └── d [age=40] └── c [age=60] >>> from bigtree import Node >>> root = Node("a", age=90) >>> b = Node("b", age=65) >>> c = Node("c", age=60) >>> d = Node("d", age=40) >>> root >> b >>> root >> c >>> d << b >>> print_tree(root, attr_list=["age"]) a [age=90] ├── b [age=65] │ └── d [age=40] └── c [age=60] Directly passing `parent` argument. >>> from bigtree import Node >>> root = Node("a") >>> b = Node("b", parent=root) >>> c = Node("c", parent=root) >>> d = Node("d", parent=b) Directly passing `children` argument. >>> from bigtree import Node >>> d = Node("d") >>> c = Node("c") >>> b = Node("b", children=[d]) >>> a = Node("a", children=[b, c]) **BaseNode Creation** Node can be created by instantiating a `BaseNode` class or by using a *dictionary*. If node is created with dictionary, all keys of dictionary will be stored as class attributes. >>> from bigtree import Node >>> root = Node.from_dict({"name": "a", "age": 90}) **BaseNode Attributes** These are node attributes that have getter and/or setter methods. Get and set other `BaseNode` 1. ``parent``: Get/set parent node 2. ``children``: Get/set child nodes Get other `BaseNode` 1. ``ancestors``: Get ancestors of node excluding self, iterator 2. ``descendants``: Get descendants of node excluding self, iterator 3. ``leaves``: Get all leaf node(s) from self, iterator 4. ``siblings``: Get siblings of self 5. ``left_sibling``: Get sibling left of self 6. ``right_sibling``: Get sibling right of self Get `BaseNode` configuration 1. ``node_path``: Get tuple of nodes from root 2. ``is_root``: Get indicator if self is root node 3. ``is_leaf``: Get indicator if self is leaf node 4. ``root``: Get root node of tree 5. ``depth``: Get depth of self 6. ``max_depth``: Get maximum depth from root to leaf node **BaseNode Methods** These are methods available to be performed on `BaseNode`. Constructor methods 1. ``from_dict()``: Create BaseNode from dictionary `BaseNode` methods 1. ``describe()``: Get node information sorted by attributes, return list of tuples 2. ``get_attr(attr_name: str)``: Get value of node attribute 3. ``set_attrs(attrs: dict)``: Set node attribute name(s) and value(s) 4. ``go_to(node: Self)``: Get a path from own node to another node from same tree 5. ``append(node: Self)``: Add child to node 6. ``extend(nodes: List[Self])``: Add multiple children to node 7. ``copy()``: Deep copy self 8. ``sort()``: Sort child nodes ---- """ def __init__( self, parent: Optional[T] = None, children: Optional[List[T]] = None, **kwargs: Any, ): self.__parent: Optional[T] = None self.__children: List[T] = [] if children is None: children = [] self.parent = parent self.children = children # type: ignore if "parents" in kwargs: raise AttributeError( "Attempting to set `parents` attribute, do you mean `parent`?" ) self.__dict__.update(**kwargs) @staticmethod def __check_parent_type(new_parent: T) -> None: """Check parent type Args: new_parent (Self): parent node """ if not (isinstance(new_parent, BaseNode) or new_parent is None): raise TypeError( f"Expect parent to be BaseNode type or NoneType, received input type {type(new_parent)}" ) def __check_parent_loop(self, new_parent: T) -> None: """Check parent type Args: new_parent (Self): parent node """ if new_parent is not None: if new_parent is self: raise LoopError("Error setting parent: Node cannot be parent of itself") if any( ancestor is self for ancestor in new_parent.ancestors if new_parent.ancestors ): raise LoopError( "Error setting parent: Node cannot be ancestor of itself" ) @property def parent(self: T) -> Optional[T]: """Get parent node Returns: (Optional[Self]) """ return self.__parent @parent.setter def parent(self: T, new_parent: T) -> None: """Set parent node Args: new_parent (Self): parent node """ if ASSERTIONS: self.__check_parent_type(new_parent) self.__check_parent_loop(new_parent) current_parent = self.parent current_child_idx = None # Assign new parent - rollback if error self.__pre_assign_parent(new_parent) try: # Remove self from old parent if current_parent is not None: if not any( child is self for child in current_parent.children ): # pragma: no cover raise CorruptedTreeError( "Error setting parent: Node does not exist as children of its parent" ) current_child_idx = current_parent.__children.index(self) current_parent.__children.remove(self) # Assign self to new parent self.__parent = new_parent if new_parent is not None: new_parent.__children.append(self) self.__post_assign_parent(new_parent) except Exception as exc_info: # Remove self from new parent if new_parent is not None: new_parent.__children.remove(self) # Reassign self to old parent self.__parent = current_parent if current_child_idx is not None: current_parent.__children.insert(current_child_idx, self) raise TreeError(exc_info) def __pre_assign_parent(self, new_parent: T) -> None: """Custom method to check before attaching parent Can be overridden with `_BaseNode__pre_assign_parent()` Args: new_parent (Self): new parent to be added """ pass def __post_assign_parent(self, new_parent: T) -> None: """Custom method to check after attaching parent Can be overridden with `_BaseNode__post_assign_parent()` Args: new_parent (Self): new parent to be added """ pass @property def parents(self) -> None: """Do not allow `parents` attribute to be accessed Raises: AttributeError: No such attribute """ raise AttributeError( "Attempting to access `parents` attribute, do you mean `parent`?" ) @parents.setter def parents(self, new_parent: T) -> None: """Do not allow `parents` attribute to be set Args: new_parent (Self): parent node Raises: AttributeError: No such attribute """ raise AttributeError( "Attempting to set `parents` attribute, do you mean `parent`?" ) def __check_children_type( self: T, new_children: List[T] | Tuple[T] | Set[T] ) -> None: """Check child type Args: new_children (Iterable[Self]): child node """ if ( not isinstance(new_children, list) and not isinstance(new_children, tuple) and not isinstance(new_children, set) ): raise TypeError( f"Expect children to be List or Tuple or Set type, received input type {type(new_children)}" ) def __check_children_loop(self: T, new_children: Iterable[T]) -> None: """Check child loop Args: new_children (Iterable[Self]): child node """ seen_children = [] for new_child in new_children: # Check type if not isinstance(new_child, BaseNode): raise TypeError( f"Expect children to be BaseNode type, received input type {type(new_child)}" ) # Check for loop and tree structure if new_child is self: raise LoopError("Error setting child: Node cannot be child of itself") if any(child is new_child for child in self.ancestors): raise LoopError( "Error setting child: Node cannot be ancestor of itself" ) # Check for duplicate children if id(new_child) in seen_children: raise TreeError( "Error setting child: Node cannot be added multiple times as a child" ) else: seen_children.append(id(new_child)) @property def children(self: T) -> Tuple[T, ...]: """Get child nodes Returns: (Tuple[Self, ...]) """ return tuple(self.__children) @children.setter def children(self: T, new_children: List[T] | Tuple[T] | Set[T]) -> None: """Set child nodes Args: new_children (List[Self]): child node """ if ASSERTIONS: self.__check_children_type(new_children) self.__check_children_loop(new_children) new_children = list(new_children) current_new_children = { new_child: (new_child.parent.__children.index(new_child), new_child.parent) for new_child in new_children if new_child.parent is not None } current_new_orphan = [ new_child for new_child in new_children if new_child.parent is None ] current_children = list(self.children) # Assign new children - rollback if error self.__pre_assign_children(new_children) try: # Remove old children from self del self.children # Assign new children to self self.__children = new_children for new_child in new_children: if new_child.parent: new_child.parent.__children.remove(new_child) new_child.__parent = self self.__post_assign_children(new_children) except Exception as exc_info: # Reassign new children to their original parent for child, idx_parent in current_new_children.items(): child_idx, parent = idx_parent child.__parent = parent parent.__children.insert(child_idx, child) for child in current_new_orphan: child.__parent = None # Reassign old children to self self.__children = current_children for child in current_children: child.__parent = self raise TreeError(exc_info) @children.deleter def children(self) -> None: """Delete child node(s)""" for child in self.children: child.parent.__children.remove(child) # type: ignore child.__parent = None def __pre_assign_children(self: T, new_children: Iterable[T]) -> None: """Custom method to check before attaching children Can be overridden with `_BaseNode__pre_assign_children()` Args: new_children (Iterable[Self]): new children to be added """ pass def __post_assign_children(self: T, new_children: Iterable[T]) -> None: """Custom method to check after attaching children Can be overridden with `_BaseNode__post_assign_children()` Args: new_children (Iterable[Self]): new children to be added """ pass @property def ancestors(self: T) -> Iterable[T]: """Get iterator to yield all ancestors of self, does not include self Returns: (Iterable[Self]) """ node = self.parent while node is not None: yield node node = node.parent @property def descendants(self: T) -> Iterable[T]: """Get iterator to yield all descendants of self, does not include self Returns: (Iterable[Self]) """ yield from preorder_iter(self, filter_condition=lambda _node: _node != self) @property def leaves(self: T) -> Iterable[T]: """Get iterator to yield all leaf nodes from self Returns: (Iterable[Self]) """ yield from preorder_iter(self, filter_condition=lambda _node: _node.is_leaf) @property def siblings(self: T) -> Iterable[T]: """Get siblings of self Returns: (Iterable[Self]) """ if self.parent is None: return () return tuple(child for child in self.parent.children if child is not self) @property def left_sibling(self: T) -> T: """Get sibling left of self Returns: (Self) """ if self.parent: children = self.parent.children child_idx = children.index(self) if child_idx: return self.parent.children[child_idx - 1] @property def right_sibling(self: T) -> T: """Get sibling right of self Returns: (Self) """ if self.parent: children = self.parent.children child_idx = children.index(self) if child_idx + 1 < len(children): return self.parent.children[child_idx + 1] @property def node_path(self: T) -> Iterable[T]: """Get tuple of nodes starting from root Returns: (Iterable[Self]) """ if self.parent is None: return [self] return tuple(list(self.parent.node_path) + [self]) @property def is_root(self) -> bool: """Get indicator if self is root node Returns: (bool) """ return self.parent is None @property def is_leaf(self) -> bool: """Get indicator if self is leaf node Returns: (bool) """ return not len(list(self.children)) @property def root(self: T) -> T: """Get root node of tree Returns: (Self) """ if self.parent is None: return self return self.parent.root @property def depth(self) -> int: """Get depth of self, indexing starts from 1 Returns: (int) """ if self.parent is None: return 1 return self.parent.depth + 1 @property def max_depth(self) -> int: """Get maximum depth from root to leaf node Returns: (int) """ return max( [self.root.depth] + [node.depth for node in list(self.root.descendants)] ) @classmethod def from_dict(cls, input_dict: Dict[str, Any]) -> BaseNode: """Construct node from dictionary, all keys of dictionary will be stored as class attributes Input dictionary must have key `name` if not `Node` will not have any name Examples: >>> from bigtree import Node >>> a = Node.from_dict({"name": "a", "age": 90}) Args: input_dict (Dict[str, Any]): dictionary with node information, key: attribute name, value: attribute value Returns: (BaseNode) """ return cls(**input_dict) def describe( self, exclude_attributes: List[str] = [], exclude_prefix: str = "" ) -> List[Tuple[str, Any]]: """Get node information sorted by attribute name, returns list of tuples Examples: >>> from bigtree.node.node import Node >>> a = Node('a', age=90) >>> a.describe() [('_BaseNode__children', []), ('_BaseNode__parent', None), ('_sep', '/'), ('age', 90), ('name', 'a')] >>> a.describe(exclude_prefix="_") [('age', 90), ('name', 'a')] >>> a.describe(exclude_prefix="_", exclude_attributes=["name"]) [('age', 90)] Args: exclude_attributes (List[str]): list of attributes to exclude exclude_prefix (str): prefix of attributes to exclude Returns: (List[Tuple[str, Any]]) """ return [ item for item in sorted(self.__dict__.items(), key=lambda item: item[0]) if (item[0] not in exclude_attributes) and (not len(exclude_prefix) or not item[0].startswith(exclude_prefix)) ] def get_attr(self, attr_name: str, default_value: Any = None) -> Any: """Get value of node attribute Returns default value if attribute name does not exist Examples: >>> from bigtree.node.node import Node >>> a = Node('a', age=90) >>> a.get_attr("age") 90 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 default_value def set_attrs(self, attrs: Dict[str, Any]) -> None: """Set node attributes Examples: >>> from bigtree.node.node import Node >>> a = Node('a') >>> a.set_attrs({"age": 90}) >>> a Node(/a, age=90) Args: attrs (Dict[str, Any]): attribute dictionary, key: attribute name, value: attribute value """ self.__dict__.update(attrs) def go_to(self: T, node: T) -> Iterable[T]: """Get path from current node to specified node from same tree Examples: >>> from bigtree import Node, print_tree >>> a = Node(name="a") >>> b = Node(name="b", parent=a) >>> c = Node(name="c", parent=a) >>> d = Node(name="d", parent=b) >>> e = Node(name="e", parent=b) >>> f = Node(name="f", parent=c) >>> g = Node(name="g", parent=e) >>> h = Node(name="h", parent=e) >>> print_tree(a) a ├── b │ ├── d │ └── e │ ├── g │ └── h └── c └── f >>> d.go_to(d) [Node(/a/b/d, )] >>> d.go_to(g) [Node(/a/b/d, ), Node(/a/b, ), Node(/a/b/e, ), Node(/a/b/e/g, )] >>> d.go_to(f) [Node(/a/b/d, ), Node(/a/b, ), Node(/a, ), Node(/a/c, ), Node(/a/c/f, )] Args: node (Self): node to travel to from current node, inclusive of start and end node Returns: (Iterable[Self]) """ if not isinstance(node, BaseNode): raise TypeError( f"Expect node to be BaseNode type, received input type {type(node)}" ) if self.root != node.root: raise TreeError( f"Nodes are not from the same tree. Check {self} and {node}" ) if self == node: return [self] self_path = [self] + list(self.ancestors) node_path = ([node] + list(node.ancestors))[::-1] common_nodes = set(self_path).intersection(set(node_path)) self_min_index, min_common_node = sorted( [(self_path.index(_node), _node) for _node in common_nodes] )[0] node_min_index = node_path.index(min_common_node) return self_path[:self_min_index] + node_path[node_min_index:] def append(self: T, other: T) -> None: """Add other as child of self Args: other (Self): other node, child to be added """ other.parent = self def extend(self: T, others: List[T]) -> None: """Add others as children of self Args: others (Self): other nodes, children to be added """ for child in others: child.parent = self def copy(self: T) -> T: """Deep copy self; clone self Examples: >>> from bigtree.node.node import Node >>> a = Node('a') >>> a_copy = a.copy() Returns: (Self) """ return copy.deepcopy(self) def sort(self: T, **kwargs: Any) -> None: """Sort children, possible keyword arguments include ``key=lambda node: node.name``, ``reverse=True`` Examples: >>> from bigtree import Node, print_tree >>> a = Node('a') >>> c = Node("c", parent=a) >>> b = Node("b", parent=a) >>> print_tree(a) a ├── c └── b >>> a.sort(key=lambda node: node.name) >>> print_tree(a) a ├── b └── c """ children = list(self.children) children.sort(**kwargs) self.__children = children def __copy__(self: T) -> T: """Shallow copy self Examples: >>> import copy >>> from bigtree.node.node import Node >>> a = Node('a') >>> a_copy = copy.deepcopy(a) Returns: (Self) """ obj: T = type(self).__new__(self.__class__) obj.__dict__.update(self.__dict__) return obj def __repr__(self) -> str: """Print format of BaseNode Returns: (str) """ class_name = self.__class__.__name__ node_dict = self.describe(exclude_prefix="_") node_description = ", ".join([f"{k}={v}" for k, v in node_dict]) return f"{class_name}({node_description})" def __rshift__(self: T, other: T) -> None: """Set children using >> bitshift operator for self >> children (other) Args: other (Self): other node, children """ other.parent = self def __lshift__(self: T, other: T) -> None: """Set parent using << bitshift operator for self << parent (other) Args: other (Self): other node, parent """ self.parent = 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=BaseNode)