Convert sorted strings to ASCII art tree, merging common prefixes

prefixtree.py

#!/usr/bin/env -S uv run --script
# /// script
# dependencies = []
# ///

import argparse
import sys


def build_tree(items, delim=None):
    """Build a tree structure from sorted strings.

    Args:
        items: List of strings to build tree from
        delim: Optional delimiter(s) to split on. If None, splits on each character.
    """
    tree = {}

    for item in items:
        # Split based on delimiter or use each character
        if delim:
            parts = []
            current = ""
            for char in item:
                if char in delim:
                    if current:
                        parts.append(current)
                    parts.append(char)
                    current = ""
                else:
                    current += char
            if current:
                parts.append(current)
        else:
            parts = list(item)

        # Navigate/build tree structure
        node = tree
        for part in parts:
            if part not in node:
                node[part] = {}
            node = node[part]

    return tree


def merge_single_branches(tree):
    """Merge branches with single children into combined labels."""
    merged = {}

    for key, subtree in tree.items():
        # Recursively merge subtree first
        merged_subtree = merge_single_branches(subtree)

        # If subtree has exactly one child, merge it
        current_key = key
        current_subtree = merged_subtree

        while len(current_subtree) == 1:
            child_key, child_subtree = next(iter(current_subtree.items()))
            current_key += child_key
            current_subtree = child_subtree

        merged[current_key] = current_subtree

    return merged


def print_tree(tree, prefix="", is_last=True, parent_prefix=""):
    """Print tree as ASCII art."""
    items = sorted(tree.items())

    for i, (key, subtree) in enumerate(items):
        is_last_item = i == len(items) - 1

        # Determine the connector and continuation
        connector = "└── " if is_last_item else "├── "
        extension = "    " if is_last_item else "│   "

        print(f"{parent_prefix}{connector}{key}")

        if subtree:
            new_prefix = parent_prefix + extension
            print_tree(subtree, prefix + extension, is_last_item, new_prefix)


def main():
    parser = argparse.ArgumentParser(
        description="Convert sorted strings to ASCII art tree, merging common prefixes"
    )
    parser.add_argument(
        "items", nargs="*", help="Items to convert to tree (if not provided, read from stdin)"
    )
    parser.add_argument("-q", "--quiet", action="store_true", help="Don't print input items")
    parser.add_argument(
        "-d",
        "--delim",
        type=str,
        default=None,
        help="Delimiter character(s) to split on (if not provided, splits on each character)",
    )

    args = parser.parse_args()

    # Get items from arguments or stdin
    if args.items:
        items = args.items
    else:
        items = [line.rstrip("\n") for line in sys.stdin]

    if not items:
        parser.print_help()
        sys.exit(1)

    tree = build_tree(items, delim=args.delim)
    tree = merge_single_branches(tree)
    print_tree(tree)


if __name__ == "__main__":
    main()