2010/11/25 Pythonでtracebackに出現するファイル名や行番号を書き換える

Category: 'python', 'testing'

メタ清水川です。ここのところなぜかメタなコードを作ることが多いんですが、昨日の Pythonでデータ定義行のファイル名と行番号を手に入れる に引き続き今回も「Pythonでtracebackに出現するファイル名や行番号を書き換え」てみたいと思います。前回に引き続き、やりたいと思ったのは以下のような理由からです。

データ定義に従ってテストするコードを書いて、テストが失敗したらそのデータ定義の行を表示したい。

前回の Pythonでデータ定義行のファイル名と行番号を手に入れる のエントリで、わざと関数呼び出しを発生させてそのときのコールスタックというかframe情報を使ってデータの定義位置(ファイル名と行番号)を取得する事は出来ました。前回はこの情報を以下のように表示させました:

AssertionError: Expect=200 but Actual=404 at 'tests.py' line 15

このように表示させるためには、self.assertEqual()などのTestCaseクラスのアサートメソッドには必ずあるmsg引数に、ファイル名や行番号を表示するよう仕込む必要があります。しかしこれは テストを書く度にわざわざ仕込んであげないといけない と言うめんどくさい手間が増えてしまうので、出来ればそんな仕込みは無しにしたいところです。

ということで、次は手に入れた定義位置の情報をtracebackに表示させたいわけです。

例えば以下のようなコード(cook.pyとします)があるとします。このプログラムはcook関数が受け取った文字列を使って何か処理をする関数_cookerを生成して返しています(が、line6で必ずエラーになります)。そのような関数オブジェクトをline13,14で2つ生成して、line17でどちらかをランダムに実行しています。

# -*- coding: utf-8 -*-
import random

def cook(base):
    def _cooker(addon):
        dish = base + x              # line 6
        # some program put here...

    return _cooker


cooks = [
    cook('spam'),                    # line 13
    cook('ham'),                     # line 14
]
f = random.choice(cooks)
f('egg')                             # line 17

cook.pyを実行したときにトレースバックに出てくる情報は以下のような内容になります:

$ python cook.py
Traceback (most recent call last):
  File "cook.py", line 17, in <module>
    f('egg')
  File "cook.py", line 6, in _cooker
    dish = base + x
NameError: global name 'x' is not defined

この上記のトレースバックを見たときに分かるのは cook.py 6行目の_cooker関数内でエラーがあったことです。しかし、実際のプログラムでは(昨日のテストのコードのように)line13が問題だったのかline14が問題だったのかを知りたいのですが、このトレースバックからそれは分かりません。実際、メタクラス使ってたりフレームワークとか作ったり使ってたりすると、時々こういうシーンに遭遇します。

そこで、_cookerの関数名がtracebackに表示されるときに spam とか ham とか表示させる方法はないか、ということで色々調べてみると、関数オブジェクトの属性を書き換えれば良さそうだと言うことが分かります。例えば:

_cooker.func_name = 'bacon' # Python 2.x
_cooker.__name__ = 'bacon' # Python 3.x (2.6以降)

上記のようにして関数名を書き換える訳です。しかし手元の環境ではこれではtracebackの表示は変わりませんでした。そこでもう少し調べると以下のように書き換えれば良さそうな事が分かりました:

_cooker.func_code.co_name = 'bacon'

しかしこの書き換えは出来ません。func_code.co_nameは読み取り専用属性なので書き換えられません。このあたりの読み取り専用属性かどうかといった情報は Python言語リファレンスのデータモデルの節 (日本語訳)に 読み出し専用の特殊属性: co_name は関数名を表します というふうに書かれています。

仕方がないので、関数オブジェクト(_cooker)とそのバイトコードが格納されているコードオブジェクト(_cooker.func_code)まるごと作り直してしまいましょう。ということで以下のように、frameから行番号を取得して、コードオブジェクトを作り直して関数名に行番号とbase名を仕込んでみました。

# -*- coding: utf-8 -*-
import sys
import random
import types

def chsig(func, **replaces):
    code_arg_names = (
        'co_argcount', 'co_nlocals', 'co_stacksize', 'co_flags', 'co_code',
        'co_consts', 'co_names', 'co_varnames', 'co_filename', 'co_name',
        'co_firstlineno', 'co_lnotab', 'co_freevars', 'co_cellvars',
    )
    old_code = func.func_code
    code_args = [replaces.get(n, getattr(old_code, n)) for n in code_arg_names]
    new_code = types.CodeType(*code_args)

    func.func_code = new_code
    return func

def cook(base):
    f = sys._getframe().f_back
    lineno = f.f_lineno
    def _cooker(addon):
        dish = base + x  # line 23
        # some program put here...

    name = "%s generated by '%s(%d)'" % (_cooker.func_name, base,lineno,)
    chsig(_cooker, co_name=name)
    return _cooker


cooks = [
    cook('spam'),        # line 32
    cook('ham'),         # line 33
]
f = random.choice(cooks)
f('egg')                 # line 36

本当はファイル名もやりたいところですが、説明を分かりやすくするために今回は省略。実行した結果は以下のようになりました:

$ python cook2.py
Traceback (most recent call last):
  File "cook2.py", line 36, in <module>
    f('egg')
  File "cook2.py", line 23, in _cooker generated by 'ham(33)'
    dish = base + x
NameError: global name 'x' is not defined

_cooker generated by 'ham(33)' が表示されるようになったおかげで、エラーが発生したときにも非常に分かりやすくなりました。これでエラーがどこで起きてるか分からないんじゃ!などというつまらないことに引っかからずにデバッグできるようになりましたね!ヒャッホウ!