求助：Python中仅用根节点遍历自定义树结构生成.dot文件的工具

你的自定义树结构因为字段命名和组织方式特殊，现成工具很难直接适配，最靠谱的方式是自己写几行递归代码生成.dot文件，比找第三方工具更灵活可控。

1. 手动实现递归生成.dot内容 #

直接写遍历函数，从根节点出发生成Graphviz格式文本，再保存为.dot文件：

class node:
    def __init__(self, attr=None, label=None):
        self.attr = attr
        self.next_attr_vals = {}
        self.label = label

def generate_dot(root):
    dot_content = ["digraph DecisionTree {", "    node [shape=box];"]
    visited = set()

    def traverse(current_node, parent_id=None, edge_label=None):
        if not current_node or id(current_node) in visited:
            return
        visited.add(id(current_node))
        # 拼接节点标签，空值自动省略
        node_label = ""
        if current_node.attr:
            node_label += f"Attr: {current_node.attr}"
        if current_node.label:
            node_label += f"\\nLabel: {current_node.label}"
        dot_content.append(f"    {id(current_node)} [label=\"{node_label}\"];")
        # 添加父节点到当前节点的边
        if parent_id is not None and edge_label is not None:
            dot_content.append(f"    {parent_id} -> {id(current_node)} [label=\"{edge_label}\"];")
        # 递归遍历所有子节点
        for key, child in current_node.next_attr_vals.items():
            traverse(child, id(current_node), key)

    traverse(root)
    dot_content.append("}")
    return "\n".join(dot_content)

# 示例用法
if __name__ == "__main__":
    # 构建测试树
    root = node(attr="Outlook")
    root.next_attr_vals["Sunny"] = node(attr="Humidity")
    root.next_attr_vals["Sunny"].next_attr_vals["High"] = node(label="No")
    root.next_attr_vals["Sunny"].next_attr_vals["Normal"] = node(label="Yes")
    root.next_attr_vals["Rain"] = node(attr="Wind")
    root.next_attr_vals["Rain"].next_attr_vals["Strong"] = node(label="No")
    root.next_attr_vals["Rain"].next_attr_vals["Weak"] = node(label="Yes")

    # 生成并保存dot文件
    dot_text = generate_dot(root)
    with open("tree.dot", "w") as f:
        f.write(dot_text)

2. 结合graphviz库生成（更简洁） #

如果不想手动拼接字符串，可以用graphviz库（先通过pip install graphviz安装），同样通过递归遍历根节点构建图：

from graphviz import Digraph

class node:
    def __init__(self, attr=None, label=None):
        self.attr = attr
        self.next_attr_vals = {}
        self.label = label

def build_graph(root):
    dot = Digraph(node_attr={'shape': 'box'})
    visited = set()

    def traverse(current_node, parent_id=None, edge_label=None):
        if not current_node or id(current_node) in visited:
            return
        visited.add(id(current_node))
        # 生成节点标签
        node_label = ""
        if current_node.attr:
            node_label += f"Attr: {current_node.attr}"
        if current_node.label:
            node_label += f"\nLabel: {current_node.label}"
        dot.node(str(id(current_node)), label=node_label)
        # 添加边
        if parent_id is not None and edge_label is not None:
            dot.edge(str(parent_id), str(id(current_node)), label=edge_label)
        # 递归处理子节点
        for key, child in current_node.next_attr_vals.items():
            traverse(child, id(current_node), key)

    traverse(root)
    return dot

# 示例用法
if __name__ == "__main__":
    # 构建测试树（同上述例子）
    root = node(attr="Outlook")
    root.next_attr_vals["Sunny"] = node(attr="Humidity")
    root.next_attr_vals["Sunny"].next_attr_vals["High"] = node(label="No")
    root.next_attr_vals["Sunny"].next_attr_vals["Normal"] = node(label="Yes")
    root.next_attr_vals["Rain"] = node(attr="Wind")
    root.next_attr_vals["Rain"].next_attr_vals["Strong"] = node(label="No")
    root.next_attr_vals["Rain"].next_attr_vals["Weak"] = node(label="Yes")

    graph = build_graph(root)
    graph.render("tree", format="dot")  # 生成tree.dot文件
    graph.view()  # 直接打开可视化窗口（可选）

这两种方式都只需要传入根节点就能完成遍历和.dot文件生成，完全适配你的自定义节点结构，避免了现成工具因结构不兼容导致的报错问题。

内容的提问来源于stack exchange，提问作者RI98