Python: ast (Abstract Syntax Tree: 抽象構文木) モジュールについて

標準ライブラリの ast モジュールを使うと Python で書かれたソースコードを構文解析できる。それによって得られるオブジェクトは AST (Abstract Syntax Tree: 抽象構文木) と呼ばれる。その使い道としては、例えば Python の lint ツールなどが考えられる。つまり、ソースコードの構造を確かめることでまずいところを見つけ出すことができる。

使った環境は次の通り。

$ sw_vers 
ProductName:    Mac OS X
ProductVersion: 10.11.4
BuildVersion:   15E65
$ python --version
Python 3.5.1

使ってみる

論よりコードということで実際にソースコードを ast モジュールで処理してみよう。

まずは次のようにシンプル極まりないサンプルを用意する。

$ cat << 'EOF' > helloworld.py
#!/usr/bin/env python
# -*- coding: utf-8 -*-


def main():
    print('Hello, World!')


if __name__ == '__main__':
    main()
EOF

ここからは Python の REPL を使って動作確認をする。先ほどのファイルと同じディレクトリで python コマンドを実行しよう。

$ python

まずは、先ほど用意したソースコードの内容を source という名前の変数に読み込んでおく。

>>> FILENAME = 'helloworld.py'
>>> with open(FILENAME, 'r') as f:
...     source = f.read()
...

読み込んだ内容を ast モジュールの parse() 関数に渡して構文解析しよう。

>>> import ast
>>> ast_object = ast.parse(source, FILENAME)

あるいは代わりに組み込み関数の compile() を使っても構わない。実際のところ、上記でも内部的には以下を呼び出している。

>>> import ast
>>> ast_object = compile(source, FILENAME, 'exec', ast.PyCF_ONLY_AST)

すると ast モジュールの解析結果として Module オブジェクトが得られる。これは Python モジュール (ようは *.py ファイル) に対応している。

>>> ast_object
<_ast.Module object at 0x103818e80>

これらのオブジェクトはノードと呼ばれる。

すべてのノードは ast.AST 抽象基底クラスのサブクラスになっている。

>>> isinstance(ast_object, ast.AST)
True

持っているアトリビュートはノードの種類によって異なる。もし入れ子になった要素があるときは、それらがリストになった body というアトリビュートを持っている。

>>> ast_object.body
[<_ast.FunctionDef object at 0x10561d490>, <_ast.If object at 0x10561d590>]

上記ではモジュールが関数と if ステートメントをひとつずつ持っていることがノードの内容から判断できる。

関数や if ステートメントのノードには、それがソースコードの何行目にあるかという情報が lineno として得られる。

>>> ast_object.body[0].lineno
5
>>> ast_object.body[1].lineno
9

関数については名前の情報も name として得られる。それに対し if ステートメントは名前がないので name アトリビュートもない。

>>> ast_object.body[0].name
'main'
>>> ast_object.body[1].name
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'If' object has no attribute 'name'

ノードがどういったアトリビュートを持っているかは ast.dump() 関数を使うと分かりやすい。

>>> ast.dump(ast_object.body[0])
"FunctionDef(name='main', args=arguments(args=[], vararg=None, kwarg=None, defaults=[]), body=[Print(dest=None, values=[Str(s='Hello, World!')], nl=True)], decorator_list=[])"

全体の構造を出力してみる

先ほどは REPL を使って ast モジュールの基本的な使い方を確認した。次は試しに AST の全体の構造を出力するプログラムを書いてみる。先ほどの helloworld.py と同じディレクトリに用意しよう。

$ cat << 'EOF' > tree.py
#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function

import ast


def walk(node, indent=0):
    # 入れ子構造をインデントで表現する
    print(' ' * indent, end='')

    # クラス名を表示する
    print(node.__class__, end='')

    # 行数の情報があれば表示する
    if hasattr(node, 'lineno'):
        msg = ': {lineno}'.format(lineno=node.lineno)
        print(msg, end='')

    # 改行を入れる
    print()

    # 再帰的に実行する
    for child in ast.iter_child_nodes(node):
        walk(child, indent=indent+4)


def main():
    FILENAME = 'helloworld.py'

    with open(FILENAME, 'r') as f:
        source = f.read()

    tree = ast.parse(source, FILENAME)
    walk(tree)


if __name__ == '__main__':
    main()
EOF

ast.iter_child_nodes() 関数を使ってノードの子要素を取り出して、それを再帰的に呼び出しながら情報を書き出している。ノードの入れ子構造はインデントで表現した。

実行してみると helloworld.py の構造が表示される。 main 関数を表す FunctionDef の中に Print が入っていて、さらにその中には Str が入っていたり。見ているだけで Python のソースコードは内部的にこう表現されているのかと面白い。

$ python tree.py
<class '_ast.Module'>
    <class '_ast.FunctionDef'>: 14
        <class '_ast.arguments'>
        <class '_ast.Print'>: 15
            <class '_ast.Str'>: 15
    <class '_ast.If'>: 28
        <class '_ast.Compare'>: 28
            <class '_ast.Name'>: 28
                <class '_ast.Load'>
            <class '_ast.Eq'>
            <class '_ast.Str'>: 28
        <class '_ast.Expr'>: 29
            <class '_ast.Call'>: 29
                <class '_ast.Name'>: 29
                    <class '_ast.Load'>

特定のノードを選択的に処理するには

先ほどの例ではノードを再帰的に処理する関数を自前で用意した。実は ast にはそれ用のヘルパが用意されているので、そちらを使っても構わない。具体的には ast.NodeVisitor クラスを継承してコールバックを実装する。これは、例えば特定のノードだけを選択的に処理したいといったときは特に便利だ。

次のサンプルコードでは FunctionVisitor という関数だけを処理するクラスを定義している。 ast.NodeVisitor クラスを継承した上で、visit_ から始まるメソッドを実装する。メソッド名の後ろには処理したいノードの名前を入れよう。今回のように関数であれば FunctionDef になる。

$ cat << 'EOF' > visitor.py
#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function

import ast


class FunctionVisitor(ast.NodeVisitor):

    def visit_FunctionDef(self, node):
        """FunctionDef オブジェクトだけに反応するコールバック"""
        print(node.name, end='')
        print(': ', end='')
        print(node.lineno)


def main():
    FILENAME = 'helloworld.py'

    with open(FILENAME, 'r') as f:
        source = f.read()

    tree = ast.parse(source, FILENAME)
    FunctionVisitor().visit(tree)


if __name__ == '__main__':
    main()
EOF