Python でパイプ演算子を使いたい <2>
ネタ記事です。/ This is a joke post which makes no practical sense.
過去にこんなエントリを書いた。
R では パイプ演算子 %>%
を使って連続した処理を記述できる。式に含まれる
x
, y
, z
は非標準評価 (NSE) によって data.frame
の列として解決される。
# R (magrittr + dplyr) df %>% mutate(x = y + z) %>% group_by(x) %>% summarize_each(funs(sum))
Python (pandas
) ではほぼ同じ処理をメソッドチェインを使って書ける。チェインとパイプ演算子でどちらが読みやすいかは好みの問題だと思うものの、式の中に 何度も df
が出てくるのはちょっとすっきりしない。
# Python (pandas) df.assign(x=df['y'] + df['z']).groupby('x').sum()
上の式をこんな風に書けるとうれしい。
df >> assign(x=y+z) >> groupby(x) >> sum()
前のエントリでは マジックメソッドと関数定義で近いことをやったが、以下のような課題がある。
- パイプ演算子で処理したいメソッドを、それぞれ関数として別に定義しなければならない
- R の非標準評価のようなことはできない (ある変数を別の環境で評価することができない)
前者は手間をかければなんとかなるが、後者については NameError
が発生してしまうためどうしようもない。
df.assign(x=y+z)
# NameError: name 'y' is not defined
上のような表現を式に含めるには、 lambda
を使って無名関数にすればよい。lambda
を使うと、下のように定義されていない関数 / 変数を式に含めることができる。
lambda: df >> assign(x=y+z) # <function <lambda> at 0x102310620>
当然、この無名関数をそのまま評価すると NameError
が発生する。定義した無名関数の中身を調べて、妥当な処理に置き換えてやればよい。
Python 標準では ast
モジュールを使って 抽象構文木 (AST) の探索や置換ができる。が、今回はそこまで複雑なことをやるわけではないので、バイトコードを直接書き換えられるとうれしい。
バイトコードの置換について簡単な方法がないか調べたところ CodeTransformer
というパッケージを見つけた。今回はこれを使いたい。
CodeTransformerとは
バイトコードに対するパターンマッチ⇨置換が簡単にできる。詳しい使い方はドキュメントを。
まずは関数がバイトコードとしてどのように表現されているのか確かめたい。これは標準の dis
モジュールを使うのが簡単。
import dis dis.dis(lambda x: x + 2) # 1 0 LOAD_FAST 0 (x) # 3 LOAD_CONST 1 (2) # 6 BINARY_ADD # 7 RETURN_VALUE
CodeTransformer
を使うと適当なパターンにマッチするバイトコードを置換することができる。ドキュメントの例にある通り、バイトコード中の加算 (BINARY_ADD
) を乗算 (BINARY_MULTIPLY
) に置き換える CodeTransformer
は以下のようになる。
import codetransformer codetransformer.__version__ # '0.6.0' from codetransformer import CodeTransformer, pattern from codetransformer.instructions import BINARY_ADD, BINARY_MULTIPLY class Multiply(CodeTransformer): @pattern(BINARY_ADD) def _add2mul(self, add_instr): yield BINARY_MULTIPLY().steal(add_instr)
pattern
デコレータで指定したバイトコードにマッチした時、対応したメソッドが呼び出される。呼び出されるメソッドから置換後のバイトコードを返せばよい。
作成した CodeTransformer
を使って関数を置換する。
mul = Multiply()(lambda x: x + 2) mul(5) # 10
Multiply
によって置き換えられた後のバイトコードを確認すると、BINARY_ADD
が BINARY_MULTIPLY
に置換されていることが確かめられる。
dis.dis(mul) # 1 0 LOAD_FAST 0 (x) # 3 LOAD_CONST 0 (2) # 6 BINARY_MULTIPLY # 7 RETURN_VALUE
パイプ演算子を使いたい
ここから、pandas
と組み合わせて使ってみる。まずはデータを準備する。
import pandas as pd df = pd.DataFrame({'A': [1, 2, 3, 4, 5, 6], 'B': list('ABABAB'), 'C': [1, 2, 3, 1, 2, 3]}) df # A B C # 0 1 A 1 # 1 2 B 2 # 2 3 A 3 # 3 4 B 1 # 4 5 A 2 # 5 6 B 3
パイプ演算子を処理するには、lambda: df >> head()
のバイトコードを lambda: df.head()
のバイトコードに置換すればよい。置換前、置換後のバイトコードをそれぞれ確認すると、
# 置換前 dis.dis(lambda: df >> head()) # 1 0 LOAD_GLOBAL 0 (df) # 3 LOAD_GLOBAL 1 (head) # 6 CALL_FUNCTION 0 (0 positional, 0 keyword pair) # 9 BINARY_RSHIFT # 10 RETURN_VALUE # 置換後 dis.dis(lambda: df.head()) # 1 0 LOAD_GLOBAL 0 (df) # 3 LOAD_ATTR 1 (head) # 6 CALL_FUNCTION 0 (0 positional, 0 keyword pair) # 9 RETURN_VALUE
比べてみると、以下のような変換を行えばよさそうだ。
- 置換前のパターン:
LOAD_GLOBAL → CALL_FUNCTION → BINARY_RSHIFT
- 置換後のパターン:
LOAD_ATTR → CALL_FUNCTION
連続したパターンとマッチさせるには、pattern
デコレータと対応するメソッドに複数の引数を指定すればよい。
from codetransformer.instructions import BINARY_RSHIFT, LOAD_GLOBAL, LOAD_ATTR, CALL_FUNCTION class Pipify(CodeTransformer): @pattern(LOAD_GLOBAL, CALL_FUNCTION, BINARY_RSHIFT) def _pipe(self, load, call, rshift): # LOAD_GLOBAL を LOAD_ATTR に置換 yield LOAD_ATTR(load.arg) yield call # BINARY_SHIFT は無視
これにパイプ演算を含む無名関数を渡すと、期待通り動いているように見える。
Pipify()(lambda: df >> head())() # A B C # 0 1 A 1 # 1 2 B 2 # 2 3 A 3 # 3 4 B 1 # 4 5 A 2 dis.dis(Pipify()(lambda: df >> head())) # 1 0 LOAD_GLOBAL 0 (df) # 3 LOAD_ATTR 1 (head) # 6 CALL_FUNCTION 0 (0 positional, 0 keyword pair) # 9 RETURN_VALUE
もっとも、上のクラスは非常に単純な条件でしか動かない。たとえば、パイプで接続される関数(メソッド)に引数がある場合、pattern
の定義とマッチしなくなるため置換が行われない。
# NG Pipify()(lambda: df >> head(2))() # NameError: name 'head' is not defined dis.dis(lambda: df >> head(2)) # 1 0 LOAD_GLOBAL 0 (df) # 3 LOAD_GLOBAL 1 (head) # 6 LOAD_CONST 1 (2) # 9 CALL_FUNCTION 1 (1 positional, 0 keyword pair) # 12 BINARY_RSHIFT # 13 RETURN_VALUE
もう少しマシな CodeTransformer
の定義を考える。パイプ演算子の記法から、 LOAD_GLOBAL
から次の BINARY_RSHIFT
までを 1 単位とした pattern
定義でマッチすればよいような気がする。
CodeTransformer
の定義は以下のようになる。(~BINARY_RSHIFT)[plus]
は BINARY_SHIFT
以外の任意のコードの繰り返しとマッチする。
from codetransformer import plus class Pipify(CodeTransformer): @pattern(LOAD_GLOBAL, (~BINARY_RSHIFT)[plus], BINARY_RSHIFT) def _pipe(self, load, *args): # メソッド呼び出しの処理 if hasattr(self, 'first_load'): yield LOAD_ATTR(load.arg) else: # 最初の LOAD_GLOBAL は置き換えない self.first_load = load yield load # メソッドへの引数の処理 for arg in args[:-1]: if isinstance(arg, LOAD_GLOBAL): # 関数の一番最初の LOAD_GLOBAL から解決 # この定義では、パイプ演算中に更新されたデータやグローバル変数には # アクセスできない yield self.first_load yield LOAD_ATTR(arg.arg) else: yield arg # 末尾は BINARY_RSHIFT なので無視
コメントで "メソッドへの引数の処理" とある部分は、 (記載のとおり限定的だが) R の非標準評価に近い処理になる。
この定義を使うと、上で失敗したメソッド呼び出しも処理できる。
Pipify()(lambda: df >> head(2))() # A B C # 0 1 A 1 # 1 2 B 2 dis.dis(Pipify()(lambda: df >> head(2))) # 1 0 LOAD_GLOBAL 0 (df) # 3 LOAD_GLOBAL 0 (df) # 6 LOAD_ATTR 1 (head) # 9 LOAD_CONST 0 (2) # 12 CALL_FUNCTION 1 (1 positional, 0 keyword pair) # 15 RETURN_VALUE
いちいち Pipify()(lambda: df >> ...)()
と書くのは面倒なので、パイプ演算を開始する関数 p
を定義する。この関数を使うとパイプ演算は p(lambda: df >> ...)
のように書ける。
def p(expr): return Pipify()(expr)()
以下、いくつかのパターンを試す。それぞれ、冒頭にコメントとして記載した処理と同じ結果になる。
# df.head() p(lambda: df >> head()) # A B C # 0 1 A 1 # 1 2 B 2 # 2 3 A 3 # 3 4 B 1 # 4 5 A 2 # df.head(2) p(lambda: df >> head(2)) # A B C # 0 1 A 1 # 1 2 B 2
パイプ演算中の B
は NameError
とはならず df.B
として解決される (非標準評価もどき)。
# df.groupby(df.B).sum() p(lambda: df >> groupby(B) >> sum()) # A C # B # A 9 6 # B 12 6
もう少し複雑な処理もできる。.assign
で df.A + df.C
の結果からなる X
列を新規で作成し、B
列の値ごとに集約して合計を取る。
# df.assign(X=df.A+df.C).groupby(df.B).sum() p(lambda: df >> assign(X=A+C) >> groupby(B) >> sum()) # A C X # B # A 9 6 15 # B 12 6 18
改行を含む場合は括弧で囲む。
p(lambda: (df >> assign(X=A+C) >> groupby(B) >> sum())) # 略
まとめ
CodeTransformer
でバイトコードを置換することで、パイプ演算のような処理が書ける。また、非標準評価に近いこともできる。- R を使ったほうがいい。