Python の ast モジュール入門 (NodeVisitor を使う)

前回の [Python の ast モジュール入門 (抽象構文木を辿る)](http://qiita.com/t2y/items/0964d01bf3db0233e3c1) では、ast モジュールのヘルパー関数を使って抽象構文木を辿ることを紹介しました。

## 抽象構文木を *NodeVisitor* で辿る

ast モジュールのヘルパー関数を使うのも1つの方法ですが、[*ast.NodeVisitor*](http://docs.python.jp/3.4/library/ast.html#ast.NodeVisitor) を使うと、もっとお手軽に抽象構文木を辿ることができます。やっていることはヘルパー関数を使うのと同じだというのを *NodeVisitor* の実装を見た方が分かりやすいのでその実装から紹介します。*NodeVisitor* は [Visitor パターン](http://qiita.com/t2y/items/3ecdc2643426a345d75a) というデザインパターンの1つです。

```py3:3.4
class NodeVisitor(object):

    def visit(self, node):
        """Visit a node."""
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)

    def generic_visit(self, node):
        """Called if no explicit visitor function exists for a node."""
        for field, value in iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, AST):
                        self.visit(item)
            elif isinstance(value, AST):
                self.visit(value)
```

*visit_NodeClassname* が定義されていないときは [*ast.iter_fields*](http://docs.python.jp/3.4/library/ast.html#ast.iter_fields) で抽象構文木を辿る *generic_visit* が実行されます。抽象構文木のノードクラスは [*ast.AST*](http://docs.python.jp/3.4/library/ast.html#ast.AST) を基底クラスに取るので、 *isinstance(value, AST)* でノードインスタンスかどうかを判定して再起的にトラバース (*self.visit()*) していくようになっています。

実際に使ってみましょう。*NodeVisitor* を継承したクラスを定義します。

```py3:3.4
>>> import ast
>>> source = """
... import sys
... def hello(s):
...     print('hello {}'.format(s))
... hello('world')
... """
>>> class PrintNodeVisitor(ast.NodeVisitor):
...     def visit(self, node):
...         print(node)
...         return super().visit(node)
...
>>> tree = ast.parse(source)
>>> PrintNodeVisitor().visit(tree)
<_ast.Module object at 0x10bec7b38>
<_ast.Import object at 0x10bec7b70>
<_ast.alias object at 0x10bec7ba8>
<_ast.FunctionDef object at 0x10bec7c18>
<_ast.arguments object at 0x10bec7c50>
<_ast.arg object at 0x10bec7c88>
<_ast.Expr object at 0x10bec7d30>
<_ast.Call object at 0x10bec7d68>
<_ast.Name object at 0x10bec7da0>
<_ast.Load object at 0x10bebe0f0>
<_ast.Call object at 0x10bec7e10>
<_ast.Attribute object at 0x10bec7e48>
<_ast.Str object at 0x10bec7e80>
<_ast.Load object at 0x10bebe0f0>
<_ast.Name object at 0x10bec7eb8>
<_ast.Load object at 0x10bebe0f0>
<_ast.Expr object at 0x10bec7f28>
<_ast.Call object at 0x10bec7f60>
<_ast.Name object at 0x10bec7f98>
<_ast.Load object at 0x10bebe0f0>
<_ast.Str object at 0x10bec7fd0>
```

簡単に抽象構文木を辿りながらノードを表示することができました。ある特定のノードをフックするには *visit_NodeClassname* のメソッドを定義します。

```py3:3.4
>>> class PrintExprNodePisitor(ast.NodeVisitor):
...     def visit_Expr(self, node):
...         print('Expr is visited')
...         return node
... 
>>> PrintExprNodePisitor().visit(tree)
Expr is visited
Expr is visited
```

*PrintNodeVisitor* での出力と見比べると *Expr* ノードを2回辿っていることが分かりますね。

## 抽象構文木を *NodeTransformer* で変更する

*NodeVisitor* はトラバースの途中でノードを変更できません。ノードを変更したい用途のときには [*ast.NodeTransformer*](http://docs.python.jp/3.4/library/ast.html#ast.NodeTransformer) を使います。

まずは簡単なソースコードから試してみましょう。

```py3:3.4
>>> import ast
>>> source = """
... print(s)
... """
>>> s = 'hello world'
>>> code = compile(source, '<string>', 'exec')
>>> exec(code)
hello world
```

[*ast.dump*](http://docs.python.jp/3.4/library/ast.html#ast.dump) を使って、このソースコードがどういった抽象構文木に展開されるかを確認します。 

```py3:3.4
>>> tree = ast.parse(source)
>>> ast.dump(tree)
"Module(body=[Expr(value=Call(func=Name(id='print', ctx=Load()), args=[Name(id='s', ctx=Load())], keywords=[], starargs=None, kwargs=None))])"
```

これを眺めながら何か適当な例を考えてみます。

ここでは例として、出力される文字列を反転してみましょう。いろいろやり方はありそうですが、*print* 文を別の関数に置き換えるのを試してみます。

```py3:3.4
>>> class ReversePrintNodeTransformer(ast.NodeTransformer):
...     def visit_Name(self, node):
...         if node.id == 'print':
...             name = ast.Name(id='reverse_print', ctx=ast.Load())
...             return ast.copy_location(name, node)
...         return node
...
>>> def reverse_print(s):
...     print(''.join(reversed(s)))
... 
>>> code = compile(ReversePrintNodeTransformer().visit(tree), '<string>', 'exec')
>>> exec(code)
dlrow olleh
>>> s = 'revese print'
>>> exec(code)
tnirp esever
```

それっぽく動きました。*print* 文が *reverse_print* 関数に置き換えられて実行されています。

[*ast.copy_location*](http://docs.python.jp/3.4/library/ast.html#ast.copy_location) を使うと、元のノードから *lineno* と *col_offset* をコピーしてくれます。この2つの属性がないと AST オブジェクトを *compile* できません。

失敗する例を試してみましょう。

```py3:3.4
>>> from ast import *
>>> expression_without_attr = dump(parse('1 + 1', mode='eval'))
>>> expression_without_attr
'Expression(body=BinOp(left=Num(n=1), op=Add(), right=Num(n=1)))'
>>> code = compile(eval(expression_without_attr), '<string>', 'eval')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: required field "lineno" missing from expr
```

ここで *ast.dump* に属性も出力するように *include_attributes=True* を渡します。

```py3:3.4
>>> expression_with_attr = dump(parse('1 + 1', mode='eval'), include_attributes=True)
>>> expression_with_attr
'Expression(body=BinOp(left=Num(n=1, lineno=1, col_offset=0), op=Add(), right=Num(n=1, lineno=1, col_offset=4), lineno=1, col_offset=0))'
>>> code = compile(eval(expression_with_attr), '<string>', 'eval')
>>> eval(code)
2
```

*lineno* や *col_offset* を出力することで *ast.dump* の出力からそのまま AST オブジェクトを生成して (eval して) コンパイルすることもできました。

また別の解決方法として [*ast.fix_missing_locations*](http://docs.python.jp/3.4/library/ast.html#ast.fix_missing_locations) を使う方法もあります。先ほどの *expression_without_attr* を使ってやってみましょう。

```py3:3.4
>>> code = compile(fix_missing_locations(eval(expression_without_attr)), '<string>', 'eval')
>>> eval(code)
2
```

今度は *compile* できましたね。*fix_missing_locations* のドキュメントによると、

> 生成されたノードに対しこれらを埋めて回るのはどちらかというと退屈な作業なので、このヘルパーが再帰的に二つの属性がセットされていないものに親ノードと同じ値をセットしていきます。

と、自動的にセットしてくれるそうです。

## 抽象構文木をいじる (*NodeTransformer* を使う) とき

実際に抽象構文木をいじって解決したい課題を見つけるのもちょっと難しいのですが、いくつか見つけたものを紹介します。

* [Pythonで正しく日本語を eval する](http://www.gembook.org/2011-02-15.html)
* [Pythonで++i](http://www.gembook.org/2012-10-24.html)
* [ast.NodeTransformerによるコード書き換えの例](http://qiita.com/wonderful_panda/items/dfb12a7de244033aae85)

データとして Python のコードを扱うとき、つまりは扱わないと難しいなにかがあるようなときに思い出すと便利なときもあるかもしれません。