Generate binary tree dot language codes by level order. Make binary tree visualization easier.

bitviz.py

#!/usr/bin/env python
# https://blog.nanpuyue.com/2019/054.html
# I use ChatGPT to translate the original Rust code into Python and made some modifications.

from typing import List, Optional, Union
from pathlib import Path
import os


class TNode:
    def __init__(self, val: Union[str, int], left: Optional["TNode"] = None, right: Optional["TNode"] = None):
        self.val = val
        self.left = left
        self.right = right


def create_tree_from_level_order(vals: List[Union[str, int]]) -> Optional[TNode]:
    if not vals:
        return None

    nodes = [None if val == "#" else TNode(val) for val in vals]
    children = nodes[::-1]

    root = children.pop()
    for node in nodes:
        if node:
            if children:
                node.left = children.pop()
            if children:
                node.right = children.pop()

    return root


def tree2dot(tree: Optional[TNode]) -> str:
    target = [None]
    distance = [0]
    res = []

    def print_node(node: TNode) -> None:
        nonlocal target, distance

        if node.left:
            left_max = node.left
            left_distance = 1
            while left_max.right:
                left_max = left_max.right
                left_distance += 1
            target[0] = left_max.val
            distance[0] = left_distance

            if node.left.left or node.left.right:
                res.append(f"    {node.left.val} [group={node.left.val}]")
            res.append(f"    {node.val} -> {node.left.val}")
            print_node(node.left)

        if node.left or node.right:
            res.append(f'    _{node.val} [group={node.val}, label="", width=0, style=invis]')
            res.append(f"    {node.val} -> _{node.val} [style=invis]")

        if node.right:
            right_min = node.right
            right_distance = 1
            while right_min.left:
                right_min = right_min.left
                right_distance += 1
            if right_distance <= distance[0]:
                target[0] = right_min.val
                distance[0] = right_distance

            if node.right.left or node.right.right:
                res.append(f"    {node.right.val} [group={node.right.val}]")
            res.append(f"    {node.val} -> {node.right.val}")
            print_node(node.right)

        if distance[0] > 1 and target[0]:
            res.append(f"    {{rank=same; _{node.val}; {target[0]}}}")

    if tree:
        res.append("digraph G {")
        res.append("    graph [nodesep=0.1]")
        res.append("    node [shape=circle, fixedsize=true, width=0.5]")
        res.append("    edge [arrowhead=none]")
        res.append(f"    {tree.val} [group={tree.val}]")
        print_node(tree)
        res.append("}")
    return "\n".join(res)

def read_input() -> str:
    filename = input("Enter the file name: ")
    content = ""
    while True:
        try:
            line = input()
            content += "\n" + line
        except EOFError:
            break
    return filename, content.split()


def create_dot_file(filename: str, content: str) -> None:
    level_order_vals = list(content)
    tree = create_tree_from_level_order(level_order_vals)
    dot_filename = f"{filename}.dot"
    svg_filename = f"{filename}.svg"
    current_dir = Path(os.getcwd())  # Get the current working directory
    dot_filepath = current_dir / dot_filename
    svg_filepath = current_dir / svg_filename
    dot = tree2dot(tree)
    with open(dot_filepath, "w") as f:
        f.write(dot)
    print(f"DOT file '{dot_filepath}' created.")
    os.system(f"dot -Tsvg '{dot_filepath}' -o '{svg_filepath}'")
    print(f"SVG file '{svg_filepath}' created.")


if __name__ == "__main__":
    filename, content = read_input()
    print(filename, content)
    create_dot_file(filename, content)