Add vprint

CHANGELOG.md

@@ -6,7 +6,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 ### Added:
-- Tree Exporter: `tree_to_pillow_graph` method to allow cmap for node background.
+- Tree Exporter: `tree_to_pillow_graph` method to allow cmap for node background and generic kwargs for yield_tree.
+- Tree Exporter: `vprint_tree` method to print trees vertically, able to use `.vshow()` as well.
 ### Changed:
 - Tree Exporter: Minor refactoring to deduplicate code in export module.
 

bigtree/__init__.py

@@ -39,6 +39,8 @@
     tree_to_pillow,
     tree_to_pillow_graph,
     tree_to_polars,
+    vprint_tree,
+    vyield_tree,
     yield_tree,
 )
 from bigtree.tree.helper import (
@@ -74,20 +76,34 @@
     findall,
 )
 from bigtree.utils.constants import (
+    ANSIBorderStyle,
     ANSIHPrintStyle,
     ANSIPrintStyle,
+    ANSIVPrintStyle,
+    ASCIIBorderStyle,
     ASCIIHPrintStyle,
     ASCIIPrintStyle,
+    ASCIIVPrintStyle,
     BaseHPrintStyle,
     BasePrintStyle,
+    BaseVPrintStyle,
+    BorderStyle,
+    ConstBoldBorderStyle,
     ConstBoldHPrintStyle,
     ConstBoldPrintStyle,
+    ConstBoldVPrintStyle,
+    ConstBorderStyle,
     ConstHPrintStyle,
     ConstPrintStyle,
+    ConstVPrintStyle,
+    DoubleBorderStyle,
     DoubleHPrintStyle,
     DoublePrintStyle,
+    DoubleVPrintStyle,
+    RoundedBorderStyle,
     RoundedHPrintStyle,
     RoundedPrintStyle,
+    RoundedVPrintStyle,
 )
 from bigtree.utils.groot import speak_like_groot, whoami
 from bigtree.utils.iterators import (

bigtree/node/node.py

@@ -74,6 +74,7 @@ class Node(basenode.BaseNode):
 
     1. ``show()``: Print tree to console
     2. ``hshow()``: Print tree in horizontal orientation to console
+    2. ``vshow()``: Print tree in vertical orientation to console
 
     ----
 
@@ -227,6 +228,12 @@ def hshow(self, **kwargs: Any) -> None:
 
         hprint_tree(self, **kwargs)
 
+    def vshow(self, **kwargs: Any) -> None:
+        """Print tree in vertical orientation to console, takes in same keyword arguments as `vprint_tree` function"""
+        from bigtree.tree.export import vprint_tree
+
+        vprint_tree(self, **kwargs)
+
     def __getitem__(self, child_name: str) -> "Node":
         """Get child by name identifier
 

bigtree/tree/export/__init__.py

@@ -11,6 +11,8 @@
     hyield_tree,
     print_tree,
     tree_to_newick,
+    vprint_tree,
+    vyield_tree,
     yield_tree,
 )
 
@@ -19,6 +21,8 @@
     "yield_tree",
     "hprint_tree",
     "hyield_tree",
+    "vprint_tree",
+    "vyield_tree",
     "tree_to_dataframe",
     "tree_to_polars",
     "tree_to_dict",

bigtree/tree/export/_stdout.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+from typing import List, Optional, TypeVar
+
+from bigtree.node import node
+from bigtree.utils.constants import BaseVPrintStyle, BorderStyle
+
+__all__ = [
+    "calculate_stem_pos",
+    "format_node",
+    "horizontal_join",
+]
+
+T = TypeVar("T", bound=node.Node)
+
+
+def calculate_stem_pos(length: int) -> int:
+    """Calculate stem position based on length
+
+    Args:
+        length (int): length of node
+
+    Returns:
+        (int) Stem position
+    """
+    if length % 2:
+        return length // 2
+    return length // 2 - 1
+
+
+def format_node(
+    _node: T,
+    alias: str,
+    intermediate_node_name: bool = True,
+    style: BaseVPrintStyle = BaseVPrintStyle.from_style("const"),
+    border_style: Optional[BorderStyle] = None,
+) -> List[str]:
+    """Format node to be same width, able to customise whether to add border
+
+    Args:
+        _node (Node): node to format
+        alias (str): node attribute to use for node name in tree as alias to `node_name`, if present.
+            Otherwise, it will default to `node_name` of node.
+        intermediate_node_name (bool): indicator if intermediate nodes have node names, defaults to True
+        style (BaseVPrintStyle): style to format node, used only if border_style is None
+        border_style (BorderStyle): border style to format node
+
+    Returns:
+        (List[str]) node display
+    """
+    if not intermediate_node_name and _node.children:
+        if border_style is None:
+            node_title: str = style.BRANCH
+            if _node.is_root:
+                node_title = style.SUBSEQUENT_CHILD
+        else:
+            node_title = ""
+    else:
+        node_title = _node.get_attr(alias) or _node.node_name
+    node_title_lines = node_title.split("\n")
+    width = max([len(node_title_lines) for node_title_lines in node_title_lines])
+
+    node_display_lines: List[str] = []
+    if border_style:
+        width += 2
+        node_display_lines.append(
+            f"{border_style.TOP_LEFT}{border_style.HORIZONTAL * width}{border_style.TOP_RIGHT}"
+        )
+        for node_title_line in node_title_lines:
+            node_display_lines.append(
+                f"{border_style.VERTICAL} {node_title_line.center(width - 2)} {border_style.VERTICAL}"
+            )
+        node_display_lines.append(
+            f"{border_style.BOTTOM_LEFT}{border_style.HORIZONTAL * width}{border_style.BOTTOM_RIGHT}"
+        )
+        node_mid = calculate_stem_pos(width + 2)
+        # If there is a parent
+        if _node.parent:
+            node_display_lines[0] = (
+                node_display_lines[0][:node_mid]
+                + style.SPLIT_BRANCH
+                + node_display_lines[0][node_mid + 1 :]  # noqa
+            )
+        # If there are subsequent children
+        if any(_node.children):
+            node_display_lines[-1] = (
+                node_display_lines[-1][:node_mid]
+                + style.SUBSEQUENT_CHILD
+                + node_display_lines[-1][node_mid + 1 :]  # noqa
+            )
+    else:
+        for node_title_line in node_title_lines:
+            node_display_lines.append(f"{node_title_line.center(width)}")
+    return node_display_lines
+
+
+def horizontal_join(node_displays: List[List[str]], spacing: int = 2) -> List[str]:
+    """Horizontally join multiple node displays
+
+    Args:
+        node_displays (List[List[str]]): multiple node displays belonging to the same row
+        spacing (int): spacing between node displays
+
+    Returns:
+        (List[str]) node display of the row
+    """
+    space = " "
+    height = max([len(node_display) for node_display in node_displays])
+    row_display = {idx: "" for idx in range(height)}
+    for node_display_idx, node_display in enumerate(node_displays):
+        width = len(node_display[0])
+
+        # Add node content
+        for row_idx, node_display_row in enumerate(node_display):
+            if node_display_idx:
+                row_display[row_idx] += spacing * space
+            row_display[row_idx] += node_display_row
+
+        # Add node buffer
+        for row_idx_buffer in range(len(node_display), height):
+            if node_display_idx:
+                row_display[row_idx_buffer] += spacing * space
+            row_display[row_idx_buffer] += width * space
+    return [row_display[k] for k in sorted(row_display)]

bigtree/tree/export/images.py

@@ -253,10 +253,12 @@ def tree_to_pillow_graph(
     rect_cmap_attr: Optional[str] = None,
     rect_outline: Union[Tuple[int, int, int], str] = "black",
     rect_width: int = 1,
+    **kwargs: Any,
 ) -> Image.Image:
     r"""Export tree to PIL.Image.Image object. Object can be
     converted to other formats, such as jpg, or png. Image will look
-    like a tree/graph-like structure.
+    like a tree/graph-like structure, accepts additional keyword arguments
+    as input to `yield_tree`.
 
     Customisations:
 
@@ -338,7 +340,7 @@ def get_node_text(_node: T, _node_content: str) -> str:
         return _node_content
 
     cmap_range: Set[Union[float, int]] = set()
-    for _, _, _node in yield_tree(tree):
+    for _, _, _node in yield_tree(tree, **kwargs):
         l, t, r, b = _draw.multiline_textbbox(
             (0, 0), get_node_text(_node, node_content), font=font
         )
@@ -375,7 +377,7 @@ def get_node_text(_node: T, _node_content: str) -> str:
     _width, _height = 0, 0
     _width_margin = 0.5 * _max_text_width + margin.get("r", 0)
     _height_margin = 0.5 * _max_text_height + margin.get("b")
-    for _, _, _node in yield_tree(tree):
+    for _, _, _node in yield_tree(tree, **kwargs):
         _width = max(_width, _node.get_attr("x") + _width_margin)
         _height = max(_height, _node.get_attr("y") + _height_margin)
     _width = int(round(_width + 0.5, 0))
@@ -385,7 +387,7 @@ def get_node_text(_node: T, _node_content: str) -> str:
     image = Image.new("RGB", (_width, _height), bg_colour)
     image_draw = ImageDraw.Draw(image)
 
-    for _, _, _node in yield_tree(tree):
+    for _, _, _node in yield_tree(tree, **kwargs):
         _x, _y = _node.get_attr("x"), _node.get_attr("y")
         x1, x2 = _x - 0.5 * _max_text_width, _x + 0.5 * _max_text_width
         y1, y2 = _y - 0.5 * _max_text_height, _y + 0.5 * _max_text_height