pandasのquery()で指定できる@ローカル変数名の謎を追う

3年ぶりだ。

仕事でpandasを使っていて、こんな書き方ができることを知った。
ローカル変数を参照させている。下の例ではfoo列の値が123であるデータをフィルタできる。

bar = 123
df_filtered = df.query('foo == @bar')

公式ドキュメントにもはっきり書かれている。
pandas.DataFrame.query — pandas 2.1.2 documentation

式を渡す時点では単なる文字列なのに、pandasはどうやってbar = 123を参照しているのだろうか？
なんかスタックを辿る仕組みがあるんだろうと見当はつくが、具体的なHowの部分を知るため取材班はアマゾンの奥地へと飛んだ。

結論: inspect.stack()を使えばできる

※この検証はPython 3.9.10で行いました

inspect.stack()を使うことで、コールスタックを辿って各フレームの情報を得ることができる。
ここにローカル変数や関数の引数も含まれている。フレームの f_locals で取得できる。

import inspect


def callee():
    stack = inspect.stack()
    for level, frame in enumerate(stack):
        print(f'\n----- frame level={level} -----')
        print(frame.function)
        print(frame.frame.f_locals)


def caller_lv3():
    callee()


def caller_lv2(arg):
    var2 = arg
    caller_lv3()


def caller_lv1():
    var1 = 1
    caller_lv2(var1)


def main():
    caller_lv1()


main()

実行結果:

----- frame level=0 -----
callee
{'stack': [FrameInfo(frame=<frame at 0x104ded900, ...(省略)...}

----- frame level=1 -----
caller_lv3
{}

----- frame level=2 -----
caller_lv2
{'arg': 1, 'var2': 1}

----- frame level=3 -----
caller_lv1
{'var1': 1}

----- frame level=4 -----
main
{}

----- frame level=5 -----
<module>
{'__name__': '__main__', ...(省略)...}

pandasで上記を実行している箇所

ここでやっている。

github.com

ずっと上の方からスタックレベルを記録してきて、ターゲットの階層を認識した上でそのフレーム情報を取得しているようだ。
これを元に@部分を値へ置換しているということですね。
APIだけ見るとさらっとしてるけど、内部ではちょっとトリッキーなことをしている。