printing tables and trees

printers.py

from collections.abc import Callable, Collection, Iterable
from enum import Enum
from typing import TypeVar


class Alignment(Enum):
    LEFT = "<"
    CENTER = "^"
    RIGHT = ">"


T = TypeVar("T")


def print_table(
    rows: Iterable[Iterable[T]],
    *,
    headers: list[str],
    alignments: list[Alignment],
    format_cell: Callable[[T], str] = str,
    use_ascii_only: bool = False,
) -> None:
    V_BAR = "|" if use_ascii_only else "│"
    H_BAR = "-" if use_ascii_only else "─"
    CROSS = "|" if use_ascii_only else "┼"
    L_END = "|" if use_ascii_only else "├"
    R_END = "|" if use_ascii_only else "┤"

    # calculate column widths
    column_widths = [len(header) for header in headers]
    for row in rows:
        for i, cell in enumerate(row):
            cell_len = len(format_cell(cell))
            if i >= len(column_widths):
                column_widths.append(cell_len)
            else:
                column_widths[i] = max(cell_len, column_widths[i])

    # add default alignment/header for unspecified columns
    alignments += [Alignment.LEFT] * (len(column_widths) - len(alignments))
    headers += [""] * (len(column_widths) - len(headers))

    # print header row
    print(V_BAR, end="")
    for i, column_width in enumerate(column_widths):
        if i > 0:
            print(V_BAR, end="")
        print(f" {headers[i]:{column_width}} ", end="")
    print(V_BAR)

    # print separator
    print(L_END, end="")
    for i, column_width in enumerate(column_widths):
        if i > 0:
            print(CROSS, end="")
        print(H_BAR * (column_width + 2), end="")
    print(R_END)

    # print data rows
    for row in rows:
        print(V_BAR, end="")
        for i, cell_value in enumerate(row):
            if i > 0:
                print(V_BAR, end="")
            column_width = column_widths[i]
            alignment = alignments[i]
            cell_text = format_cell(cell_value)
            print(f" {cell_text:{alignment.value}{column_width}} ", end="")
        print(V_BAR)


def print_tree(
    root_node: T,
    get_children: Callable[[T], Collection[T] | None],
    *,
    format_node: Callable[[T], str] = str,
    use_ascii_only: bool = False,
) -> None:
    V_BAR = "|" if use_ascii_only else "│"
    H_BAR = "-" if use_ascii_only else "─"
    CHILD = "|" if use_ascii_only else "├"
    CHILD_LAST = "\\" if use_ascii_only else "└"
    ELLIPSIS = "..." if use_ascii_only else "⋯"

    def print_node(node: T, prefix: str) -> None:
        print(format_node(node))
        children = get_children(node)
        if children is None:
            print(f"{prefix}{CHILD_LAST}{H_BAR} {ELLIPSIS}")
            return

        for i, child in enumerate(children):
            last_child = i == len(children) - 1
            print(f"{prefix}{CHILD_LAST if last_child else CHILD}{H_BAR}", end="")

            inner_prefix = f"{prefix}{' ' if last_child else V_BAR} "
            print_node(child, inner_prefix)

    print_node(root_node, "")