PySnooper の実装の仕組みを調べる

概要

cool-RR/PySnooper: Never use print for debugging again

PySnooper というデバッグ用の Python ライブラリについて、どのように実装されているかソースコードを読んで調べた結果をまとめます。

PySnooper とは

関数定義にデコレーターをつけるだけで、実行中にソースコードの各行でローカル変数がどんな値を取っているか確認することができます。
例えば、以下のようなスクリプトを実行すると：

from pysnooper import snoop

@snoop()
def func(x):
    a = x + 1
    b = a * 4
    return a + 2

func(1)

標準出力に以下のようなトレース結果が表示されます：

Starting var:.. x = 1
23:06:03.976593 call         4 def func(x):
23:06:03.976593 line         5     a = x + 1
New var:....... a = 2
23:06:03.976593 line         6     b = a * 4
New var:....... b = 8
23:06:03.976593 line         7     return a + 2
23:06:03.977593 return       7     return a + 2
Return value:.. 4

ソースコード探訪

snoop 関数

PySnooper でユーザーが使用する API は基本的にデコレーター関数の snoope だけです。
というわけで、まず snoope が実装されているソースコードから見ていきましょう。

最初の数行は、トレース結果の出力先（stderr など）を指定するためのもので、本筋とはあまり関係が無いので無視します。

その直後の decorate 関数の定義を見ると、ほとんどの処理を Tracer クラスに委譲していることがわかります：

    def decorate(function):
        target_code_object = function.__code__
        tracer = Tracer(target_code_object=target_code_object, write=write,
                        truncate=truncate, variables=variables, depth=depth,
                        prefix=prefix, overwrite=overwrite)

        def inner(function_, *args, **kwargs):
            with tracer:
                return function(*args, **kwargs)
        return decorator.decorate(function, inner)

Tracer クラス

__enter__, __exit__ 関数

上述の decorate 関数の中では、 Tracer クラスのオブジェクトが with 構文で使われています。
つまり、デコレートする関数を呼び出す前後に Tracer クラスの __enter__, __exit__ 関数がそれぞれ呼び出されるということです。
そこで、これらの関数の定義を確認すると、 __enter__ で sys.settrace を使ってトレース関数を設定し、 __exit__ でそれを元に戻していることがわかります：

    def __enter__(self):
        self.original_trace_function = sys.gettrace()
        sys.settrace(self.trace)

    def __exit__(self, exc_type, exc_value, exc_traceback):
        sys.settrace(self.original_trace_function)

sys.settrace について

sys — System-specific parameters and functions — Python 3.7.3 documentation

sys.settrace(trace) を実行すると、システムのトレース関数として trace を設定します。その後以下のいずれかのイベントが発生したときに trace が呼び出されます：

• call : 任意の関数を呼び出すとき
• line : ソースコードの任意の1行を実行するとき
• return : 任意の関数またはコードブロックが値を返すとき
• exception : 例外が発生するとき
• opcode : opcode を実行するとき

このとき、 trace は3つの引数 frame, event, arg と共に呼び出されます：

• frame : 呼び出し時のスタックフレームの状態を表すオブジェクト。このドキュメントに記載されている frame 型の属性をもつ
• event : 上述のいずれかのイベントを表す文字列
• arg : イベントが return のときは関数の返り値、イベントが exception のときはタプル (exception, value, traceback)

trace 関数

トレースすべき関数か判定

snoope(depth=k) でデコレータされた関数の呼び出しと、そこから k 層分のスタックだけがトレースの対象になります。
すなわち、以下のいずれかの条件を満たす場合だけトレースが行われます：

• 現在実行中の関数 frame.f_code とデコレータされた関数 self.target_code_object が一致する
• 現在のスタック frame から呼び出し元 f_back を k 回以下辿ることで、デコレータされた関数に到達できる

変更された変数を表示

以下の部分で、前回トレース時の変数の状態 old_local_reprs と現在の変数の状態 local_reprs を取得します：

        self.frame_to_old_local_reprs[frame] = old_local_reprs = \
                                               self.frame_to_local_reprs[frame]
        self.frame_to_local_reprs[frame] = local_reprs = \
                               get_local_reprs(frame, variables=self.variables)

frame.f_locals.items() によって、現在のスコープ中にあるローカル変数の名前と値を全て取得できます。
また、 self.variables は snoop(variables=(x, y)) のようにオプション指定されたスコープ外の変数を表します。

old_local_reprs と local_reprs が得られたら、これらを比較して変更または新規に追加された変数を特定し、結果を出力します。

実行中のソースコードの行を表示

frame からソースコードを取得する処理は get_source_from_frame 関数で定義されています。

キャッシュの参照や、 IPython の場合その他の例外処理を除いて、重要な部分だけに着目すると：

module_name = frame.f_globals.get('__name__')
loader = frame.f_globals.get('__loader__')
source = loader.get_source(module_name)

これが source = source.splitlines() で行ごとのリストに変換して返されます。その後 trace 関数の中で以下のようにして実行中の行だけを抽出します：

        line_no = frame.f_lineno
        source_line = get_source_from_frame(frame)[line_no - 1]

トレース関数の引数から取得できる情報を実験的に確認

PySnooper の挙動を大幅に簡略化して再現し、実際にソースコードやローカル変数を取得できるか確認しましょう。
以下の Python スクリプトを保存して実行します：

import sys
    
def func(x):
    a = x + 1
    b = a * 4
    return a + 2

def trace(frame, event, arg):
    if event == 'line':
        module = frame.f_globals.get('__name__')
        loader = frame.f_globals.get('__loader__')
        source = loader.get_source(module)
        lineno = frame.f_lineno
        line = source.splitlines()[lineno - 1]
        print(line)

    if event == 'return':
        print(frame.f_locals.items())

    return trace

sys.settrace(trace)
func(1)

すると標準出力に以下の結果が print され、確かにソースコードやローカル変数を取得できることがわかります：

    a = x + 1
    b = a * 4
    return a + 2
dict_items([('x', 1), ('a', 2), ('b', 8)])
dict_items([])

まとめ

• snoop でデコレートされた関数の呼び出し前に sys.settrace を設定し、呼び出し後に元に戻す。
• トレース関数の frame 引数から、現在のスタックフレームの状態を取得する。
  • デコレートされた関数から k 層以内のスタックか否か
  • 実行中の関数が定義されたソースコード
  • 実行中のソースコードの行
  • ローカル変数の名前と値
• 実行中の行を表示
• ローカル変数に変更があれば表示