0

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

在此处输入图像描述

非常感谢您的帮助！

4

1 回答 1

1

AnytreeDotExporter有一个nodeattrfunc参数，您可以在其中传递一个函数，该函数接受 aNode并返回应该在 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

于 2022-02-14T22:59:40.540 回答