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

Question

我正在寻找一种构造和绘制树的简单方法（在 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 是正方形的。

非常感谢您的帮助！

Answer 1

Anytree 的 DotExporter 有一个 nodeattrfunc 参数，您可以在其中传递接受 Node 和 returns 应该以 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

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

Visualizing a tree with nodes of different colors with anytree

python

anytree

google-colaboratory