使用anytree可视化具有不同颜色节点的树答案

【问题标题】：Visualizing a tree with nodes of different colors with anytree使用anytree可视化具有不同颜色节点的树
【发布时间】：2022-02-15 06:59:40
【问题描述】：

我正在寻找一种简单的方法来构造和绘制一棵树（在 google colab 上）。

重要的是，我想要不同颜色和形状的节点。理想情况下，我想要以下内容。

from anytree import Node, RenderTree
from anytree.exporter import DotExporter
from IPython.display import Image


# construct tree
ceo = Node("CEO") #root

vp_1 = Node("VP_1", parent=ceo, color="red")
vp_2 = Node("VP_2", parent=ceo)

gm_1 = Node("GM_1", parent=vp_1, shape="square", color="red")
gm_2 = Node("GM_2", parent=vp_2, shape="square")

m_1 = Node("M_1", parent=gm_2)

# draw tree
DotExporter(ceo).to_picture("ceo.png")

# show image
Image('ceo.png')

由于color 和shape 不是Node 的真正参数，因此此代码当前生成以下图像。我希望 VP_1 和 GM_1 为红色，GM_1 和 GM_2 为正方形。

非常感谢您的帮助！

【问题讨论】：

anytree 不支持根据我在文档中看到的内容进行整形和重新着色
我可以使用任何树以外的东西。
然后stackoverflow.com/questions/2442014/tree-libraries-in-python 可能是搜索支持不同可视化格式的库的良好起点。
谢谢！我最终使用了graphviz：graphviz.readthedocs.io/en/stable/examples.html

标签： python google-colaboratory anytree

【解决方案1】：

Anytree 的DotExporter 有一个nodeattrfunc 参数，您可以在其中传递一个函数，该函数接受Node 并返回它应该在DOT 语言字符串中给出的属性。因此，如果您已将颜色和形状信息存储为节点属性，则可以使用自定义 nodeattrfunc 如下所示将它们转换为 DOT 属性：

def set_color_shape(node):
    attrs = []
    attrs += [f'color={node.color}'] if hasattr(node, 'color') else []
    attrs += [f'shape={node.shape}'] if hasattr(node, 'shape') else []
    return ', '.join(attrs)

DotExporter(ceo, nodeattrfunc=set_color_shape).to_picture('example.png')

结果见：https://i.stack.imgur.com/lhPqT.png

【讨论】：