Python の ast モジュール入門 (抽象構文木を辿る)

概要

Python で 抽象構文木 (Abstract Syntax Tree) を扱うのに、その名の通り ast モジュール が標準ライブラリとして提供されています。

PEP で AST について言及しているのは PEP 339 - Design of the CPython Compiler のようです。Python 2.5 で ast モジュールが追加され、2.6 でヘルパー関数が提供されたようです。余談ですが、2.6 の What's New によると、

The ast module provides an Abstract Syntax Tree representation of Python code, and Armin Ronacher contributed a set of helper functions that perform a variety of common tasks.

これらのヘルパー関数は Armin Ronacher 氏がコントリビュートしたようです。

ast モジュールのヘルパー関数

普通の開発用途でも ast.literal_eval を使ったことがある方もいるのではないでしょうか。eval との違いは、評価対象をリテラルの式のみに限定することで安全に評価できる点です。簡単に紹介すると、

3.4

>>> import ast
>>> eval('1 + 1')
2
>>> ast.literal_eval('1 + 1')
2

>>> x = 1
>>> eval('x + 1')
2
>>> ast.literal_eval('x + 1')
...
ValueError: malformed node or string: <_ast.BinOp object at 0x104c1e160>

>>> eval('max([1, 3, 5])')
5
>>> ast.literal_eval('max([1, 3, 5])')
...
ValueError: malformed node or string: <_ast.Call object at 0x104c4cd68>

といったように実行環境から変数にアクセスしたり、関数を呼び出したりといったことはできないように制御されています。これで要件を満たせるなら、セキュリティの観点からも literal_eval を使うのが望ましいです。

さらに余談ですが、literal_eval の実装を覗いてみると、一旦 AST オブジェクトに compile した後で諸々のチェックを行っているので、文字列だけでなく AST オブジェクトを直接渡すこともできます。

3.4

>>> ast.literal_eval(ast.parse('1 + 1', mode='eval'))
2

いきなり ast.parse というヘルパー関数が出てきました。この関数の実装は単純に組み込み関数の compile を使っています。

3.4

def parse(source, filename='<unknown>', mode='exec'):                               
    """                                                                             
    Parse the source into an AST node.                                              
    Equivalent to compile(source, filename, mode, PyCF_ONLY_AST).                   
    """                                                                             
    return compile(source, filename, mode, PyCF_ONLY_AST)              

ここでドキュメントには、

抽象構文木を作成するには、 ast.PyCF_ONLY_AST を組み込み関数 compile() のフラグとして渡すか、あるいはこのモジュールで提供されているヘルパー関数 parse() を使います。

と、さらっと書いてあります。あれ！？組み込み関数の compile ってコードオブジェクト返すんじゃなかったっけ？と思っていたら、フラグとして ast.PyCF_ONLY_AST を渡すことで Expression オブジェクトを返すという挙動が違うようです。なるほど、それを意識させないための ast.parse でもあるわけですね。

3.4

>>> compile('1 + 1', '<string>', 'eval')
<code object <module> at 0x104c01c00, file "<string>", line 1>

>>> compile('1 + 1', '<string>', 'eval', ast.PyCF_ONLY_AST)
<_ast.Expression object at 0x104c4ceb8>

コード解析の事例 (pyflakes)

ast モジュールのドキュメントには、Node クラスとヘルパー関数の説明が一通りありますが、ざっと眺めてみても、(チュートリアル的に) これで何をすればいいの？って感じです。おそらくは ast モジュールを必要とする開発者にはそれで十分なのでしょう。

そこで具体例をみながらもう少し掘り下げてみます。先ほど紹介した 2.6 の What's New によると、

These will be useful for HTML templating packages, code analyzers, and similar tools that process Python code.

Python コードを処理するテンプレート、コード解析、そういった類のツールに便利だとあります。ここでは、lint ツールの1つである pyflakes で ast モジュールの実用例をみていきます。

pyflakes の使い方は、

pyflakes-err1.py

import sys

def f(sequence):
    value = 2
    for i in sequence:
        yield i + 1

g = f(range(3))

のようなコードがあったとしてコマンドラインだと pyflakes コマンドで

$ pip install pyflakes
$ pyflakes pyflakes-err1.py 
pyflakes-err1.py:1: 'sys' imported but unused
pyflakes-err1.py:4: local variable 'value' is assigned to but never used

このようにエラーチェックができます。インタラクティブシェルからでも以下のように実行できます。

3.4

>>> import ast
>>> from pyflakes.checker import Checker
>>> source = open('pyflakes-err1.py').read()
>>> tree = ast.parse(source)
>>> checker = Checker(tree, 'pyflakes-err1.py')
>>> for msg in checker.messages:
...   print(str(msg))
... 
pyflakes-err1.py:4: local variable 'value' is assigned to but never used
pyflakes-err1.py:1: 'sys' imported but unused

compile した AST ノードを Checker に渡すのみなので簡単ですね。これらの処理をみながら ast モジュールの使い方を調べてみます。当初は実際の checker.py のコードを見ながら解説しようと思ったのですが、あまり分かりやすいコードではないので ast モジュールを使う上での本質的なところのみを解説してみます。

3.4

>>> tree
<_ast.Module object at 0x10abe0390>
>>> tree.body
[<_ast.Import object at 0x10abe04e0>,
 <_ast.FunctionDef object at 0x10abe6470>,
 <_ast.Assign object at 0x10abfb2b0>]

compile が返す Module ノードオブジェクトには body という属性があり、モジュール内のそれぞれのノードを含むリストを持っています。これを汎用的に扱うには、ノードの _fields の属性名を調べてそれぞれのノードごとに処理を実装します。

Checker がやっているメイン処理を疑似コードで簡単に表すと以下のようなイメージです。

>>> for name in tree._fields:
...   for node in getattr(tree, name):
...     get_node_handler(node.__class__.__name__.upper())(node)

Checker の内部処理はスコープやスタックの制御が煩雑なのでここでは触れません。抽象構文木を辿りながらノードごとのハンドラーを呼び出すような作りになっています。例えば Import ノードに対するハンドラーメソッドは以下のように定義されています。

checker.py

    def IMPORT(self, node):
        for alias in node.names:
            name = alias.asname or alias.name
            importation = Importation(name, node)
            self.addBinding(node, importation)

pyflakes では ast モジュールのヘルパー関数を全く使っていないため、自前でそういった処理を実装しています (そのためコードの見通しが悪いようにもみえます) 。

例えば、先ほどの Module ノードの子ノードを取得する処理は ast.iter_child_nodes で取得できます。

3.4

>>> list(ast.iter_child_nodes(tree))
[<_ast.Import object at 0x10abe04e0>,
 <_ast.FunctionDef object at 0x10abe6470>,
 <_ast.Assign object at 0x10abfb2b0>]

他にも node._fields のそれぞれのフィールドを返す ast.iter_fields というヘルパー関数もあります。例えば FunctionDef ノードで違いをみると分かりやすいです。

3.4

>>> func_node = tree.body[1]
>>> list(ast.iter_fields((func_node)))
[('name', 'f'),
 ('args', <_ast.arguments object at 0x10abe64a8>),
 ('body', [<_ast.Assign object at 0x10abf9f28>,
           <_ast.For object at 0x10abfb080>]),
 ('decorator_list', []),
 ('returns', None)]

>>> list(ast.iter_child_nodes((func_node)))
[<_ast.arguments object at 0x10abe64a8>,
 <_ast.Assign object at 0x10abf9f28>,
 <_ast.For object at 0x10abfb080>]

lint ツールであれば、これらを使って抽象構文木を辿りながら処理を書いて行けば良さそうというのがイメージとして分かります。

もう1つ覚えておくと良さそうなのが ast.walk というヘルパー関数です。引数に与えたノードの全ての子孫ノードを 順不同 で返します。

3.4

>>> list(ast.walk(func_node))
[<_ast.FunctionDef object at 0x10abe6470>,
 <_ast.arguments object at 0x10abe64a8>,
 <_ast.Assign object at 0x10abf9f28>,
 <_ast.For object at 0x10abfb080>,
 <_ast.arg object at 0x10abf9e80>,
 <_ast.Name object at 0x10abf9f60>,
 <_ast.Num object at 0x10abfb048>,
 <_ast.Name object at 0x10abfb0b8>,
 <_ast.Name object at 0x10abfb0f0>,
 <_ast.Expr object at 0x10abfb128>,
 <_ast.Store object at 0x10abefbe0>,
 <_ast.Store object at 0x10abefbe0>,
 <_ast.Load object at 0x10abefac8>,
 <_ast.Yield object at 0x10abfb160>,
 <_ast.BinOp object at 0x10abfb198>,
 <_ast.Name object at 0x10abfb1d0>,
 <_ast.Add object at 0x10abf19e8>,
 <_ast.Num object at 0x10abfb208>,
 <_ast.Load object at 0x10abefac8>]

コード解析の事例 (flake8)

ast.walk を使う具体例として、Hacking Python AST: checking methods declaration で flake8 (lint ツールの統合ツールみたいなもの) に staticmethod のチェック処理を追加するプラグインが紹介されています。

3.4

for stmt in ast.walk(self.tree):
    # Ignore non-class
    if not isinstance(stmt, ast.ClassDef):
        continue
    # If it's a class, iterate over its body member to find methods
    for body_item in stmt.body:
        # Not a method, skip
        if not isinstance(body_item, ast.FunctionDef):
            continue
        # Check that it has a decorator
        ...

このように isinstance で AST のノードオブジェクトを調べながらチェック処理を実装していくという流れは pyflakes でも同様です。こんな風に使えば良いんだというイメージは伝わりますね。

まとめ

ast モジュールの入門として lint ツールから実用例を学びました。

次の記事

• [Python の ast モジュール入門 (NodeVisitor を使う)]
  (http://qiita.com/t2y/items/c8877cf5d3d22cdcf2a8)