added 3rd party packages, elog, bigtree

python37/packages/bigtree/tree/helper.py

@@ -0,0 +1,201 @@
+from typing import Optional, Type
+
+import numpy as np
+
+from bigtree.node.basenode import BaseNode
+from bigtree.node.binarynode import BinaryNode
+from bigtree.node.node import Node
+from bigtree.tree.construct import dataframe_to_tree
+from bigtree.tree.export import tree_to_dataframe
+from bigtree.tree.search import find_path
+from bigtree.utils.exceptions import NotFoundError
+
+__all__ = ["clone_tree", "prune_tree", "get_tree_diff"]
+
+
+def clone_tree(tree: BaseNode, node_type: Type[BaseNode]) -> BaseNode:
+    """Clone tree to another `Node` type.
+    If the same type is needed, simply do a tree.copy().
+
+    >>> from bigtree import BaseNode, Node, clone_tree
+    >>> root = BaseNode(name="a")
+    >>> b = BaseNode(name="b", parent=root)
+    >>> clone_tree(root, Node)
+    Node(/a, )
+
+    Args:
+        tree (BaseNode): tree to be cloned, must inherit from BaseNode
+        node_type (Type[BaseNode]): type of cloned tree
+
+    Returns:
+        (BaseNode)
+    """
+    if not isinstance(tree, BaseNode):
+        raise ValueError(
+            "Tree should be of type `BaseNode`, or inherit from `BaseNode`"
+        )
+
+    # Start from root
+    root_info = dict(tree.root.describe(exclude_prefix="_"))
+    root_node = node_type(**root_info)
+
+    def recursive_add_child(_new_parent_node, _parent_node):
+        for _child in _parent_node.children:
+            if _child:
+                child_info = dict(_child.describe(exclude_prefix="_"))
+                child_node = node_type(**child_info)
+                child_node.parent = _new_parent_node
+                recursive_add_child(child_node, _child)
+
+    recursive_add_child(root_node, tree.root)
+    return root_node
+
+
+def prune_tree(tree: Node, prune_path: str, sep: str = "/") -> Node:
+    """Prune tree to leave only the prune path, 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.
+
+    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, print_tree
+    >>> root = Node("a")
+    >>> b = Node("b", parent=root)
+    >>> c = Node("c", parent=root)
+    >>> print_tree(root)
+    a
+    ├── b
+    └── c
+
+    >>> root_pruned = prune_tree(root, "a/b")
+    >>> print_tree(root_pruned)
+    a
+    └── b
+
+    Args:
+        tree (Node): existing tree
+        prune_path (str): prune path, all siblings along the prune path will be removed
+        sep (str): path separator
+
+    Returns:
+        (Node)
+    """
+    prune_path = prune_path.replace(sep, tree.sep)
+    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):
+        while child.parent:
+            child.parent.children = [child, None]
+            child = child.parent
+        return tree_copy
+
+    while child.parent:
+        child.parent.children = [child]
+        child = child.parent
+    return tree_copy
+
+
+def get_tree_diff(
+    tree: Node, other_tree: Node, only_diff: bool = True
+) -> Optional[Node]:
+    """Get difference of `tree` to `other_tree`, changes are relative to `tree`.
+
+    (+) and (-) will be added relative to `tree`.
+      - For example: (+) refers to nodes that are in `other_tree` but not `tree`.
+      - For example: (-) refers to nodes that are in `tree` but not `other_tree`.
+
+    Note that only leaf nodes are compared and have (+) or (-) indicator. Intermediate parent nodes are not compared.
+
+    Function can return all original tree nodes and differences, or only the differences.
+
+    >>> from bigtree import Node, get_tree_diff, print_tree
+    >>> root = Node("a")
+    >>> b = Node("b", parent=root)
+    >>> c = Node("c", parent=root)
+    >>> d = Node("d", parent=b)
+    >>> e = Node("e", parent=root)
+    >>> print_tree(root)
+    a
+    ├── b
+    │   └── d
+    ├── c
+    └── e
+
+    >>> root_other = Node("a")
+    >>> b_other = Node("b", parent=root_other)
+    >>> c_other = Node("c", parent=b_other)
+    >>> d_other = Node("d", parent=root_other)
+    >>> e_other = Node("e", parent=root_other)
+    >>> print_tree(root_other)
+    a
+    ├── b
+    │   └── c
+    ├── d
+    └── e
+
+    >>> tree_diff = get_tree_diff(root, root_other)
+    >>> print_tree(tree_diff)
+    a
+    ├── b
+    │   ├── c (+)
+    │   └── d (-)
+    ├── c (-)
+    └── d (+)
+
+    >>> tree_diff = get_tree_diff(root, root_other, only_diff=False)
+    >>> print_tree(tree_diff)
+    a
+    ├── b
+    │   ├── c (+)
+    │   └── d (-)
+    ├── c (-)
+    ├── d (+)
+    └── e
+
+    Args:
+        tree (Node): tree to be compared against
+        other_tree (Node): tree to be compared with
+        only_diff (bool): indicator to show all nodes or only nodes that are different (+/-), defaults to True
+
+    Returns:
+        (Node)
+    """
+    tree = tree.copy()
+    other_tree = other_tree.copy()
+    name_col = "name"
+    path_col = "PATH"
+    indicator_col = "Exists"
+
+    data = tree_to_dataframe(tree, name_col=name_col, path_col=path_col, leaf_only=True)
+    data_other = tree_to_dataframe(
+        other_tree, name_col=name_col, path_col=path_col, leaf_only=True
+    )
+    data_both = data[[path_col, name_col]].merge(
+        data_other[[path_col, name_col]], how="outer", indicator=indicator_col
+    )
+
+    data_both[name_col] = np.where(
+        data_both[indicator_col] == "left_only",
+        data_both[name_col] + " (-)",
+        np.where(
+            data_both[indicator_col] == "right_only",
+            data_both[name_col] + " (+)",
+            data_both[name_col],
+        ),
+    )
+
+    if only_diff:
+        data_both = data_both.query(f"{indicator_col} != 'both'")
+    data_both = data_both.drop(columns=indicator_col).sort_values(path_col)
+    if len(data_both):
+        return dataframe_to_tree(
+            data_both,
+            node_type=tree.__class__,
+        )