prestringの表現をASTから生成する機能を追加した

python prestring ast

github.com

個人的にコード生成をするときにはprestringというパッケージをベースにして行うことが多い。テンプレートエンジンによるコード生成やASTを直接触る方法との差分は別の機会に書くとして今回追加しようとした機能についてまとめておく。

コード生成

コード生成というのは何かの入力を元に何らかのコードを生成するもののこと。

<コード生成用のコード> + <設定> => <生成されたコード>

swagger関係のものなど巷にはコード生成用のツールが溢れていたりはする。しかし自分たちの利用している環境（コード）に則さないなどちょっと不便さを感じたときには調整する必要が出てきたりする。この調整に結構手間取ることが多い。

自分で作るとなると今度はコード生成用のコードが必要になる（当たり前だけれど）。

コード生成用のコード

ここでのコード生成用のコードというのは実行するとソースコードになるコードのこと。例えば以下の様なコードを実行するとpythonのコードを出力する。

hello.py

from prestring.python import Module
m = Module()  # noqa
m.sep()


with m.def_('hello', 'name', '*', 'message: str =  "hello world"'):
    m.docstring('greeting message')
    m.stmt('print(f"{name}: {message}")')

withがたくさん出てくる以外はある程度出力結果が透けて見える程度の構文。そうではあるのだけど手書きするのが面倒。

生成後のコード

def hello(name, *, message: str = "hello world"):
    """greeting message"""
    print(f"{name}: {message}")

このようにコードを出力するようなコードをコード生成用のコードと呼んでいる（フォーマッターが必要、入力データの走査、並行実行、複数ファイルへの出力、実行速度など考えるべきことは他にも色々あるけれど今回は省略）。

コードからコード生成用のコードを生成する

コード生成用のコードを書き始めるタイミングでは、出力後のコードのサンプルをリテラル的にそのまま書き下したいと言うことが多い。慣れればコード生成用のコードをそのまま直に書くということもできるようになっては来るのだけれど。やっぱりめんどくさい（ちなみにテンプレートエンジンでのコード生成は、このあたりのフェーズではけっこう高速に作業することができて、単にテンプレートに直接出力後のコードの文字列を記述するだけで良い）。この部分をどうにかしようという話。

入力として出力後のコード例を受け取って、ASTを作り、ASTから出力用のコードを生成すれば良さそう。ということで作った。

生成されたコード(を模したもの)
=> (AST)-> <コード生成用のコード> // ここの部分
=> <コード生成用のコード>  + <設定>
=> <生成されたコード>

例えば先程のhello.pyのコードを生成したい場合には以下の様なコードを書く。

def main():
    import inspect
    from prestring.python import Module
    from prestring.python.transform import transform_string

    m = Module()
    source = inspect.getsource(hello)
    m = transform_string(source, m=m)
    print(m)

transform_stringがコードの文字列を渡すと良い感じにやってくれる。inspectパッケージのgetsourceで自分自身のソースコードを取ってこれるのでこれを渡している。 以下のような出力になる。

with m.def_('hello', 'name', '*', 'message: str =  "hello world"'):
    m.docstring('greeting message')
    m.stmt('print(f"{name}: {message}")')

ちなみにtransform系のコマンドは3つある

  • transform_file -- ファイル名を渡して変換
  • transform_string -- ソースコード文字列を渡して変換
  • transform -- ASTを渡して変換

(transform_objectなどもあって良いかもしれない)

コマンドとして実行

毎回コードを書くのも面倒なのでコマンドとしても使えるようにした。入力としてソースコードのファイルを取る。

$ python -m prestring.python.transform hello.py | tee hello.template.py
from prestring.python import Module
m = Module()  # noqa
m.sep()


with m.def_('hello', 'name', '*', 'message: str =  "hello world"'):
    m.docstring('greeting message')
    m.stmt('print(f"{name}: {message}")')


print(m)

直接実行できるようなコードが生成されるので生成された結果を確かめられる。

$ python hello.template.py | tee hello.output.py


def hello(name, *, message: str =  "hello world"):
    """
    greeting message
    """
    print(f"{name}: {message}")

transform_string() などにmを渡さなくても同様の振る舞いになるがこのあたりの挙動は変更するかもしれない。

詳細

途中で力尽きたのでおまけに近い。

ASTからのコード生成

pythonでASTを触るライブラリはたくさんあるのだけれど。とりあえず標準ライブラリの範囲でやることにした。ただし標準ライブラリのastパッケージはたしかコメントなどの情報が消えてしまったような記憶がある。そんなわけで2to3(python2.xからpython3.xへの変換用のツール)に使われているライブラリのlib2to3を拝借して使うことにした。ちなみにyapfというコードフォーマッターもこのライブラリを使っている。

prestring.python.parseに小さなコード片を書いた。

transformと同様の感じで以下の様な関数がある。

  • parse_file
  • parse_string

lib2to3のnodeはけっこう便利で、文字列化すると自動でソースコードになる。

from prestring.python.parse import parse_file, dump_tree

t = parse_file("./hello.py")
print(str(t))

# def hello(name, *, message: str = "hello world"):
#     """greeting message"""
#     print(f"{name}: {message}")

作ったASTはdump_treeで出力できる。

from prestring.python.parse import parse_file, dump_tree

t = parse_file("./hello.py")
dump_tree(t)

小さめのコードを渡しても出力が長い。

file_input [2 children]
  funcdef[name='hello'] [5 children]
    NAME('def') [lineno=1, column=0, prefix='']
    NAME('hello') [lineno=1, column=4, prefix=' ']
    parameters [3 children]
      LPAR('(') [lineno=1, column=9, prefix='']
      typedargslist[args='name' ',' '*' ',' 'tname' '=' '"hello world"'] [7 children]
        NAME('name') [lineno=1, column=10, prefix='']
        COMMA(',') [lineno=1, column=14, prefix='']
        STAR('*') [lineno=1, column=16, prefix=' ']
        COMMA(',') [lineno=1, column=17, prefix='']
        tname [3 children]
          NAME('message') [lineno=1, column=19, prefix=' ']
          COLON(':') [lineno=1, column=26, prefix='']
          NAME('str') [lineno=1, column=28, prefix=' ']
        EQUAL('=') [lineno=1, column=32, prefix=' ']
        STRING('"hello world"') [lineno=1, column=34, prefix=' ']
      RPAR(')') [lineno=1, column=47, prefix='']
    COLON(':') [lineno=1, column=48, prefix='']
    suite [5 children]
      NEWLINE('\n') [lineno=1, column=49, prefix='']
      INDENT('    ') [lineno=2, column=0, prefix='']
      simple_stmt [2 children]
        STRING('"""greeting message"""') [lineno=2, column=4, prefix='']
        NEWLINE('\n') [lineno=2, column=26, prefix='']
      simple_stmt [2 children]
        power [2 children]
          NAME('print') [lineno=3, column=4, prefix='    ']
          trailer [3 children]
            LPAR('(') [lineno=3, column=9, prefix='']
            STRING('f"{name}: {message}"') [lineno=3, column=10, prefix='']
            RPAR(')') [lineno=3, column=30, prefix='']
        NEWLINE('\n') [lineno=3, column=31, prefix='']
      DEDENT('') [lineno=4, column=0, prefix='']
  ENDMARKER('') [lineno=4, column=0, prefix='']

ちなみに

loggingライブラリのコードを変換するとこんな感じになります。