Source code for masterpiece.treevisualizer
"""
TreeVisualizer, for rendering instance hierarchy.
Author: Juha Meskanen
Date: 2024-10-26
"""
from typing import Any
from colorama import Fore, Style, init
from .composite import Composite
from .masterpiece import MasterPiece
# Initialize Colorama
init(autoreset=True)
[docs]
class TreeVisualizer:
"""
The `TreeVisualizer` class is designed to visually represent hierarchical structures
using ASCII art and colors, making it easy to understand the relationships between
nodes in a tree. It is particularly useful for visualizing instances of the
`MasterPiece` and `Composite` classes, which represent elements in a tree structure.
Features:
---------
- Supports customizable colors for node representation using the Colorama library.
- Prints the hierarchy of nodes with clear visual indicators (├─, └─, and │) to represent parent-child relationships.
- Automatically resets colors after each print to maintain consistent output.
Usage:
------
To use the `TreeVisualizer`, first create an instance by specifying the desired color.
Then, call the `print_tree` method with the root node of your tree.
Example:
--------
.. code-block:: python
from masterpiece.core import TreeVisualizer, MasterPiece, Composite
# Create a sample hierarchy
parent = Composite("parent")
child1 = MasterPiece("child1")
child2 = MasterPiece("child2")
parent.add(child1)
parent.add(child2)
# Initialize the visualizer with a specified color
visualizer = TreeVisualizer("green")
# Print the hierarchy
visualizer.print_tree(parent)
"""
[docs]
def __init__(self, color: str):
"""Initialize the TreeVisualizer with a color.
Args:
color (str): The color to use for visualizing the tree.
"""
self.color = color
[docs]
def get_color(self) -> Any:
"""Return the corresponding Colorama color code for the specified color."""
colors = {
"red": Fore.RED,
"green": Fore.GREEN,
"yellow": Fore.YELLOW,
"blue": Fore.BLUE,
"magenta": Fore.MAGENTA,
"cyan": Fore.CYAN,
"white": Fore.WHITE,
"reset": Style.RESET_ALL,
}
return colors.get(
self.color.lower(), Fore.WHITE
) # Default to white if color not found
[docs]
def print_tree(
self, node: MasterPiece, prefix: str = "", is_last: bool = True
) -> None:
"""Print the hierarchy of the node using ├─, └─, and │ with the specified color.
Args:
node (MasterPiece): The root node to print.
prefix (str, optional): The prefix for the current level.
is_last (bool, optional): Whether this node is the last child.
"""
# Get the color code for the current node
color_code = self.get_color()
# Print the current node's name with the color
print(prefix, end="")
if prefix:
print(f"{color_code}└─ " if is_last else f"{color_code}├─ ", end="")
print(f"{color_code}{node.name}{Style.RESET_ALL}") # Reset color after printing
# Prepare the new prefix for the child nodes
if prefix:
prefix += " " if is_last else f"{color_code}│ "
else:
prefix = " " if is_last else f"{color_code}│ "
# Recursively print all the children
if isinstance(node, Composite):
for i, child in enumerate(node.children):
is_last_child = i == len(node.children) - 1
self.print_tree(child, prefix, is_last_child)