pythonのastモジュールに不満がでたらlib2to3のコードを使ってみよう

python ast

はじめに

pythonのコードをparseするためにastモジュールが用意されていたりする。 このモジュールはpythonのコードをparseしてvisitor的なものでastをtravarseしてなどと便利ではあるのだけれど。 コメントなどの情報が消えてしまうなどの不満が出ることがある。 このような時にlib2to3用のコードを使ってみると良いのではないかという話。

lib2to3?

2to3というツールがあったりする。これはpython2.x用のコードをpython3.x用のコードに変換してくれるツール。 よく考えてみて欲しいのだけれど、2to3によってコメントの情報が失われることはない。そして2to3もおおよそASTを取り出してからの変換ということになっているはず。ということは2to3の内部のコードを覗いてみればコメントなどの情報を失うことなくAST変換を行う術が分かるはず。

何が言いたいかというと、コメント情報などの失われを防ぐためにlexerなどから作るなどということは不要ということ(ちなみに完全にフルでparserを再実装したbaronというものもあったりする。ただもう少し抽象度の高いredbaronから使う事がおすすめされていたりする)。

そして2to3の内部で使われているコードがlib2to3というもの。ちなみにこのlib2to3はyapfというコードフォーマッター(gofmtのようなもの)にも使われていたりする。

ちょっとしたコード変換

試しにlib2to3を利用してちょっとしたコード変換をしてみる。

例えば以下の様なコードがあるとする。

hello.py

def hello():
    # this is comment
    return "hello"

これを以下の様に変換してみる。

def *replaced*():
    # this is *replaced* comment
    return "hello"

こういう感じのコードを書けば良い。

from lib2to3 import pytree
from lib2to3 import pygram
from lib2to3.pgen2 import driver

default_driver = driver.Driver(pygram.python_grammar_no_print_statement, convert=pytree.convert)


def parse(code, parser_driver=default_driver):
    return parser_driver.parse_string(code, debug=True)


with open("hello.py") as rf:
    t =  parse(rf.read())
print(t)
t.children[0].children[1].value = "*replaced*"
t.children[0].children[4].children[1].prefix = "    # this is *replaced* comment\n"
print(t)

treeを直接触っているので何をやっているかはものすごく分かりづらいものではあるけれど。テキトウに関数名やコメント部分に *replaced* という文字列を挿入している。tree自体を文字列として出力するとおおよそそのままpythonコードとして出力されるというのも便利。

ちなみにファイルからtreeを作る際は以下でも良い。

t = default_driver.parse_file("hello.py")

もう少し真面目にするなら

もう少し真面目にするなら、このlib2to3用のASTに対するvisitorを作ってあげると良い。 yapfのpytree_visitor.pyなどが参考になる。

もう少し詳しいことは気が向いたら書くかもしれない。