ast.NodeTransformerによるコード書き換えの例

Converterを書くのが面倒だというはなし の関係でちょっと試した内容を書いておく。
Pythonでast触る機会はほとんどないので、備忘録として。

やりたかったこと

たとえば "v * 2" という文字列が与えられた時に、引数vを2倍して返すようなfunctionを生成したい。
(引数の名前がvであることは既知とする)

実現の方針

Single Statementの場合

上の例のように、与えられた文字列がSingle Statementの場合は非常に簡単で、単に文字列操作でlambdaにしてevalすればいい。

single.py

text = "v * 2"
func = eval("lambda v:" + text)

Multiple Statementsへの対応

与えられた文字列に複数の式が含まれる可能性がある場合は、上の方法は使えない。
例えば

from System.Windows import Thickness
Thickness(v, 0, 0, 0)

と指定された場合のことを考える。
とりあえずぱっと思いつくのは、同じように文字列操作でfunction defの形にして、execする方法。
(このコードだと生成されたfunctionはNoneしか返さないんだけど、その話は後で)

multi.py

text = """
from System.Windows import Thickness
Thickness(v, 0, 0, 0)
"""
import re
source = "def test(v):\n" + re.sub("^", "  ", text, flags=re.M)
exec(source)
return test

ほとんどの場合はこれでいいんだけど、例えば以下のように、複数行リテラルが与えられた場合に、上のような正規表現による単純なインデント処理ではリテラルの中身までインデントされてしまうという問題がある。

"""value of v is:
{v}""".format(v=v)

ほとんど考慮する必要もないようなレアケースなんだけど、ダメなパターンがあるのが分かっててそのままというのもちょっと座りが悪いので、別案としてast変換でfunction def化する方針でいく。

1. まず、 def test(v): pass という空のfunction defに対するastを生成する。
2. 生成したastのPassノードの部分を、与えられた文字列から生成したastノードに置き換える。

という手順。

コードで書くとこんな感じ。

functionize.py

import ast

class _MakeFunction(ast.NodeTransformer):
    def __init__(self, body):
        self.body = body

    def run(self):
        template_ast = ast.parse('def test(v): pass', mode='exec')
        return self.visit(template_ast)
 
    def visit_Pass(self, node):
        return ast.parse(self.body, mode='exec').body
 
if __name__ == '__main__':
    body = "from System.Windows import Thickness\nThickness(v, 0, 0, 0)"
    functionized_ast = _MakeFunction(body).run()
    exec(compile(functionized_ast, '<string>', mode='exec'))

戻り値の生成

これでfunctionにはなったけど、先に書いたようにこのままでは戻り値がない。
なのでさらにast変換を使って、Rubyのように最後に評価した式の結果が戻り値として返るように書き換える。
手順は以下の通り。

1. functionの最初に　_retval_ = None というコードを追加する
2. functionの最後に return _retval_ をコード追加する
3. function中の出てくる式を、すべて結果を _retval_ に対する代入文に書き換える。

これを実行するTransformerをコードで書くとこう。

functionize.py

class _ReturnLastResult(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        self.generic_visit(node)
        retval_assign = \
                ast.Assign(targets=[ast.Name(id='_retval_', ctx=ast.Store())],
                           value=ast.NameConstant(value=None))               
        retval_return = \
                ast.Return(value=ast.Name(id='_retval_', ctx=ast.Load()))
        node.body.insert(0, retval_assign)
        node.body.append(retval_return)
        return ast.fix_missing_locations(node)

    def visit_Expr(self, node):
        target = ast.Name(id='_retval_', ctx=ast.Store())
        assign = ast.copy_location(
                    ast.Assign(targets=[target], value=node.value), node)
        return assign

コード

最終的なコードはこうなった。
https://gist.github.com/wonderful-panda/8a22b74248a60cc8bb22

結局元のエントリではSingle Statementのみ考慮することにしてこれは使わなかったけど。