Python に end を導入してみる

この記事は Python Advent Calendar 2016 9日目の記事です。

モチベーション

みなさん御存知の通り、Python は文法上のブロックがインデントによって表現されています。たとえば次のような 3 重ループを含む Python プログラムを考えてみましょう。

def warshall_floyd(g):
    n = len(g)
    for j in range(n):
        for i in range(n):
            for k in range(n):
                g[i][k] = min(g[i][k], g[i][j] + g[j][k])
    return g[0][n - 1]

for ループの終了はインデントが減ることで表現されています。これによりブロックが簡潔に表現できるなどの利点はある一方で、複数レベルのブロックが一度に終わった時いくつのレベルが終わったか分かりづらいのが難点の一つでもあります。

ところで Ruby ではブロックの終了を end というキーワードで表現します。

n = g.length
n.times do |j|
  n.times do |i|
    n.times do |k|
      g[i][k] = [g[i][k], g[i][j] + g[j][k]].min
    end
  end
end

(注: 私は Ruby にまったく詳しくないのでこの素人っぽいコードにはあまり深く突っ込まないでください……) これなら3つのブロックが終了したことが明確ですね。

では Python にも end を導入すれば Ruby みたいにブロックの終了が見やすくなるのでは？という小学生並みの発想により、そのようなモジュールをメタプログラミング的に作ってみました、というのが今回の記事の内容になります。

使用例

何はともあれ、使用例です。ちなみにこのモジュールは PyPI に登録してあるので、pip install end で簡単にインストールできます。

import end

def func():
    for i in range(3):
        print(i)
    end
end

with open('a.txt') as f:
    print(f.read())
end

try:
    time.sleep(3)
except KeyboardInterrupt:
    print('Interrupted')
finally:
    print('Slept')
end

if 28 > 3:
    print('nya-n')
# end がここに無いので SyntaxError が投げられる！

スクリプト内で import end と書くことにより、そのスクリプト内でのみチェックが有効になり、ブロックの終了を end で示すことが強制されます。end をブロックの末尾に書き忘れたり、ブロックの末尾でないところに end を置くと SyntaxError が発生します。

end モジュールのソースコードは https://github.com/nya3jp/end に置いてあります。ちなみに end モジュール自身とそのユニットテストも end を利用して書かれていますので、もっと実用的(?)なコードでの使用例はそちらで見ることが出来ます。

実装方法

さて、以下では end モジュールをどのように実装したかを解説していきます。解説で用いる Python 処理系は CPython 3.3+ を仮定します。CPython 2.7 でも動作しますが、標準モジュールの名前などに多少変更があります。

基本的なアイデア

end なるキーワードは Python の文法に存在しないので、何も考えずに end と書くと NameError が発生します。まずはこれを回避したいわけですが、これは簡単で、単にグローバル変数として end という名前のものを定義してあげればいいだけです。

end = object()

def main():
    for i in range(3):
        print(i)
    end  # OK
end  # OK

変数 end の中身はなんでも構いません。そこで end.py という空のファイルを作成し、import end としてやることにします。すると end はモジュールを指すグローバル変数となり、先ほどの例と同様に end と書くことができるようになります。

# end.py
# ここにはまだ何も書いていない

# sample.py
import end

def main():
    for i in range(3):
        print(i)
    end  # OK
end  # OK

とりあえずこれだけでも見た目はそれっぽいのですが、このままでは end と書き忘れても何も怒られません。エラーが出ないと誰も end なんて書かないのであまり意味がないですね。

そこで、end.py に色々書くことにより、end モジュールが import された時に import 元のモジュール(上記の例の場合は sample.py) に対して文法チェックを走らせることにします。これは大まかに言って次の2ステップからなります。

end.py が import されるたびに import 元のソースコードを読み出す
読み出したソースコードを文法解析して end が正しく使われているかチェックする

それぞれ解説していきます。

import 元のソースコードを読む

まずは import 元のソースコードを読むことを考えましょう。

import のフック

Python の import 文はいろいろな方法でフックすることができます。大まかに言って 2 種類くらい方法があります。

ビルトイン関数 __import__() を置き換える
sys.meta_path に meta path loader を登録する

実は import 文は __import__() というビルトイン関数の呼び出しの単なるラッパーです。なので、この関数を自分で定義したものに置き換えることで import をフックすることができます。これが 1 番目の方法になります。

また、デフォルトの __import__() ビルトイン関数はユーザがカスタマイズ可能なフックの手段を用意しており、それが 2 番目の sys.meta_path を使う方法になります。こちらについてはわたしが Python Advent Calendar 2015 で書いた Hacking import mechanism という記事で解説しているので、気になる方はそちらもご覧ください。

今回は 1 番目の __import__() を置き換える方法を利用します。これは、デフォルトの __import__() は meta path loader を呼び出す前に sys.modules に既にモジュールがロードされているかチェックして既にロードされていればそれを返すようなキャッシュ機構を備えているため、2 番目の sys.meta_path による方法では end モジュールの 2 回目以降の import をフックするのが難しいからです。

というわけで肩慣らしに __import__() ビルトイン関数を置き換えてみましょう。

# first_hack.py
import builtins
import functools

original_import = builtins.__import__

@functools.wraps(original_import)
def my_import(name, *args, **kwargs):
    print('import is hacked! name=%r' % name)
    return original_import(name, *args, **kwargs)

builtins.__import__ = my_import

import os  # => import is hacked! name="os"

Python 3 ではビルトイン関数は builtins モジュールに入っています(Python 2 では __builtin__)。この中にある __import__() を自分で定義した関数で置き換えることで、あっさりと import 文を乗っ取ることができました。

スタックフレームの走査

import のフックが出来たので、次に import 元がどこかを調べます。これは import が関数呼び出しに帰着される以上、スタックフレームを辿れば分かるはずです。

スタックフレーム(あるいは単にフレーム)は Python インタプリタの実行中、各関数呼び出しの各レベルにおいて今どこが実行されておりどのような環境を持っているかを示しているものです。現在のスタックフレームは inspect.currentframe() 関数で取得することが出来て、バックリンクにより関数の呼び出し元を辿っていくことができます。

これを使って、我々の __import__() を呼び出したコードのファイル名と呼び出し元の行数を出力してみましょう。

# second_hack.py
import builtins
import functools
import inspect

original_import = builtins.__import__

@functools.wraps(original_import)
def my_import(name, *args, **kwargs):
    caller = inspect.currentframe().f_back
    print('called from %s:%d' % (caller.f_globals['__name__'], caller.f_lineno))
    return original_import(name, *args, **kwargs)

builtins.__import__ = my_import

import os  # => called from second_hack.py:15

inspect.currentframe() で手に入るフレームは my_import() 内を表すフレームなので、f_back を一度辿って呼び出し元のフレームを取得します。フレームからはそこにおけるグローバル変数テーブルが frame.f_globals で参照できるので、これを使って __name__ を読んでソースコードのパスを取得します。ついでに実行中の行番号も frame.f_lineno で取得可能です。フレームオブジェクトに入っている情報の一覧は inspect モジュールのドキュメントに書いてあります。

CPython バイトコードの解析

もう import 元のファイル名が手に入ったように見えますが、実はこれでは不完全です。というのも、__import__() をフックしているのが end.py だけであればこのままで十分なのですが、他にもフックしているコードがあった場合、我々の定義した __import__() の呼び出し元は import 文を実行したユーザーコードではなく、他のフックコードになってしまう可能性があるからです。実際、Python 3 でいろいろテストをしていたところ、__import__() を置き換えるコードが標準ライブラリ内にあり、単純に親フレームを見るだけでは正しく import 元を特定できない現象に遭遇しました。

ではスタックフレームの列から import を行っているユーザーコードを正しく特定するにはどうしたらいいでしょうか。これは、各スタックフレームで実行中のバイトコードを解析することで解決します。

スタックフレームには、そのフレームで実行中のバイトコードとインストラクションポインタ（いまどの命令を実行しようとしているかを示すカウンタ）が入っています。これを使って、今まさに import 文を実行しているフレームを探し出すのです。__import__() をフックするようなコードはその性質上 import 文ではなく関数呼び出しを実行しているはずなので、そのようなフレームはスキップすることができます。

CPython には dis というモジュールが用意されており、これを使ってバイトコードを逆アセンブルすることができます。試しに import 文だけを含むような関数 f() を逆アセンブルしてみましょう。

def f():
    import os

これは次のようなバイトコードになります。

 4           0 LOAD_CONST               1 (0)
             3 LOAD_CONST               0 (None)
             6 IMPORT_NAME              0 (os)
             9 STORE_FAST               0 (os)
            12 LOAD_CONST               0 (None)
            15 RETURN_VALUE

CPython のバイトコードに馴染みのない方も多いかと思いますが、IMPORT_NAME が import 文に対応する命令であり、そのパラメータ（オペランド）として “os” という文字列が渡されている、ということだけ分かれば十分です。

import 文に対応する命令が分かったので、あとはスタックフレームを走査してバイトコードを見つつ import end を実行しているフレームを探します。そのコードは以下のような感じになります(注: いろいろ端折っています)。

import dis
import inspect

def find_importer_frame():
    frame = inspect.currentframe()
    while frame:
        code = frame.f_code
        lasti = frame.f_lasti
        if code.co_code[lasti] == dis.opmap['IMPORT_NAME']:
            arg = code.co_code[lasti + 1] + code.co_code[lasti + 2] * 256
            name = code.co_names[arg]
            if name == 'end':
                break
        frame = frame.f_back
    return frame

スタックフレームが特定できれば、ソースコードを取得するのは inspect.getsource(frame) を呼び出すだけで終わりです。

ソースコードの解析

ソースコードが手に入ったので、あとはこれを解析して、end が正しく置かれているかをチェックします。

Python には標準で Python ソースコードを解析して AST (Abstract Syntax Tree) を出力する ast モジュールがバンドルされています。これを使うことでコードの解析を手軽に行なうことができます。

たとえばこのようなシンプルなソースコードを考えましょう。

def main():
    for i in range(3):
        print(i)
    print('done')

これは ast.parse() により次のような AST に変換されます。

Module(
  body=[
    FunctionDef(
      name='main',
      args=arguments(args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]),
      body=[
        For(
          target=Name(id='i', ctx=Store()),
          iter=Call(func=Name(id='range', ctx=Load()), args=[Num(n=3)], keywords=[]),
          body=[
            Expr(value=Call(func=Name(id='print', ctx=Load()), args=[Name(id='i', ctx=Load())], keywords=[])),
          ],
          orelse=[]),
        Expr(value=Call(func=Name(id='print', ctx=Load()), args=[Str(s='done')], keywords=[])),
       ],
       decorator_list=[],
       returns=None),
  ],
)

Module.body や FunctionDef.body がブロックの中身に対応しています。また、単に end と書くと AST では Expr(value=Name(id='end', ctx=Load())) と表現されます。よって、ブロックを見つけたらその次に end が来ること、またそれ以外のところで end が単体で現れないことを確認すれば終わりです。Python の文法要素すべてに対して場合分けをしたりしないといけないので地味に面倒な作業ではありますが、やるだけなので頑張ります。チェックに失敗した場合は SyntaxError を raise します。

以上で import 元の文法チェックが出来るようになりました！

だいたい完成

end.py が初めて import された時は、まず __import__() フックをインストールし、その場で呼び出し元の文法チェックを実行します。以降は import end が実行される度に __import__() フック内から文法チェックを呼び出します。

これでほとんど完成です。細かなところは端折っているので、詳しくは実際のソースコードをご覧ください: https://github.com/nya3jp/end/blob/master/end.py

落ち穂拾い

この end チェックは実は不完全で、正しく解析できないケースがあることが分かっています。一番の問題は一行にブロックを押し込めた時 (if a > 28: a = 3 など) の扱いで、Python の物理行・論理行の仕様などを考えると字句解析レベルの情報が必要になり、ast モジュールの出力だけではどうやっても正しく解析できません。

あと、そもそもソースコードがない場合はどうにもなりません。コンパイル済みのバイトコードを実行している時とか、REPL を実行している時とかですね。REPL で使えないのはちょっと悲しい。

まとめ

以上、Python に end を導入する話でした。end がソースコードの見やすさに与える影響はさておいて、Python における(変な)メタプログラミングのサンプルとしては多少は役に立つかな？と思います。

ちなみに実際に end.py やそのユニットテストを end を使って書いてみたわけですが、感想としては、まぁ確かにブロックの終了は分かりやすくなるものの、すごく野暮ったくなるなという印象でした。別に Ruby を dis っているわけではなくて、見慣れていないせいで Python のコードっぽく見えないというか。あとエディタのシンタックスハイライトはことごとく end を無視するのでそれもマイナスですね。

そして念のため最後にひとこと: end を実用してはいけません。これを仕事のコードに導入したりすると同僚との関係が悪化すること請け合いです :)

参考文献

The import system - Python 3 documentation
- import 文が __import__() への呼び出しで実装されていることや、sys.meta_path によるフックについて書かれています。
Hacking import mechanism
- わたしが去年書いた Python Advent Calendar 2015 の記事で、import フックについて解説しています。