Python Graphviz 的使用-绘制树形图-菜鸟笔记

摘要本文会介绍使用 Graphviz 来绘制树形图。同时介绍使用 Python 中的 graphviz 来简单 Graphviz 的使用。

简介

Graphviz 是一个绘图工具，可以根据 dot 脚本画出树形图等。本文会介绍 graphviz 的简单使用，使用 graphviz 来创建 dot 文件来绘制树形图。本文还会通过几个 graphviz 的例子，来介绍 graphviz 的使用。

参考资料

Graphviz 的中文文档，Graphviz 中文文档
Dot 语言介绍（想要更多了解的时候可以查看），Graphviz Dot 语言介绍
Python Graphviz 的入门文档，Graphviz User Guide
Python Graphviz 的 Example 介绍，Graphviz 的例子
Graphviz 的下载，Graphviz Download
使用 Graphviz 的例子，生成项目uml框架图-pyreverse介绍

Graphviz 的简单介绍

Graphviz 编译的源文件是 dot 文件。我们通过书写 dot 文件，来获得不同的树状图。首先需要安装 Graphviz，可以通过下面的链接进行安装，Graphviz Download。关于 Graphviz 的说明，可以参考这个链接，有详细的例子解释，Graphviz 中文文档。

绘制简单树形图

我们将如下代码保存为 graph01.dot 文件：

digraph G {
main -> parse -> execute;
main -> init;
main -> cleanup;
execute -> make_string;
execute -> printf;
execute -> compare;
}

接着使用命令 dot -Tjpg ./awesome_project/graph.dot -o graph01.jpg 编译，可以得到名为 graph01.jpg 的图片，图片内容如下所示。

绘制更加复杂的树形图-node 和 edge 的样式

上面是一个简单的树状图的绘制，下面我们来看一个稍微复杂的情况，有以下的几个需求：

给 node 和 edge 添加格式，颜色，加粗，是否虚线；
在箭头上添加文字，例如使用 a->b [style="dashed", color="skyblue"]；
支持中文，在 node 上指定 fontname 即可，fontname="Microsoft Yahei"
指定箭头的方向，这里在连接部分使用 dir=none, dir=both, dir=forward, dir=back；

我们直接看一下最终的代码：

digraph G {
main [shape=box];
main -> parse [weight=8];
parse-> execute;
main -> init [style=dotted, dir = none];
main -> cleanup [dir = both];
execute -> {make_string, printf};
init -> make_string;
edge [color=red];
main -> printf [style=bold, label="100 times"];
make_string [label = "make a\nstring"];
node [shape=box, style=filled,color=".7, .3, 1.0", fontname="Microsoft Yahei"];
execute -> 比较;
}

最终的效果如下所示：

定义子图和子图的样式

graphviz 支持子图，即图中的部分节点和边相对对立（软件的模块划分经常如此）。比如在下面的例子中，我们将「流量生成」，「批量流量还原」和「真实流量还原」是一个子图。「仿真环境」等是一个模块，完整的代码如下（需要注意，子图的名称必须以 cluster 开头，否则 graphviz 无法设别）：

digraph G {
SUMO_RL [style="filled", fontsize = 20, color="black", fillcolor="chartreuse"];
SUMO_RL -> 流量生成 [color="red"];
subgraph cluster_traffic{
bgcolor="mintcream";
label = "强化学习数据生成";
流量生成 -> 真实流量还原;
流量生成 -> 批量生成流量;
}
强化学习训练 [style="dashed", color="yellowgreen"];
强化学习测试 [style="dashed", color="yellowgreen"];
批量生成流量 -> 强化学习训练 [style="dashed", color="skyblue"];
真实流量还原 -> 强化学习测试 [style="dashed", color="skyblue"];
SUMO_RL -> 仿真环境 [color="red"];
subgraph cluster_rl_env{
bgcolor = "mintcream";
label = "强化学习交互环境";
仿真环境 -> 获得车辆属性;
获得车辆属性 -> 提取环境特征 [label="observation and reward"];
仿真环境 -> 异步控制不同信号灯 [label="action"];
}
与SUMO交互 [style="dashed", color="yellowgreen"];
提取环境特征 -> 与SUMO交互 [style="dashed", color="skyblue"];
异步控制不同信号灯 -> 与SUMO交互 [style="dashed", color="skyblue"];
}

最终的效果如下图所示，可以看到我们设置了子图的背景色（同时也可以看一下 node 和 edge 的样式的设计）：

在 VS Code 中显示 dot 文件

我们可以安装 Graphviz Preview 插件来实时显示图像。在安装之前，我们需要确保 Graphviz 已经下载安装好了（Graphviz Download），如果是 Ubuntu，直接使用 sudo apt install graphviz 安装即可。下图是 Graphviz Preview 插件，在全部安装完毕之后，就可以显示了：

Python Graphviz 简单介绍

上面我们直接使用 dot 文件来生成树形图。但是这样还是有些不是很方便，特别是当树形图比较复杂的时候。这个时候就可以结合 python 的 graphviz 库来进行绘制。首先需要使用 pip 来安装 graphviz。

pip install graphviz

最基础的 Python Graphviz 图像

我们用 graphviz 的例子来进行说明。首先创建一个最基础的 Graphviz 图像，创建一个 Hello -> World 的连接图。

from graphviz import Digraph
g = Digraph('G', filename='hello.gv')
g.node('node1', label='Hello')
g.node('node2', label='World')
g.edge('node1', 'node2')
g.view()

上面我们首先初始化了一个 Digraph 类，它可以让我们设置这个 graph 的名称，上面是 'G'，和最后 dot 文件的名称，即 filename 这个参数。
接着我们使用 node 来创建节点，其中第一个参数是 node 的 name，第二个参数是 node 的 label，也就是最终会显示在图中的文字。
最后使用 edge 来将 node 连起来。edge 的参数为 node 的 name，第一个为起点，第二个为终点。

最终绘制出的结果如下所示：

同时，我们可以通过 g.source 来获取 dot 代码。获取的结果如下所示：

给树形图增加样式

在上面的基础上，我们可以增加树形图的样式。例如我们可以修改 edge 连接的箭头的样式。直接在上面初始化 Digraph 之后进行修改。

g.edge_attr.update(arrowhead='vee', arrowsize='2') # edge 的样式

此时的树形图如下所示：

上面是从整体上修改整个图的样式。我们也可以单独修改某个 node 的样子。下面我们把某个 node 的样子修改为五角星。只需要直接在 node 中增加 shape 的参数，如下所示：

g.node('node1', label='Hello', shape='star')
g.node('node2', label='World', shape='egg')

最终的效果图如下所示，关于 node 支持的形状，可以查看链接，Polygon-based Nodes：

我们还可以是设置更加复杂的 node 的样式，可以使用类似 html 的 label。这些被写在 label 下，使用<,> 来包裹内容。下面看一个例子，需要注意的是，下面在 Digraph 中使用了 node_atrr，会对整个图的 node 全局有效，这里是使得 node 只显示文字，不显示边框：

from graphviz import Digraph
g = Digraph('G', node_attr={'shape': 'plaintext'}, filename='hello.gv')
g.graph_attr['rankdir'] = 'LR'
g.edge_attr.update(arrowhead='vee', arrowsize='2') # edge 的样式
g.node('node1', label='Hello', shape='star')
g.node('node2', label='World', shape='egg')
g.node('tab', label='''<

left right

>''')
g.edge('node1', 'node2')
g.edge('node1', 'tab')
g.view()

最终的效果如下图所示：

Python Graphviz 例子

这里会记录一些 graphviz 的例子，一些其他的树形图可以基于这些例子做进一步的扩展。这一部分的内容参考自，Graphviz 的例子。

自定义 node 样式-fsm

这个例子中，我们修改 node 的样式，使得如果通过 node 来进行定义，node 为双圆的样式。

from graphviz import Digraph
f = Digraph('finite_state_machine', filename='fsm.gv')
f.attr(rankdir='LR', size='20,5')
# 单独定义的 node 会有双圆结构
f.attr('node', shape='doublecircle')
f.node('LR_0')
f.node('LR_3')
f.node('LR_4')
f.node('LR_8')
f.attr('node', shape='circle')
f.edge('LR_0', 'LR_2', label='SS(B)')
f.edge('LR_0', 'LR_1', label='SS(S)')
f.edge('LR_1', 'LR_3', label='S($end)')
f.edge('LR_2', 'LR_6', label='SS(b)')
f.edge('LR_2', 'LR_5', label='SS(a)')
f.edge('LR_2', 'LR_4', label='S(A)')
f.edge('LR_5', 'LR_7', label='S(b)')
f.edge('LR_5', 'LR_5', label='S(a)')
f.edge('LR_6', 'LR_6', label='S(b)')
f.edge('LR_6', 'LR_5', label='S(a)')
f.edge('LR_7', 'LR_8', label='S(b)')
f.edge('LR_7', 'LR_5', label='S(a)')
f.edge('LR_8', 'LR_6', label='S(b)')
f.edge('LR_8', 'LR_5', label='S(a)')
f.view()

最终的结果如下所示，可以看到上面单独定义的，LR_0、LR_3、LR_4、LR_8 都是两个圆：

包含子图的情况

我们可以创建多个 graph，不同的 graph 的样式不同，最后将他们合并。直接看一下下面的例子。

from graphviz import Digraph
g = Digraph('G', filename='cluster.gv')
# NOTE: the subgraph name needs to begin with 'cluster' (all lowercase)
# so that Graphviz recognizes it as a special cluster subgraph
with g.subgraph(name='cluster_0') as c:
c.attr(style='filled', color='lightgrey')
c.node_attr.update(style='filled', color='white')
c.edges([('a0', 'a1'), ('a1', 'a2'), ('a2', 'a3')])
c.attr(label='process #1')
with g.subgraph(name='cluster_1') as c:
c.attr(color='blue')
c.node_attr['style'] = 'filled'
c.edges([('b0', 'b1'), ('b1', 'b2'), ('b2', 'b3')])
c.attr(label='process #2')
g.edge('start', 'a0')
g.edge('start', 'b0')
g.edge('a1', 'b3')
g.edge('b2', 'a3')
g.edge('a3', 'a0')
g.edge('a3', 'end')
g.edge('b3', 'end')
g.node('start', shape='Mdiamond')
g.node('end', shape='Msquare')
g.view()

最终的结果如下所示，左侧的使用灰色进行填充，右侧的为蓝色的边框：

使用类 html 标签

上面我们介绍过类 html 标签来进行绘制，实际上我们可以进行一些简化，如下所示：

from graphviz import Digraph
s = Digraph('structs', filename='structs_revisited.gv',
node_attr={'shape': 'record'})
s.node('struct1', ' left| middle| right')
s.node('struct2', ' one| two')
s.node('struct3', r'hello\nworld |{ b |{c| d|e}| f}| g | h')
s.edges([('struct1:f1', 'struct2:f0'), ('struct1:f2', 'struct3:here')])
s.view()

最终的结果如下所示，上面的 f0 等其实都是 label，连线的时候会根据这些 label 进行连线。

我们把上面的 f0 等都去掉，做一些变化：

from graphviz import Digraph
s = Digraph('structs', filename='structs_revisited.gv',
node_attr={'shape': 'record'})
s.node('struct1', 'left| middle| right')
s.node('struct2', '{a|{b1|b2|b3}|c}')
s.node('struct3', r'hello\nworld |{ b |{c|d|e}| f}| g | h')
s.edges([('struct1', 'struct2'), ('struct1', 'struct3')])
s.view()

可以得到下面的图，有横向和纵向的分割：

绘制树结构

下面使用 graphviz 来绘制二叉树：

from graphviz import Digraph, nohtml
g = Digraph('g', filename='btree.gv',
node_attr={'shape': 'record', 'height': '.1'})
g.node('node0', nohtml(' | G|'))
g.node('node1', nohtml(' | E|'))
g.node('node2', nohtml(' | B|'))
g.node('node3', nohtml(' | F|'))
g.node('node4', nohtml(' | R|'))
g.node('node5', nohtml(' | H|'))
g.node('node6', nohtml(' | Y|'))
g.node('node7', nohtml(' | A|'))
g.node('node8', nohtml(' | C|'))
g.edge('node0:f2', 'node4:f1')
g.edge('node0:f0', 'node1:f1')
g.edge('node1:f0', 'node2:f1')
g.edge('node1:f2', 'node3:f1')
g.edge('node2:f2', 'node8:f1')
g.edge('node2:f0', 'node7:f1')
g.edge('node4:f2', 'node6:f1')
g.edge('node4:f0', 'node5:f1')
g.view()

最终的结果如下所示：

Python Graphviz 的使用-绘制树形图

简介