Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/treelib/tree.py: 21%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3# Copyright (C) 2011
4# Brett Alistair Kromkamp - brettkromkamp@gmail.com
5# Copyright (C) 2012-2025
6# Xiaming Chen - chenxm35@gmail.com
7# and other contributors.
8# All rights reserved.
9"""
10Tree structure in `treelib`.
12The :class:`Tree` object defines the tree-like structure based on :class:`Node` objects.
13Trees provide hierarchical data organization with efficient operations for traversal,
14modification, search, and visualization.
16A new tree can be created from scratch without any parameter or a shallow/deep copy of another tree.
17When deep=True, a deepcopy operation is performed on feeding tree parameter and more memory
18is required to create the tree.
20Key Features:
21 - O(1) node lookup and access
22 - Multiple traversal modes (depth-first, breadth-first, zigzag)
23 - Flexible tree modification (add, remove, move, paste)
24 - Rich display options with customizable formatting
25 - JSON/dict export and import capabilities
26 - Subtree operations and filtering
27 - Tree metrics and analysis tools
28"""
29from __future__ import print_function, unicode_literals
31try:
32 from builtins import str as text
33except ImportError:
34 from __builtin__ import str as text # type: ignore
36import codecs
37import json
38import sys
39import uuid
40from copy import deepcopy
41from typing import Any, Callable, List, Optional, Union, cast
43from six import iteritems, python_2_unicode_compatible
45try:
46 from StringIO import StringIO # type: ignore
47except ImportError:
48 from io import StringIO
50from .exceptions import (
51 DuplicatedNodeIdError,
52 InvalidLevelNumber,
53 LinkPastRootNodeError,
54 LoopError,
55 MultipleRootError,
56 NodeIDAbsentError,
57)
58from .node import Node
60if sys.version_info >= (3, 9):
61 StrList = list[str]
62 StrLList = list[list[str]]
63 NodeList = list[Node]
64else:
65 StrList = List[str] # Python 3.8 and earlier
66 StrLList = List[List[str]]
67 NodeList = List[Node]
70__author__ = "chenxm"
73@python_2_unicode_compatible
74class Tree(object):
75 """
76 Hierarchical tree data structure.
78 A Tree is a collection of Node objects organized in a hierarchical structure
79 with exactly one root node and zero or more child nodes. Each node (except root)
80 has exactly one parent, but can have multiple children.
82 The tree provides efficient operations for:
83 - Adding, removing, and moving nodes
84 - Traversing the tree in different orders
85 - Searching and filtering nodes
86 - Displaying tree structure
87 - Exporting to various formats
89 Attributes:
90 root (str): Identifier of the root node, or None if tree is empty.
91 nodes (dict): Dictionary mapping node identifiers to Node objects.
93 Constants:
94 ROOT (int): Tree traversal starting from root level.
95 DEPTH (int): Depth-first tree traversal mode.
96 WIDTH (int): Breadth-first tree traversal mode.
97 ZIGZAG (int): Zigzag tree traversal mode.
99 Example:
100 Creating and manipulating a tree::
102 tree = Tree()
104 # Build structure
105 tree.create_node("Company", "company")
106 tree.create_node("Engineering", "eng", parent="company")
107 tree.create_node("Sales", "sales", parent="company")
109 # Add team members
110 tree.create_node("Alice", "alice", parent="eng")
111 tree.create_node("Bob", "bob", parent="eng")
112 tree.create_node("Carol", "carol", parent="sales")
114 # Display tree
115 tree.show()
117 # Find specific nodes
118 eng_team = tree.children("eng")
119 all_employees = [tree[node].tag for node in tree.expand_tree()
120 if tree.level(node) > 1]
122 # Tree operations
123 tree.move_node("alice", "sales") # Alice moves to sales
124 subtree = tree.subtree("eng") # Get engineering subtree
125 """
127 #: ROOT, DEPTH, WIDTH, ZIGZAG constants :
128 (ROOT, DEPTH, WIDTH, ZIGZAG) = list(range(4))
129 node_class = Node
131 def __contains__(self, identifier: str) -> bool:
132 """
133 Check if a node with the given identifier exists in this tree.
135 Implements the 'in' operator for tree membership testing, providing
136 a convenient way to check node existence without raising exceptions.
138 Args:
139 identifier: Node identifier to check for existence.
141 Returns:
142 bool: True if node exists in tree, False otherwise.
144 Example:
145 Checking node membership::
147 if "user123" in tree:
148 print("User node exists")
149 user = tree["user123"]
150 else:
151 print("User node not found")
153 # More concise than try/except approach
154 exists = "node_id" in tree # True or False
155 """
156 return identifier in self.nodes.keys()
158 def __init__(
159 self,
160 tree=None,
161 deep: bool = False,
162 node_class=None,
163 identifier: Optional[str] = None,
164 ) -> None:
165 """
166 Initialize a new tree or copy another tree.
168 Creates either an empty tree or a copy of an existing tree. When copying,
169 you can choose between shallow copy (references to same nodes) or deep copy
170 (completely independent nodes and data).
172 Args:
173 tree: Existing Tree object to copy from. If None, creates empty tree.
174 deep: If True, perform deep copy of all nodes and their data.
175 If False, creates shallow copy sharing node references.
176 node_class: Custom Node class to use instead of default Node.
177 Must be subclass of Node.
178 identifier: Unique identifier for this tree instance.
179 If None, generates UUID automatically.
181 Example:
182 Different ways to create trees::
184 # Empty tree
185 tree1 = Tree()
187 # Tree with custom identifier
188 tree2 = Tree(identifier="my_tree")
190 # Shallow copy of existing tree
191 tree3 = Tree(tree1)
193 # Deep copy with independent data
194 tree4 = Tree(tree1, deep=True)
196 # Custom node class
197 class MyNode(Node):
198 def __init__(self, tag, identifier=None):
199 super().__init__(tag, identifier)
200 self.custom_attr = "value"
202 tree5 = Tree(node_class=MyNode)
204 Raises:
205 AssertionError: If node_class is not a subclass of Node.
206 """
207 self._identifier: Optional[str] = None
208 self._set_identifier(identifier)
210 if node_class:
211 assert issubclass(node_class, Node)
212 self.node_class = node_class
214 #: dictionary, identifier: Node object
215 self._nodes = {}
217 #: Get or set the identifier of the root. This attribute can be accessed and modified
218 #: with ``.`` and ``=`` operator respectively.
219 self.root: Optional[str] = None
221 if tree is not None:
222 self.root = tree.root
223 for nid, node in iteritems(tree.nodes):
224 new_node = deepcopy(node) if deep else node
225 self._nodes[nid] = new_node
226 if tree.identifier != self._identifier:
227 new_node.clone_pointers(tree.identifier, self._identifier)
229 def _clone(
230 self,
231 identifier: Optional[str] = None,
232 with_tree: bool = False,
233 deep: bool = False,
234 ):
235 """
236 Create a clone of this tree instance with optional content copying.
238 This method is designed to be overridden by subclasses to enable
239 proper cloning of extended tree classes. It provides the foundation
240 for subtree and remove_subtree operations while maintaining
241 polymorphic behavior.
243 Args:
244 identifier: Unique identifier for the new tree instance.
245 If None, generates UUID automatically.
246 with_tree: If True, copy all nodes from current tree.
247 If False, create empty tree with same class.
248 deep: If True and with_tree=True, perform deep copy of node data.
249 If False, create shallow copy sharing node references.
251 Returns:
252 Tree: New tree instance of the same class as this tree.
254 Example:
255 Subclassing with custom clone behavior::
257 class EnhancedTree(Tree):
258 def __init__(self, metadata=None, **kwargs):
259 super().__init__(**kwargs)
260 self.metadata = metadata or {}
262 def _clone(self, identifier=None, with_tree=False, deep=False):
263 return EnhancedTree(
264 metadata=self.metadata.copy(),
265 identifier=identifier,
266 tree=self if with_tree else None,
267 deep=deep
268 )
270 # Custom tree operations preserve metadata
271 enhanced = EnhancedTree(metadata={"version": "1.0"})
272 subtree = enhanced.subtree("node_id") # Preserves metadata
273 """
274 return self.__class__(identifier=identifier, tree=self if with_tree else None, deep=deep)
276 @property
277 def identifier(self) -> Optional[str]:
278 """
279 Get the unique identifier of this tree instance.
281 Each tree has its own unique identifier used to distinguish it from
282 other trees, especially when nodes exist in multiple trees simultaneously.
283 This identifier is automatically generated if not provided during creation.
285 Returns:
286 str or None: The unique identifier of this tree instance.
288 Example:
289 Using tree identifiers::
291 tree1 = Tree(identifier="main_tree")
292 tree2 = Tree() # Auto-generated identifier
294 print(tree1.identifier) # "main_tree"
295 print(tree2.identifier) # UUID string like "abc123..."
297 # Used internally for node relationships
298 node.predecessor(tree1.identifier) # Parent in tree1
299 """
300 return self._identifier
302 def _set_identifier(self, nid: Optional[str]) -> None:
303 """
304 Initialize tree identifier with given value or auto-generate one.
306 Private method used during tree creation to set the unique identifier.
307 If no identifier is provided, generates a UUID automatically to ensure
308 uniqueness across tree instances.
310 Args:
311 nid: Desired tree identifier, or None to auto-generate.
313 Note:
314 This is an internal method used during tree initialization.
315 The identifier should not be changed after tree creation.
316 """
317 if nid is None:
318 self._identifier = str(uuid.uuid1())
319 else:
320 self._identifier = nid
322 def __getitem__(self, key: str) -> Node:
323 """
324 Get node by identifier using bracket notation.
326 Provides convenient dictionary-style access to nodes. Raises exception
327 if node doesn't exist, making it clear when invalid identifiers are used.
328 For safer access that returns None instead of raising, use get_node().
330 Args:
331 key: Node identifier to retrieve.
333 Returns:
334 Node: The node object with the specified identifier.
336 Raises:
337 NodeIDAbsentError: If no node with the given identifier exists.
339 Example:
340 Accessing nodes by identifier::
342 # Direct access (raises exception if not found)
343 user_node = tree["user123"]
344 profile_node = tree["user123_profile"]
346 # Use with in operator for safety
347 if "user123" in tree:
348 user_node = tree["user123"]
349 print(user_node.tag)
350 """
351 try:
352 return self._nodes[key]
353 except KeyError:
354 raise NodeIDAbsentError("Node '%s' is not in the tree" % key)
356 def __len__(self) -> int:
357 """
358 Get the total number of nodes in this tree.
360 Returns the count of all nodes currently in the tree, including
361 the root node. Useful for tree size analysis and iteration bounds.
363 Returns:
364 int: Total number of nodes in the tree.
366 Example:
367 Getting tree size::
369 tree_size = len(tree)
370 print(f"Tree has {tree_size} nodes")
372 # Empty tree check
373 if len(tree) == 0:
374 print("Tree is empty")
376 # Use in comparisons
377 if len(tree1) > len(tree2):
378 print("Tree1 is larger")
379 """
380 return len(self._nodes)
382 def __str__(self) -> str:
383 """
384 Get string representation of the tree structure.
386 Returns a formatted tree visualization showing the hierarchical structure
387 with default display settings. Useful for debugging, logging, and quick
388 inspection of tree contents.
390 Returns:
391 str: Formatted tree structure as string.
393 Example:
394 Tree string representation::
396 tree = Tree()
397 tree.create_node("Root", "root")
398 tree.create_node("Child A", "a", parent="root")
399 tree.create_node("Child B", "b", parent="root")
401 print(str(tree))
402 # Output:
403 # Root
404 # ├── Child A
405 # └── Child B
407 # Use in logging
408 logger.info(f"Current tree structure:\n{tree}")
409 """
410 self._reader = ""
412 def write(line):
413 self._reader += line.decode("utf-8") + "\n"
415 self.__print_backend(func=write)
416 return self._reader
418 def __print_backend(
419 self,
420 nid: Optional[str] = None,
421 level: int = ROOT,
422 idhidden: bool = True,
423 filter: Optional[Callable[[Node], bool]] = None,
424 key: Optional[Callable[[Node], Node]] = None,
425 reverse: bool = False,
426 line_type="ascii-ex",
427 data_property=None,
428 sorting: bool = True,
429 func: Callable[[bytes], None] = print,
430 ):
431 """
432 Another implementation of printing tree using Stack
433 Print tree structure in hierarchy style.
435 For example:
437 .. code-block:: bash
439 Root
440 |___ C01
441 | |___ C11
442 | |___ C111
443 | |___ C112
444 |___ C02
445 |___ C03
446 | |___ C31
448 A more elegant way to achieve this function using Stack
449 structure, for constructing the Nodes Stack push and pop nodes
450 with additional level info.
452 UPDATE: the @key @reverse is present to sort node at each
453 level.
454 """
455 # Factory for proper get_label() function
456 if data_property:
457 if idhidden:
459 def get_label(node: Node) -> str:
460 return getattr(node.data, data_property)
462 else:
464 def get_label(node: Node) -> str:
465 return "%s[%s]" % (
466 getattr(node.data, data_property),
467 node.identifier,
468 )
470 else:
471 if idhidden:
473 def get_label(node: Node) -> str:
474 return node.tag
476 else:
478 def get_label(node: Node) -> str:
479 return "%s[%s]" % (node.tag, node.identifier)
481 # legacy ordering
482 if sorting and key is None:
484 def key(node):
485 return node
487 # iter with func
488 for pre, node in self.__get(nid, level, filter, key, reverse, line_type, sorting):
489 label = get_label(node)
490 func("{0}{1}".format(pre, label).encode("utf-8"))
492 def __get(
493 self,
494 nid: Optional[str],
495 level: int,
496 filter_: Optional[Callable[[Node], bool]],
497 key: Optional[Callable[[Node], Node]],
498 reverse: bool,
499 line_type: str,
500 sorting: bool,
501 ):
502 # default filter
503 if filter_ is None:
505 def filter_(node):
506 return True
508 # render characters
509 dt = {
510 "ascii": ("|", "|-- ", "+-- "),
511 "ascii-ex": ("\u2502", "\u251c\u2500\u2500 ", "\u2514\u2500\u2500 "),
512 "ascii-exr": ("\u2502", "\u251c\u2500\u2500 ", "\u2570\u2500\u2500 "),
513 "ascii-em": ("\u2551", "\u2560\u2550\u2550 ", "\u255a\u2550\u2550 "),
514 "ascii-emv": ("\u2551", "\u255f\u2500\u2500 ", "\u2559\u2500\u2500 "),
515 "ascii-emh": ("\u2502", "\u255e\u2550\u2550 ", "\u2558\u2550\u2550 "),
516 }[line_type]
518 return self.__get_iter(nid, level, filter_, key, reverse, dt, [], sorting)
520 def __get_iter(
521 self,
522 nid: Optional[str],
523 level: int,
524 filter_: Callable[[Node], bool],
525 key: Optional[Callable[[Node], Node]],
526 reverse: bool,
527 dt, # tuple[str, str, str]
528 is_last: list,
529 sorting: bool,
530 ):
531 dt_vertical_line, dt_line_box, dt_line_corner = dt
533 nid = cast(str, self.root if nid is None else nid)
534 node = self[nid]
536 if level == self.ROOT:
537 yield "", node
538 else:
539 leading = "".join(
540 map(
541 lambda x: dt_vertical_line + " " * 3 if not x else " " * 4,
542 is_last[0:-1],
543 )
544 )
545 lasting = dt_line_corner if is_last[-1] else dt_line_box
546 yield leading + lasting, node
548 if filter_(node) and node.expanded:
549 children = [self[i] for i in node.successors(self._identifier) if filter_(self[i])]
550 idxlast = len(children) - 1
551 if sorting:
552 if key:
553 children.sort(key=key, reverse=reverse)
554 elif reverse:
555 children = list(reversed(children))
556 level += 1
557 for idx, child in enumerate(children):
558 is_last.append(idx == idxlast)
559 for item in self.__get_iter(child.identifier, level, filter_, key, reverse, dt, is_last, sorting):
560 yield item
561 is_last.pop()
563 def __update_bpointer(self, nid, parent_id):
564 """set self[nid].bpointer"""
565 self[nid].set_predecessor(parent_id, self._identifier)
567 def __update_fpointer(self, nid: str, child_id: str, mode: int):
568 if nid is None:
569 return
570 else:
571 self[nid].update_successors(child_id, mode, tree_id=self._identifier)
573 def add_node(self, node: Node, parent: Optional[Union[Node, str]] = None) -> None:
574 """
575 Add an existing node object to the tree.
577 Integrates a pre-created Node instance into the tree structure by
578 establishing parent-child relationships and updating internal mappings.
579 For creating and adding nodes simultaneously, use create_node() instead.
581 Args:
582 node: Existing Node instance to add to the tree.
583 Must be an instance of the tree's node_class.
584 parent: Parent node (Node object or identifier string), or None
585 to make this node the root. If None, tree must be empty.
587 Raises:
588 OSError: If node is not an instance of the expected node class.
589 DuplicatedNodeIdError: If node identifier already exists in tree.
590 MultipleRootError: If parent is None but tree already has a root.
591 NodeIDAbsentError: If parent identifier doesn't exist in tree.
593 Example:
594 Adding pre-created nodes::
596 # Create nodes first
597 root_node = Node("Company", "company")
598 dept_node = Node("Engineering", "eng")
600 # Add to tree
601 tree.add_node(root_node) # Root node
602 tree.add_node(dept_node, parent="company") # Child node
604 # Adding with Node object as parent
605 mgr_node = Node("Manager", "mgr")
606 tree.add_node(mgr_node, parent=dept_node)
607 """
608 if not isinstance(node, self.node_class):
609 raise OSError("First parameter must be object of {}".format(self.node_class))
611 if node.identifier in self._nodes:
612 raise DuplicatedNodeIdError("Can't create node " "with ID '%s'" % node.identifier)
614 pid = parent.identifier if isinstance(parent, self.node_class) else parent
616 if pid is None:
617 if self.root is not None:
618 raise MultipleRootError("A tree takes one root merely.")
619 else:
620 self.root = node.identifier
621 elif not self.contains(pid):
622 raise NodeIDAbsentError("Parent node '%s' " "is not in the tree" % pid)
624 pid = cast(str, pid)
625 self._nodes.update({node.identifier: node})
626 self.__update_fpointer(pid, node.identifier, self.node_class.ADD)
627 self.__update_bpointer(node.identifier, pid)
628 node.set_initial_tree_id(cast(str, self._identifier))
630 def all_nodes(self) -> NodeList:
631 """
632 Get all nodes in the tree as a list.
634 Returns a list containing all Node objects currently in the tree.
635 The order is not guaranteed. For ordered traversal, use expand_tree().
636 Useful for operations that need to process all nodes simultaneously.
638 Returns:
639 list[Node]: List of all Node objects in the tree.
641 Example:
642 Processing all nodes::
644 # Get all nodes
645 all_nodes = tree.all_nodes()
646 print(f"Total nodes: {len(all_nodes)}")
648 # Process each node
649 for node in all_nodes:
650 print(f"Node: {node.tag}")
652 # Filter nodes by property
653 leaf_nodes = [node for node in tree.all_nodes()
654 if node.is_leaf(tree.identifier)]
655 """
656 return list(self._nodes.values())
658 def all_nodes_itr(self) -> Any:
659 """
660 Get all nodes in the tree as an iterator.
662 Returns an iterator over all Node objects in the tree, providing
663 memory-efficient access for large trees. Useful when you need to
664 process nodes one at a time without loading all into memory.
666 Returns:
667 Iterator[Node]: Iterator over all Node objects in the tree.
669 Example:
670 Memory-efficient node processing::
672 # Iterate without loading all nodes
673 for node in tree.all_nodes_itr():
674 if node.tag.startswith("temp_"):
675 process_temporary_node(node)
677 # Use with filter operations
678 filtered = filter(lambda n: n.data is not None,
679 tree.all_nodes_itr())
681 Note:
682 Added by William Rusnack for memory efficiency.
683 """
684 return self._nodes.values()
686 def ancestor(self, nid, level=None):
687 """
688 Get the ancestor node at a specific level above the given node.
690 Traverses up the tree hierarchy from the specified node to find an
691 ancestor at the desired level. If no level is specified, returns
692 the immediate parent. Useful for navigating tree hierarchies.
694 Args:
695 nid: Identifier of the node to start from.
696 level: Target level of the ancestor (0 = root). If None,
697 returns immediate parent.
699 Returns:
700 str or None: Identifier of the ancestor node, or None if not found.
702 Raises:
703 NodeIDAbsentError: If nid doesn't exist in the tree.
704 InvalidLevelNumber: If level is invalid (>= descendant level).
706 Example:
707 Finding ancestors::
709 # Get immediate parent
710 parent_id = tree.ancestor("grandchild")
712 # Get ancestor at specific level
713 root_id = tree.ancestor("grandchild", level=0) # Root
714 grandparent_id = tree.ancestor("grandchild", level=1)
716 # Use in hierarchy navigation
717 if tree.ancestor("node", level=0) == "root":
718 print("Node is in main hierarchy")
719 """
720 if not self.contains(nid):
721 raise NodeIDAbsentError("Node '%s' is not in the tree" % nid)
723 descendant = self[nid]
724 ascendant = self[nid]._predecessor[self._identifier]
725 ascendant_level = self.level(ascendant)
727 if level is None:
728 return ascendant
729 elif nid == self.root:
730 return self[nid]
731 elif level >= self.level(descendant.identifier):
732 raise InvalidLevelNumber(
733 "Descendant level (level %s) must be greater \
734 than its ancestor's level (level %s)"
735 % (str(self.level(descendant.identifier)), level)
736 )
738 while ascendant is not None:
739 if ascendant_level == level:
740 return self[ascendant]
741 else:
742 descendant = ascendant
743 ascendant = self[descendant]._predecessor[self._identifier]
744 ascendant_level = self.level(ascendant)
745 return None
747 def children(self, nid: str) -> NodeList:
748 """
749 Get all direct children of the specified node.
751 Returns a list of Node objects that are immediate children of the
752 given node. The order follows the insertion order unless the tree
753 has been sorted. Returns empty list if node has no children.
755 Args:
756 nid: Identifier of the parent node.
758 Returns:
759 list[Node]: List of child Node objects. Empty if no children.
761 Raises:
762 NodeIDAbsentError: If nid doesn't exist in the tree.
764 Example:
765 Working with child nodes::
767 # Get all children
768 child_nodes = tree.children("parent_id")
769 print(f"Parent has {len(child_nodes)} children")
771 # Process each child
772 for child in tree.children("parent_id"):
773 print(f"Child: {child.tag}")
775 # Check if node has children
776 if tree.children("node_id"):
777 print("Node has children")
778 else:
779 print("Node is a leaf")
780 """
781 return [self[i] for i in self.is_branch(nid)]
783 def contains(self, nid):
784 """
785 Check if the tree contains a node with the given identifier.
787 Determines whether a node with the specified identifier exists
788 in this tree. Equivalent to using the 'in' operator but provided
789 as an explicit method for clarity and consistency.
791 Args:
792 nid: Node identifier to check for existence.
794 Returns:
795 bool: True if node exists in tree, False otherwise.
797 Example:
798 Checking node existence::
800 # Explicit method call
801 if tree.contains("user123"):
802 print("User node exists")
804 # Equivalent using 'in' operator
805 if "user123" in tree:
806 print("User node exists")
808 # Use before operations
809 if tree.contains("node_id"):
810 tree.move_node("node_id", "new_parent")
811 """
812 return True if nid in self._nodes else False
814 def create_node(
815 self,
816 tag: Optional[str] = None,
817 identifier: Optional[str] = None,
818 parent: Optional[Union[Node, str]] = None,
819 data: Any = None,
820 ) -> Node:
821 """
822 Create and add a new node to the tree.
824 This is the primary method for building tree structures. Creates a new node
825 with the specified properties and attaches it to the given parent. If no
826 parent is specified, the node becomes the root (only allowed if tree is empty).
828 Args:
829 tag: Human-readable label for display. If None, uses identifier.
830 identifier: Unique ID for the node. If None, generates UUID automatically.
831 parent: Parent node identifier, Node object, or None for root.
832 Must be None if tree is empty, must exist if tree has nodes.
833 data: Optional user data to associate with this node.
835 Returns:
836 Node: The newly created Node object.
838 Raises:
839 DuplicatedNodeIdError: If identifier already exists in the tree.
840 MultipleRootError: If parent is None but tree already has a root.
841 NodeIDAbsentError: If parent identifier doesn't exist in the tree.
843 Example:
844 Building a family tree::
846 tree = Tree()
848 # Create root (no parent)
849 tree.create_node("Grandpa", "grandpa")
851 # Add children
852 tree.create_node("Dad", "dad", parent="grandpa")
853 tree.create_node("Uncle", "uncle", parent="grandpa")
855 # Add grandchildren
856 tree.create_node("Me", "me", parent="dad")
857 tree.create_node("Sister", "sister", parent="dad")
859 # Add node with custom data
860 tree.create_node("Baby", "baby", parent="me",
861 data={"age": 1, "cute": True})
862 """
863 node = self.node_class(tag=tag, identifier=identifier, data=data)
864 self.add_node(node, parent)
865 return node
867 def depth(self, node: Optional[Node] = None) -> int:
868 """
869 Get the maximum depth of the tree or the level of a specific node.
871 When called without arguments, returns the maximum depth of the entire
872 tree (distance from root to deepest leaf). When called with a node,
873 returns the level of that specific node (distance from root).
875 Args:
876 node: Node object or identifier string. If None, returns tree depth.
877 If provided, returns the level of this specific node.
879 Returns:
880 int: Tree depth (max level) or node level. Root is at level 0.
882 Raises:
883 NodeIDAbsentError: If specified node doesn't exist in the tree.
885 Example:
886 Measuring tree dimensions::
888 # Get overall tree depth
889 max_depth = tree.depth()
890 print(f"Tree depth: {max_depth}")
892 # Get specific node level
893 node_level = tree.depth("some_node")
894 node_level2 = tree.depth(tree["some_node"]) # Same result
896 # Use for tree analysis
897 if tree.depth() > 5:
898 print("Tree is quite deep")
899 """
900 ret = 0
901 if node is None:
902 # Get maximum level of this tree
903 leaves = self.leaves()
904 for leave in leaves:
905 level = self.level(leave.identifier)
906 ret = level if level >= ret else ret
907 else:
908 # Get level of the given node
909 if not isinstance(node, self.node_class):
910 nid = node
911 else:
912 nid = node.identifier
913 if not self.contains(nid):
914 raise NodeIDAbsentError("Node '%s' is not in the tree" % nid)
915 ret = self.level(nid)
916 return ret
918 def expand_tree(
919 self,
920 nid: Optional[str] = None,
921 mode: int = DEPTH,
922 filter: Optional[Callable[[Node], bool]] = None,
923 key: Optional[Callable[[Node], Node]] = None,
924 reverse: bool = False,
925 sorting: bool = True,
926 ):
927 """
928 Traverse the tree and yield node identifiers in specified order.
930 This is the primary method for tree iteration, providing flexible traversal
931 with multiple algorithms, filtering, and sorting options. Essential for most
932 tree processing operations.
934 Args:
935 nid: Starting node identifier. If None, starts from tree root.
936 Must exist in the tree if specified.
937 mode: Traversal algorithm to use:
938 Tree.DEPTH (0) - Depth-first search (default)
939 Tree.WIDTH (1) - Breadth-first search
940 Tree.ZIGZAG (2) - ZigZag traversal (alternating levels)
941 filter: Optional function to filter nodes during traversal.
942 Takes Node object, returns bool. If False, node and its
943 entire subtree are skipped.
944 key: Sorting function for nodes at each level. Takes Node object,
945 returns comparison key. If None and sorting=True, sorts by node.
946 reverse: If True, reverse the sorting order at each level.
947 sorting: If True, sort nodes at each level using key function.
948 If False, preserve insertion order (key/reverse ignored).
950 Yields:
951 str: Node identifiers in traversal order.
953 Raises:
954 NodeIDAbsentError: If nid doesn't exist in the tree.
955 ValueError: If mode is not a valid traversal constant.
957 Example:
958 Different traversal patterns::
960 tree = Tree()
961 tree.create_node("A", "a")
962 tree.create_node("B", "b", parent="a")
963 tree.create_node("C", "c", parent="a")
964 tree.create_node("D", "d", parent="b")
965 tree.create_node("E", "e", parent="c")
967 # Depth-first (default): A, B, D, C, E
968 for node_id in tree.expand_tree():
969 print(f"DFS: {tree[node_id].tag}")
971 # Breadth-first: A, B, C, D, E
972 for node_id in tree.expand_tree(mode=Tree.WIDTH):
973 print(f"BFS: {tree[node_id].tag}")
975 # ZigZag traversal
976 for node_id in tree.expand_tree(mode=Tree.ZIGZAG):
977 print(f"ZigZag: {tree[node_id].tag}")
979 # Filtered traversal (only nodes with vowels)
980 vowel_filter = lambda node: any(v in node.tag.lower() for v in 'aeiou')
981 for node_id in tree.expand_tree(filter=vowel_filter):
982 print(f"Vowels: {tree[node_id].tag}")
984 # Sorted by tag, reversed
985 for node_id in tree.expand_tree(key=lambda x: x.tag, reverse=True):
986 print(f"Sorted: {tree[node_id].tag}")
988 # From specific subtree
989 for node_id in tree.expand_tree(nid="b"):
990 print(f"Subtree: {tree[node_id].tag}")
992 Common Usage Patterns::
994 # Process all nodes
995 for node_id in tree.expand_tree():
996 process_node(tree[node_id])
998 # Get all node tags
999 tags = [tree[nid].tag for nid in tree.expand_tree()]
1001 # Find specific nodes
1002 matching_nodes = [nid for nid in tree.expand_tree()
1003 if tree[nid].tag.startswith("prefix")]
1005 # Level-order processing
1006 for node_id in tree.expand_tree(mode=Tree.WIDTH):
1007 level = tree.level(node_id)
1008 print(f"Level {level}: {tree[node_id].tag}")
1010 Performance:
1011 - Time complexity: O(n) where n is number of visited nodes
1012 - Memory: O(h) where h is tree height (for traversal stack)
1013 - Generator-based: memory efficient for large trees
1015 Note:
1016 This is a generator function - it yields results lazily. Perfect for
1017 large trees where you might not need to visit all nodes, or when
1018 memory efficiency is important.
1019 """
1020 nid = cast(str, self.root if nid is None else nid)
1021 if not self.contains(nid):
1022 raise NodeIDAbsentError("Node '%s' is not in the tree" % nid)
1024 filter = (lambda x: True) if (filter is None) else filter
1025 if filter(self[nid]):
1026 yield nid
1027 queue = [self[i] for i in self[nid].successors(self._identifier) if filter(self[i])]
1028 if mode in [self.DEPTH, self.WIDTH]:
1029 if sorting:
1030 queue.sort(key=key, reverse=reverse)
1031 while queue:
1032 yield queue[0].identifier
1033 expansion = [self[i] for i in queue[0].successors(self._identifier) if filter(self[i])]
1034 if sorting:
1035 expansion.sort(key=key, reverse=reverse)
1036 if mode is self.DEPTH:
1037 queue = expansion + queue[1:] # depth-first
1038 elif mode is self.WIDTH:
1039 queue = queue[1:] + expansion # width-first
1041 elif mode is self.ZIGZAG:
1042 # Suggested by Ilya Kuprik (ilya-spy@ynadex.ru).
1043 stack_fw: NodeList = []
1044 queue.reverse()
1045 stack = stack_bw = queue
1046 direction = False
1047 while stack:
1048 expansion = [self[i] for i in stack[0].successors(self._identifier) if filter(self[i])]
1049 yield stack.pop(0).identifier
1050 if direction:
1051 expansion.reverse()
1052 stack_bw = expansion + stack_bw
1053 else:
1054 stack_fw = expansion + stack_fw
1055 if not stack:
1056 direction = not direction
1057 stack = stack_fw if direction else stack_bw
1059 else:
1060 raise ValueError("Traversal mode '{}' is not supported".format(mode))
1062 def filter_nodes(self, func: Callable[[Node], bool]):
1063 """
1064 Filter all nodes in the tree using a custom function.
1066 Applies the given function to every node and returns an iterator over
1067 nodes where the function returns True. Useful for finding specific
1068 types of nodes or nodes with certain properties.
1070 Args:
1071 func: Function that takes a Node object and returns bool.
1072 True means include the node in results.
1074 Returns:
1075 Iterator[Node]: Iterator over nodes where func returns True.
1077 Example:
1078 Filtering nodes by criteria::
1080 # Find all leaf nodes
1081 leaves = list(tree.filter_nodes(lambda n: n.is_leaf(tree.identifier)))
1083 # Find nodes with specific data
1084 important_nodes = list(tree.filter_nodes(
1085 lambda n: hasattr(n.data, 'priority') and n.data.priority == 'high'
1086 ))
1088 # Find nodes by tag pattern
1089 temp_nodes = list(tree.filter_nodes(
1090 lambda n: n.tag.startswith('temp_')
1091 ))
1093 # Memory efficient processing
1094 for node in tree.filter_nodes(lambda n: n.data is not None):
1095 process_node_with_data(node)
1097 Note:
1098 Added by William Rusnack for flexible node filtering.
1099 """
1100 return filter(func, self.all_nodes_itr())
1102 def get_node(self, nid: Optional[str]) -> Optional[Node]:
1103 """
1104 Safely get a node by identifier without raising exceptions.
1106 Provides null-safe access to nodes, returning None if the node
1107 doesn't exist instead of raising an exception. Preferred over
1108 bracket notation when node existence is uncertain.
1110 Args:
1111 nid: Node identifier to retrieve, or None.
1113 Returns:
1114 Node or None: The node object if found, None otherwise.
1116 Example:
1117 Safe node access::
1119 # Safe access (no exception)
1120 node = tree.get_node("might_not_exist")
1121 if node is not None:
1122 print(f"Found: {node.tag}")
1123 else:
1124 print("Node not found")
1126 # Compare with bracket notation
1127 try:
1128 node = tree["might_not_exist"] # Raises exception
1129 except NodeIDAbsentError:
1130 node = None
1132 # Use in conditional logic
1133 user_node = tree.get_node("user123")
1134 if user_node and user_node.data.active:
1135 process_active_user(user_node)
1136 """
1137 if nid is None or not self.contains(nid):
1138 return None
1139 return self._nodes[nid]
1141 def is_branch(self, nid):
1142 """
1143 Get the list of child node identifiers for the specified node.
1145 Returns the identifiers of all direct children of the given node.
1146 Despite the name suggesting a boolean, this method actually returns
1147 the list of child identifiers. Use children() for Node objects.
1149 Args:
1150 nid: Identifier of the parent node. Cannot be None.
1152 Returns:
1153 list[str]: List of child node identifiers. Empty if no children.
1155 Raises:
1156 OSError: If nid is None.
1157 NodeIDAbsentError: If nid doesn't exist in the tree.
1159 Example:
1160 Getting child identifiers::
1162 # Get child IDs
1163 child_ids = tree.is_branch("parent_id")
1164 for child_id in child_ids:
1165 print(f"Child ID: {child_id}")
1167 # Check if node has children
1168 if tree.is_branch("node_id"):
1169 print("Node has children")
1171 # Use with node objects
1172 child_nodes = [tree[child_id] for child_id in tree.is_branch("parent")]
1173 """
1174 if nid is None:
1175 raise OSError("First parameter can't be None")
1176 if not self.contains(nid):
1177 raise NodeIDAbsentError("Node '%s' is not in the tree" % nid)
1179 try:
1180 fpointer = self[nid].successors(self._identifier)
1181 except KeyError:
1182 fpointer = []
1183 return fpointer
1185 def leaves(self, nid: Optional[str] = None) -> NodeList:
1186 """
1187 Get all leaf nodes in the tree or subtree.
1189 Returns all nodes that have no children. If a starting node is specified,
1190 only considers leaves within that subtree. Leaf nodes are terminal nodes
1191 in the tree structure and often represent end points in hierarchies.
1193 Args:
1194 nid: Root of subtree to search. If None, searches entire tree.
1196 Returns:
1197 list[Node]: List of all leaf Node objects in the specified scope.
1199 Example:
1200 Finding leaf nodes::
1202 # Get all leaves in tree
1203 all_leaves = tree.leaves()
1204 print(f"Tree has {len(all_leaves)} leaf nodes")
1206 # Get leaves in specific subtree
1207 subtree_leaves = tree.leaves("department_root")
1209 # Process leaf nodes
1210 for leaf in tree.leaves():
1211 print(f"Leaf: {leaf.tag}")
1213 # Find deepest nodes
1214 max_level = max(tree.level(leaf.identifier) for leaf in tree.leaves())
1215 deepest_leaves = [leaf for leaf in tree.leaves()
1216 if tree.level(leaf.identifier) == max_level]
1217 """
1218 leaves = []
1219 if nid is None:
1220 for node in self._nodes.values():
1221 if node.is_leaf(self._identifier):
1222 leaves.append(node)
1223 else:
1224 for node in self.expand_tree(nid):
1225 if self[node].is_leaf(self._identifier):
1226 leaves.append(self[node])
1227 return leaves
1229 def level(self, nid, filter=None):
1230 """
1231 Get the level (depth) of a node in the tree hierarchy.
1233 Returns the distance from the root to the specified node, where
1234 the root is at level 0. Optionally filters the path calculation
1235 by excluding certain nodes.
1237 Args:
1238 nid: Identifier of the node to get level for.
1239 filter: Optional function to filter nodes during path calculation.
1240 Takes Node object, returns bool. Filtered nodes are excluded.
1242 Returns:
1243 int: Level of the node (0 for root, 1 for root's children, etc.)
1245 Example:
1246 Getting node levels::
1248 # Basic level calculation
1249 root_level = tree.level("root") # 0
1250 child_level = tree.level("child") # 1
1251 grandchild_level = tree.level("gc") # 2
1253 # Use for tree analysis
1254 max_depth = max(tree.level(node.identifier)
1255 for node in tree.all_nodes())
1257 # Filtered level calculation
1258 level = tree.level("node", filter=lambda n: n.tag != "skip")
1259 """
1260 return len([n for n in self.rsearch(nid, filter)]) - 1
1262 def link_past_node(self, nid: str):
1263 """
1264 Remove a node while preserving connections between its parent and children.
1266 Deletes the specified node but maintains the tree structure by directly
1267 connecting the node's parent to all of its children. This operation
1268 effectively "links past" the node, removing it from the hierarchy
1269 without breaking the tree connections.
1271 Args:
1272 nid: Identifier of the node to link past and remove.
1274 Raises:
1275 NodeIDAbsentError: If nid doesn't exist in the tree.
1276 LinkPastRootNodeError: If attempting to link past the root node.
1278 Example:
1279 Linking past nodes::
1281 # Before: Root -> Manager -> Employee1, Employee2
1282 tree.link_past_node("Manager")
1283 # After: Root -> Employee1, Employee2
1285 # Useful for flattening hierarchies
1286 middle_managers = ["mgr1", "mgr2", "mgr3"]
1287 for mgr in middle_managers:
1288 if tree.contains(mgr):
1289 tree.link_past_node(mgr)
1291 # Cannot link past root
1292 try:
1293 tree.link_past_node(tree.root)
1294 except LinkPastRootNodeError:
1295 print("Cannot link past root node")
1296 """
1297 if not self.contains(nid):
1298 raise NodeIDAbsentError("Node '%s' is not in the tree" % nid)
1299 if self.root == nid:
1300 raise LinkPastRootNodeError("Cannot link past the root node, " "delete it with remove_node()")
1301 # Get the parent of the node we are linking past
1302 parent = self[self[nid].predecessor(self._identifier)]
1303 # Set the children of the node to the parent
1304 for child in self[nid].successors(self._identifier):
1305 self[child].set_predecessor(parent.identifier, self._identifier)
1306 # Link the children to the parent
1307 for id_ in self[nid].successors(self._identifier) or []:
1308 parent.update_successors(id_, tree_id=self._identifier)
1309 # Delete the node
1310 parent.update_successors(nid, mode=parent.DELETE, tree_id=self._identifier)
1311 del self._nodes[nid]
1313 def move_node(self, source, destination):
1314 """
1315 Move a node from its current parent to a new parent.
1317 Changes the parent-child relationship by moving the source node
1318 (and its entire subtree) from its current parent to the specified
1319 destination parent. This operation preserves the subtree structure
1320 while reorganizing the tree hierarchy.
1322 Args:
1323 source: Identifier of the node to move.
1324 destination: Identifier of the new parent node.
1326 Raises:
1327 NodeIDAbsentError: If source or destination node doesn't exist.
1328 LoopError: If moving would create a circular reference
1329 (destination is a descendant of source).
1331 Example:
1332 Reorganizing tree structure::
1334 # Move employee between departments
1335 tree.move_node("employee123", "sales_dept")
1337 # Move entire department
1338 tree.move_node("engineering_dept", "new_division")
1340 # Prevent circular moves (this would raise LoopError)
1341 try:
1342 tree.move_node("parent", "child_of_parent")
1343 except LoopError:
1344 print("Cannot create circular reference")
1346 # Bulk reorganization
1347 employees = ["emp1", "emp2", "emp3"]
1348 for emp in employees:
1349 tree.move_node(emp, "new_manager")
1350 """
1351 if not self.contains(source) or not self.contains(destination):
1352 raise NodeIDAbsentError
1353 elif self.is_ancestor(source, destination):
1354 raise LoopError
1356 parent = self[source].predecessor(self._identifier)
1357 self.__update_fpointer(parent, source, self.node_class.DELETE)
1358 self.__update_fpointer(destination, source, self.node_class.ADD)
1359 self.__update_bpointer(source, destination)
1361 def is_ancestor(self, ancestor, grandchild):
1362 """
1363 Check if one node is an ancestor of another node.
1365 Determines whether the ancestor node appears in the path from
1366 the grandchild node to the root. An ancestor is any node that
1367 appears above another node in the tree hierarchy.
1369 Args:
1370 ancestor: Identifier of the potential ancestor node.
1371 grandchild: Identifier of the potential descendant node.
1373 Returns:
1374 bool: True if ancestor is indeed an ancestor of grandchild,
1375 False otherwise.
1377 Example:
1378 Checking ancestry relationships::
1380 # Direct parent-child
1381 is_parent = tree.is_ancestor("parent", "child") # True
1383 # Multi-level ancestry
1384 is_grandparent = tree.is_ancestor("grandparent", "grandchild") # True
1386 # Not related
1387 is_ancestor = tree.is_ancestor("sibling1", "sibling2") # False
1389 # Prevent circular moves
1390 if not tree.is_ancestor("node_to_move", "new_parent"):
1391 tree.move_node("node_to_move", "new_parent")
1392 else:
1393 print("Cannot create circular reference")
1394 """
1395 parent = self[grandchild].predecessor(self._identifier)
1396 child = grandchild
1397 while parent is not None:
1398 if parent == ancestor:
1399 return True
1400 else:
1401 child = self[child].predecessor(self._identifier)
1402 parent = self[child].predecessor(self._identifier)
1403 return False
1405 @property
1406 def nodes(self):
1407 """
1408 Get the internal dictionary mapping node identifiers to Node objects.
1410 Returns the underlying dictionary that stores all nodes in the tree,
1411 mapping from node identifiers to Node instances. Useful for direct
1412 access to the node collection and bulk operations.
1414 Returns:
1415 dict[str, Node]: Dictionary mapping node IDs to Node objects.
1417 Example:
1418 Working with the nodes dictionary::
1420 # Get all node identifiers
1421 all_ids = list(tree.nodes.keys())
1423 # Get all node objects
1424 all_nodes = list(tree.nodes.values())
1426 # Direct access to nodes dict
1427 nodes_dict = tree.nodes
1428 for node_id, node in nodes_dict.items():
1429 print(f"ID: {node_id}, Tag: {node.tag}")
1431 # Bulk operations
1432 total_nodes = len(tree.nodes)
1433 has_node = "some_id" in tree.nodes
1434 """
1435 return self._nodes
1437 def parent(self, nid: str) -> Optional[Node]:
1438 """
1439 Get the parent Node object of the specified node.
1441 Returns the immediate parent of the given node, or None if the node
1442 is the root or has no parent. Provides object-based access to the
1443 parent, unlike ancestor() which returns identifiers.
1445 Args:
1446 nid: Identifier of the node whose parent to retrieve.
1448 Returns:
1449 Node or None: Parent Node object, or None if node is root.
1451 Raises:
1452 NodeIDAbsentError: If nid doesn't exist in the tree.
1454 Example:
1455 Working with parent relationships::
1457 # Get parent node
1458 parent_node = tree.parent("child_id")
1459 if parent_node:
1460 print(f"Parent: {parent_node.tag}")
1461 else:
1462 print("Node is root or orphaned")
1464 # Navigate up the hierarchy
1465 current = "leaf_node"
1466 while tree.parent(current):
1467 current = tree.parent(current).identifier
1468 print(f"Ancestor: {tree[current].tag}")
1470 # Check parent properties
1471 parent = tree.parent("employee")
1472 if parent and hasattr(parent.data, 'department'):
1473 print(f"Department: {parent.data.department}")
1474 """
1475 if not self.contains(nid):
1476 raise NodeIDAbsentError("Node '%s' is not in the tree" % nid)
1478 pid = self[nid].predecessor(self._identifier)
1479 if pid is None or not self.contains(pid):
1480 return None
1482 return self[pid]
1484 def merge(self, nid: str, new_tree, deep: bool = False):
1485 """Patch @new_tree on current tree by pasting new_tree root children on current tree @nid node.
1487 Consider the following tree:
1488 >>> current.show()
1489 root
1490 ├── A
1491 └── B
1492 >>> new_tree.show()
1493 root2
1494 ├── C
1495 └── D
1496 └── D1
1497 Merging new_tree on B node:
1498 >>>current.merge('B', new_tree)
1499 >>>current.show()
1500 root
1501 ├── A
1502 └── B
1503 ├── C
1504 └── D
1505 └── D1
1507 Note: if current tree is empty and nid is None, the new_tree root will be used as root on current tree. In all
1508 other cases new_tree root is not pasted.
1509 """
1510 if new_tree.root is None:
1511 return
1513 if nid is None:
1514 if self.root is None:
1515 new_tree_root = new_tree[new_tree.root]
1516 self.add_node(new_tree_root)
1517 nid = new_tree.root
1518 else:
1519 raise ValueError('Must define "nid" under which new tree is merged.')
1520 for child in new_tree.children(new_tree.root):
1521 self.paste(nid=nid, new_tree=new_tree.subtree(child.identifier), deep=deep)
1523 def paste(self, nid: str, new_tree, deep: bool = False):
1524 """
1525 Paste a @new_tree to the original one by linking the root
1526 of new tree to given node (nid).
1528 Update: add @deep copy of pasted tree.
1529 """
1530 assert isinstance(new_tree, Tree)
1532 if new_tree.root is None:
1533 return
1535 if nid is None:
1536 raise ValueError('Must define "nid" under which new tree is pasted.')
1538 if not self.contains(nid):
1539 raise NodeIDAbsentError("Node '%s' is not in the tree" % nid)
1541 set_joint = set(new_tree._nodes) & set(self._nodes) # joint keys
1542 if set_joint:
1543 raise ValueError("Duplicated nodes %s exists." % list(map(text, set_joint)))
1545 for cid, node in iteritems(new_tree.nodes):
1546 if deep:
1547 node = deepcopy(new_tree[node])
1548 self._nodes.update({cid: node})
1549 node.clone_pointers(new_tree.identifier, self._identifier)
1551 self.__update_bpointer(new_tree.root, nid)
1552 self.__update_fpointer(nid, new_tree.root, self.node_class.ADD)
1554 def paths_to_leaves(self) -> StrLList:
1555 """
1556 Use this function to get the identifiers allowing to go from the root
1557 nodes to each leaf.
1559 :return: a list of list of identifiers, root being not omitted.
1561 For example:
1563 .. code-block:: python
1565 Harry
1566 |___ Bill
1567 |___ Jane
1568 | |___ Diane
1569 | |___ George
1570 | |___ Jill
1571 | |___ Mary
1572 | |___ Mark
1574 Expected result:
1576 .. code-block:: python
1578 [['harry', 'jane', 'diane', 'mary'],
1579 ['harry', 'jane', 'mark'],
1580 ['harry', 'jane', 'diane', 'george', 'jill'],
1581 ['harry', 'bill']]
1583 """
1584 res = []
1586 for leaf in self.leaves():
1587 res.append([nid for nid in self.rsearch(leaf.identifier)][::-1])
1589 return res
1591 def remove_node(self, identifier: str) -> int:
1592 """Remove a node indicated by 'identifier' with all its successors.
1593 Return the number of removed nodes.
1594 """
1595 if not self.contains(identifier):
1596 raise NodeIDAbsentError("Node '%s' " "is not in the tree" % identifier)
1598 parent = self[identifier].predecessor(self._identifier)
1600 # Remove node and its children
1601 removed = list(self.expand_tree(identifier))
1603 for id_ in removed:
1604 if id_ == self.root:
1605 self.root = None
1606 self.__update_bpointer(id_, None)
1607 for cid in self[id_].successors(self._identifier) or []:
1608 self.__update_fpointer(id_, cid, self.node_class.DELETE)
1610 # Update parent info
1611 self.__update_fpointer(parent, identifier, self.node_class.DELETE)
1612 self.__update_bpointer(identifier, None)
1614 for id_ in removed:
1615 self.nodes.pop(id_)
1616 return len(removed)
1618 def remove_subtree(self, nid: str, identifier: Optional[str] = None):
1619 """
1620 Get a subtree with ``nid`` being the root. If nid is None, an
1621 empty tree is returned.
1623 For the original tree, this method is similar to
1624 `remove_node(self,nid)`, because given node and its children
1625 are removed from the original tree in both methods.
1626 For the returned value and performance, these two methods are
1627 different:
1629 * `remove_node` returns the number of deleted nodes;
1630 * `remove_subtree` returns a subtree of deleted nodes;
1632 You are always suggested to use `remove_node` if your only to
1633 delete nodes from a tree, as the other one need memory
1634 allocation to store the new tree.
1636 :return: a :class:`Tree` object.
1637 """
1638 st = self._clone(identifier)
1639 if nid is None:
1640 return st
1642 if not self.contains(nid):
1643 raise NodeIDAbsentError("Node '%s' is not in the tree" % nid)
1644 st.root = nid
1646 # in original tree, the removed nid will be unreferenced from its
1647 # parents children
1648 parent = self[nid].predecessor(self._identifier)
1650 removed = list(self.expand_tree(nid))
1651 for id_ in removed:
1652 if id_ == self.root:
1653 self.root = None
1654 st._nodes.update({id_: self._nodes.pop(id_)})
1655 st[id_].clone_pointers(cast(str, self._identifier), cast(str, st.identifier))
1656 st[id_].reset_pointers(self._identifier)
1657 if id_ == nid:
1658 st[id_].set_predecessor(None, st.identifier)
1659 self.__update_fpointer(parent, nid, self.node_class.DELETE)
1660 return st
1662 def rsearch(self, nid: str, filter: Optional[Callable[[Node], bool]] = None):
1663 """
1664 Traverse the tree branch along the branch from nid to its
1665 ancestors (until root).
1667 :param filter: the function of one variable to act on the :class:`Node` object.
1668 """
1669 if nid is None:
1670 return
1672 if not self.contains(nid):
1673 raise NodeIDAbsentError("Node '%s' is not in the tree" % nid)
1675 filter = (lambda x: True) if (filter is None) else filter
1677 current = nid
1678 while current is not None:
1679 if filter(self[current]):
1680 yield current
1681 # subtree() hasn't update the bpointer
1682 current = self[current].predecessor(self._identifier) if self.root != current else None
1684 def save2file(
1685 self,
1686 filename: str,
1687 nid: Optional[str] = None,
1688 level: int = ROOT,
1689 idhidden: bool = True,
1690 filter: Optional[Callable[[Node], bool]] = None,
1691 key: Optional[Callable[[Node], Node]] = None,
1692 reverse: bool = False,
1693 line_type: str = "ascii-ex",
1694 data_property: Optional[str] = None,
1695 sorting: bool = True,
1696 ):
1697 """
1698 Save the tree into file for offline analysis.
1699 """
1701 def _write_line(line, f):
1702 f.write(line + b"\n")
1704 def handler(x):
1705 return _write_line(x, open(filename, "ab"))
1707 self.__print_backend(
1708 nid,
1709 level,
1710 idhidden,
1711 filter,
1712 key,
1713 reverse,
1714 line_type,
1715 data_property,
1716 sorting,
1717 func=handler,
1718 )
1720 def show(
1721 self,
1722 nid: Optional[str] = None,
1723 level: int = ROOT,
1724 idhidden: bool = True,
1725 filter: Optional[Callable[[Node], bool]] = None,
1726 key: Optional[Callable[[Node], Node]] = None,
1727 reverse: bool = False,
1728 line_type: str = "ascii-ex",
1729 data_property: Optional[str] = None,
1730 stdout: bool = True,
1731 sorting: bool = True,
1732 ):
1733 """
1734 Display the tree structure in a beautiful hierarchical format.
1736 This method provides rich visualization options for tree structures, allowing
1737 customization of appearance, content filtering, and display properties. Perfect
1738 for debugging, documentation, and user interfaces.
1740 Args:
1741 nid: Starting node identifier. If None, starts from root.
1742 level: Starting level (0 for root). Mainly for internal use.
1743 idhidden: If True, hide node identifiers in output (cleaner display).
1744 If False, show both tag and identifier as "tag[id]".
1745 filter: Function to selectively display nodes. Takes Node object,
1746 returns bool. Filtered nodes and their subtrees are hidden.
1747 key: Sorting function for nodes at each level. Takes Node object,
1748 returns comparison key. If None and sorting=True, sorts by Node.
1749 reverse: If True, reverse the sorting order at each level.
1750 line_type: Visual style for tree connections. Options:
1751 'ascii' - Simple |-- style
1752 'ascii-ex' - Unicode ├── style (default)
1753 'ascii-exr' - Unicode with rounded corners
1754 'ascii-em' - Double-line Unicode ╠══ style
1755 'ascii-emv' - Mixed vertical Unicode style
1756 'ascii-emh' - Mixed horizontal Unicode style
1757 data_property: If specified, display this property from node.data
1758 instead of node.tag. Useful for rich data objects.
1759 stdout: If True, print to console. If False, return as string.
1760 sorting: If True, sort nodes at each level. If False, preserve
1761 insertion order (key and reverse are ignored).
1763 Returns:
1764 str or None: If stdout=False, returns formatted string.
1765 If stdout=True, prints to console and returns None.
1767 Example:
1768 Different display styles::
1770 tree = Tree()
1771 tree.create_node("Root", "root")
1772 tree.create_node("Child A", "a", parent="root")
1773 tree.create_node("Child B", "b", parent="root")
1775 # Basic display
1776 tree.show()
1778 # Fancy style with IDs
1779 tree.show(idhidden=False, line_type="ascii-em")
1781 # Filtered display (only nodes starting with 'C')
1782 tree.show(filter=lambda x: x.tag.startswith('C'))
1784 # Sorted by tag, reversed
1785 tree.show(key=lambda x: x.tag, reverse=True)
1787 # Show custom data property
1788 tree.show(data_property="name") # If node.data.name exists
1790 # Return as string instead of printing
1791 tree_str = tree.show(stdout=False)
1793 Output styles::
1795 ascii: Root
1796 |-- Child A
1797 |-- Child B
1799 ascii-ex: Root
1800 ├── Child A
1801 └── Child B
1803 ascii-em: Root
1804 ╠══ Child A
1805 ╚══ Child B
1807 Note:
1808 The tree display automatically handles complex hierarchies with proper
1809 indentation and connection lines. Very large trees may be truncated
1810 for performance reasons.
1811 """
1812 self._reader = ""
1814 def write(line):
1815 self._reader += line.decode("utf-8") + "\n"
1817 try:
1818 self.__print_backend(
1819 nid,
1820 level,
1821 idhidden,
1822 filter,
1823 key,
1824 reverse,
1825 line_type,
1826 data_property,
1827 sorting,
1828 func=write,
1829 )
1830 except NodeIDAbsentError:
1831 print("Tree is empty")
1833 if stdout:
1834 print(self._reader)
1835 else:
1836 return self._reader
1838 def siblings(self, nid: str) -> NodeList:
1839 """
1840 Return the siblings of given @nid.
1842 If @nid is root or there are no siblings, an empty list is returned.
1843 """
1844 siblings = []
1846 if nid != self.root:
1847 pid = self[nid].predecessor(self._identifier)
1848 siblings = [self[i] for i in self[pid].successors(self._identifier) if i != nid]
1850 return siblings
1852 def size(self, level: Optional[int] = None) -> int:
1853 """
1854 Get the number of nodes of the whole tree if @level is not
1855 given. Otherwise, the total number of nodes at specific level
1856 is returned.
1858 @param level The level number in the tree. It must be between
1859 [0, tree.depth].
1861 Otherwise, InvalidLevelNumber exception will be raised.
1862 """
1863 if level is None:
1864 return len(self._nodes)
1865 else:
1866 try:
1867 level = int(level)
1868 return len([node for node in self.all_nodes_itr() if self.level(node.identifier) == level])
1869 except Exception:
1870 raise TypeError("level should be an integer instead of '%s'" % type(level))
1872 def subtree(self, nid: str, identifier: Optional[str] = None):
1873 """
1874 Return a shallow COPY of subtree with nid being the new root.
1875 If nid is None, return an empty tree.
1876 If you are looking for a deepcopy, please create a new tree
1877 with this shallow copy, e.g.,
1879 .. code-block:: python
1881 new_tree = Tree(t.subtree(t.root), deep=True)
1883 This line creates a deep copy of the entire tree.
1884 """
1885 st = self._clone(identifier)
1886 if nid is None:
1887 return st
1889 if not self.contains(nid):
1890 raise NodeIDAbsentError("Node '%s' is not in the tree" % nid)
1892 st.root = nid
1893 for node_n in self.expand_tree(nid):
1894 st._nodes.update({self[node_n].identifier: self[node_n]})
1895 # define nodes parent/children in this tree
1896 # all pointers are the same as copied tree, except the root
1897 st[node_n].clone_pointers(cast(str, self._identifier), cast(str, st.identifier))
1898 if node_n == nid:
1899 # reset root parent for the new tree
1900 st[node_n].set_predecessor(None, st.identifier)
1901 return st
1903 def update_node(self, nid: str, **attrs) -> None:
1904 """
1905 Update node's attributes.
1907 :param nid: the identifier of modified node
1908 :param attrs: attribute pairs recognized by Node object
1909 :return: None
1910 """
1911 cn = self[nid]
1912 for attr, val in iteritems(attrs):
1913 if attr == "identifier":
1914 # Updating node id meets following contraints:
1915 # * Update node identifier property
1916 # * Update parent's followers
1917 # * Update children's parents
1918 # * Update tree registration of var _nodes
1919 # * Update tree root if necessary
1920 cn = self._nodes.pop(nid)
1921 setattr(cn, "identifier", val)
1922 self._nodes[val] = cn
1924 if cn.predecessor(self._identifier) is not None:
1925 self[cn.predecessor(self._identifier)].update_successors(
1926 nid,
1927 mode=self.node_class.REPLACE,
1928 replace=val,
1929 tree_id=self._identifier,
1930 )
1932 for fp in cn.successors(self._identifier):
1933 self[fp].set_predecessor(val, self._identifier)
1935 if self.root == nid:
1936 self.root = val
1937 else:
1938 setattr(cn, attr, val)
1940 def to_dict(self, nid=None, key=None, sort=True, reverse=False, with_data=False):
1941 """Transform the whole tree into a dict."""
1943 nid = self.root if (nid is None) else nid
1944 ntag = self[nid].tag
1945 tree_dict = {ntag: {"children": []}}
1946 if with_data:
1947 tree_dict[ntag]["data"] = self[nid].data
1949 if self[nid].expanded:
1950 queue = [self[i] for i in self[nid].successors(self._identifier)]
1951 key = (lambda x: x) if (key is None) else key
1952 if sort:
1953 queue.sort(key=key, reverse=reverse)
1955 for elem in queue:
1956 tree_dict[ntag]["children"].append(
1957 self.to_dict(elem.identifier, with_data=with_data, sort=sort, reverse=reverse)
1958 )
1959 if len(tree_dict[ntag]["children"]) == 0:
1960 tree_dict = self[nid].tag if not with_data else {ntag: {"data": self[nid].data}}
1961 return tree_dict
1963 def to_json(self, with_data: bool = False, sort: bool = True, reverse: bool = False):
1964 """
1965 Export the tree structure as a JSON string.
1967 Converts the entire tree into a JSON representation that can be easily
1968 saved, transmitted, or reconstructed. Perfect for data persistence,
1969 API responses, and cross-system integration.
1971 Args:
1972 with_data: If True, include node.data in the JSON output.
1973 If False, only include tree structure and tags.
1974 Default False for smaller output size.
1975 sort: If True, sort children at each level by node comparison.
1976 If False, preserve insertion order.
1977 reverse: If True, reverse the sorting order at each level.
1978 Only applies when sort=True.
1980 Returns:
1981 str: JSON string representation of the tree.
1983 Example:
1984 Basic JSON export::
1986 tree = Tree()
1987 tree.create_node("Root", "root")
1988 tree.create_node("Child A", "a", parent="root")
1989 tree.create_node("Child B", "b", parent="root")
1991 # Basic structure only
1992 json_str = tree.to_json()
1993 print(json_str)
1994 # Output: {"Root": {"children": [{"Child A": {}}, {"Child B": {}}]}}
1996 # Include node data
1997 tree.create_node("Data Node", "data", parent="root",
1998 data={"type": "important", "value": 42})
1999 json_with_data = tree.to_json(with_data=True)
2000 print(json_with_data)
2001 # Output includes: "data": {"type": "important", "value": 42}
2003 # Sorted output
2004 sorted_json = tree.to_json(sort=True, reverse=True)
2006 JSON Structure::
2008 The output follows this hierarchical format:
2010 {
2011 "root_tag": {
2012 "children": [
2013 {
2014 "child1_tag": {
2015 "children": [...],
2016 "data": {...} // if with_data=True
2017 }
2018 },
2019 {
2020 "child2_tag": {} // leaf node
2021 }
2022 ],
2023 "data": {...} // if with_data=True
2024 }
2025 }
2027 Usage with APIs::
2029 # Save to file
2030 with open('tree.json', 'w') as f:
2031 f.write(tree.to_json(with_data=True))
2033 # Send via HTTP
2034 import requests
2035 response = requests.post('/api/trees',
2036 json={'tree': tree.to_json()})
2038 # Pretty print
2039 import json
2040 formatted = json.dumps(json.loads(tree.to_json()), indent=2)
2041 print(formatted)
2043 Note:
2044 The JSON format preserves the complete tree structure and can be
2045 used to reconstruct equivalent trees. However, complex data objects
2046 in node.data must be JSON-serializable (no functions, custom classes
2047 without __dict__, etc.).
2048 """
2049 return json.dumps(self.to_dict(with_data=with_data, sort=sort, reverse=reverse))
2051 def to_graphviz(
2052 self,
2053 filename: Optional[str] = None,
2054 shape: str = "circle",
2055 graph: str = "digraph",
2056 filter=None,
2057 key=None,
2058 reverse: bool = False,
2059 sorting: bool = True,
2060 ):
2061 """Exports the tree in the dot format of the graphviz software"""
2062 nodes, connections = [], []
2063 if self.nodes:
2064 for n in self.expand_tree(
2065 mode=self.WIDTH,
2066 filter=filter,
2067 key=key,
2068 reverse=reverse,
2069 sorting=sorting,
2070 ):
2071 nid = self[n].identifier
2072 label = str(self[n].tag).translate(str.maketrans({'"': r"\""}))
2073 state = '"{0}" [label="{1}", shape={2}]'.format(nid, label, shape)
2074 nodes.append(state)
2076 for c in self.children(nid):
2077 cid = c.identifier
2078 edge = "->" if graph == "digraph" else "--"
2079 connections.append(('"{0}" ' + edge + ' "{1}"').format(nid, cid))
2081 # write nodes and connections to dot format
2082 is_plain_file = filename is not None
2083 if is_plain_file:
2084 f = codecs.open(cast(str, filename), "w", "utf-8")
2085 else:
2086 f = StringIO()
2088 f.write(graph + " tree {\n")
2089 for n in nodes:
2090 f.write("\t" + n + "\n")
2092 if len(connections) > 0:
2093 f.write("\n")
2095 for cns in connections:
2096 f.write("\t" + cns + "\n")
2098 f.write("}")
2100 if not is_plain_file:
2101 print(f.getvalue())
2103 f.close()
2105 @classmethod
2106 def from_map(cls, child_parent_dict, id_func=None, data_func=None):
2107 """
2108 takes a dict with child:parent, and form a tree
2109 """
2110 tree = Tree()
2111 if tree is None or tree.size() > 0:
2112 raise ValueError("need to pass in an empty tree")
2113 id_func = id_func if id_func else lambda x: x
2114 data_func = data_func if data_func else lambda x: None
2115 parent_child_dict = {}
2116 root_node = None
2117 for k, v in child_parent_dict.items():
2118 if v is None and root_node is None:
2119 root_node = k
2120 elif v is None and root_node is not None:
2121 raise ValueError("invalid input, more than 1 child has no parent")
2122 else:
2123 if v in parent_child_dict:
2124 parent_child_dict[v].append(k)
2125 else:
2126 parent_child_dict[v] = [k]
2127 if root_node is None:
2128 raise ValueError("cannot find root")
2130 tree.create_node(root_node, id_func(root_node), data=data_func(root_node))
2131 queue = [root_node]
2132 while len(queue) > 0:
2133 parent_node = queue.pop()
2134 for child in parent_child_dict.get(parent_node, []):
2135 tree.create_node(
2136 child,
2137 id_func(child),
2138 parent=id_func(parent_node),
2139 data=data_func(child),
2140 )
2141 queue.append(child)
2142 return tree