tree_to_mermaid.py

import sys
import re
from anytree import Node, RenderTree

MAX_NODES = 500

def parse_tree_output(lines):
    root = Node("root")
    stack = [root]
    prev_indent = -1
    node_count = 0
    
    line_pattern = re.compile(r'^([│├└─\s]*)(.*)')

    for line in lines:
        if node_count >= MAX_NODES:
            print(f"Warning: Node limit of {MAX_NODES} reached. Some nodes may be omitted.", file=sys.stderr)
            break

        line = line.rstrip()
        if not line or line.startswith('pomoria'):
            continue

        match = line_pattern.match(line)
        if not match:
            continue

        indent_str, name = match.groups()
        indent = len(indent_str)
        name = name.lstrip()

        if indent > prev_indent:
            parent = stack[-1]
        elif indent == prev_indent:
            stack.pop()
            parent = stack[-1]
        else:
            while indent < prev_indent and len(stack) > 1:
                stack.pop()
                prev_indent = len(stack) * 2 - 2
            parent = stack[-1]

        node = Node(name, parent=parent)
        stack.append(node)
        prev_indent = indent
        node_count += 1

    return root if root.children else None

def generate_mermaid_tree(root):
    if not root:
        return "graph TD\n    id0[\"No data\"]"

    mermaid_code = ["graph TD"]
    node_id = 0
    node_ids = {}

    def add_node(node, parent_id=None):
        nonlocal node_id
        if node_id >= MAX_NODES:
            return

        current_id = f"id{node_id}"
        node_ids[node] = current_id
        label = node.name.replace('"', '\\"')  # Escape double quotes in labels
        mermaid_code.append(f'    {current_id}["{label}"]')
        
        if parent_id:
            mermaid_code.append(f'    {parent_id} --- {current_id}')
        
        node_id += 1

        for child in node.children:
            add_node(child, current_id)

    for child in root.children:
        add_node(child)

    if node_id >= MAX_NODES:
        mermaid_code.append(f'    id{MAX_NODES}["... (Node limit reached)"]')

    return "\n".join(mermaid_code)

def main():
    tree_output = sys.stdin.readlines()
    root = parse_tree_output(tree_output)
    mermaid_code = generate_mermaid_tree(root)
    print(mermaid_code)

if __name__ == "__main__":
    main()