Pythonの抽象構文木をGraphvizで可視化する

はじめに

Pythonには抽象構文木を扱うastというモジュールがある。これを使って抽象構文木を可視化してみる。

スクリプト

import ast
import sys

from graphviz import Digraph


def visit(node, nodes, pindex, g):
    name = str(type(node).__name__)
    index = len(nodes)
    nodes.append(index)
    g.node(str(index), name)
    if index != pindex:
        g.edge(str(index), str(pindex))
    for n in ast.iter_child_nodes(node):
        visit(n, nodes, index, g)


filename = sys.argv[0]
if len(sys.argv) > 1:
    filename = sys.argv[1]

with open(filename) as f:
    graph = Digraph(format="png")
    tree = ast.parse(f.read())
    visit(tree, [], 0, graph)
    graph.render("test")

実行

例えばこんなファイルを食わせる。

def func(a, b):
    return a + b


def func2(a):
    print(a)

print("Hello")

$ python parse.py test.py

すると、こんな構文木が得られる。

image0.png

ちなみに引数を指定しないと自分自身をパースする。

$ python parse.py

得られる結果はこんな感じ。

image1.png

動作原理

基本的にはast.parseで得られた木から再帰的にたどっていくだけ。ast.iter_child_nodes(node)でnodeにぶら下がるノードのイテレータが得られるので、それに対してforを回す。

Graphivzのノードのキャプションはtype(node).__name__としたが、そのままg.edgeでつないでしまうと同じ名前が同じノードとみなされてしまうので、一工夫が必要。

そこで、ノードをリストに追加することで、それぞれに一意な通し番号をふる。

    name = str(type(node).__name__)
    index = len(nodes)
    nodes.append(index)
    g.node(str(index), name)

で、この通し番号をGraphvizのノード番号として、繋げば良い。

 g.edge(str(index), str(pindex))

最初は、オブジェクトをリストではなくハッシュに登録してたんだけど、抽象構文オブジェクト、例えばast.Loadはなぜか同じオブジェクトが使い回されているようで、同じノードとみなされてしまう。上記のように一意な通し番号を振れば大丈夫。