🕌

Pythonの@オペレータについて調べてみた

に公開

今回はPythonで@をオペレータとして利用する方法を調べてみました。Python、特にNumPyを利用して行列を取り扱っていると、計算の時に@をオペレータとして使うことがしばしばあります。というか私はNumPy以外で@をオペレータとして使っていることをみたことはありませんでした。+や-など一般的なさん術オペレータはよく利用しますが@はどのようにすれば使えるようになるのかふと気になり調べてみました。

NumPyにおける@オペレータ

NumPyでは@は行列積を計算するためのオペレータとなっています。例えば以下にAとB二つの行列が会った時に、その積を計算できます。以下の例はBを単位行列にしているので、AとBの行列積の結果はAと一致します。

numpy_at_operator.py
import numpy as np

A = np.array(
    [
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9],
    ]
)
B = np.identity(3)

print(f"{A @ B =}")
uv run numpy_at_operator.py

# 結果
A @ B =array([[1., 2., 3.],
       [4., 5., 6.],
       [7., 8., 9.]])

実装はどうなっているのか?

調べたところ、__matmul__という特殊関数をクラスで定義することにより演算子として定義できるようです。例えば__add__という演算子を定義すればそのクラスに対して加算+の演算子が定義できます。この特殊関数を低議すことにより、そのクラスのインスタンス同士の計算に@オペレータを採用することができるということになります。特殊関数の名前が__matmul__という名前からも想像できますが、おそらくこのオペレータは行列積を計算するために導入された物だと思いますが、その実態は行列積である必要はありません。例えばpathlib.Pathではhoge / fugaみたいな形で割り算の演算子であるスラッシュを利用するとパスを連結できるので、使い方自体は自由です!

それでは早速@オペレータを使えるようにしてみます。今回は以下のように変数xをメンバーにもつHogeクラスがあり、そのクラスに@オペレータを実装します。@オペレータの役割は二つのHogeインスタンスの変数xの掛け算を返すものとします。

matmul_impl.py
class Hoge:
    def __init__(self, x):
        self.x = x

    def __matmul__(self, other):
        return self.x * other.x


a1 = Hoge(10)
a2 = Hoge(20)
print(f"{a1 @ a2 =}")

今回大事なのは__matmul__関数を定義しているところと、その引数にotherをおいているところです。今回型チェックは入れてないですが、otherには同じHoge型を想定しています。実装のように、二つのインスタンスの変数xの積を返します。

def __matmul__(self, other):
    return self.x * other.x

それでは実際にこれを実行してみましょう。まず初めに、仮にHogeから__matmul__の実装を無くした状態で実行したらどうなるかやってみます。想像の通り@オペレータの実装がされていないので、Hogeクラス同士の@演算は存在しないと怒られます。

uv run matmul_impl.py  # 裏で実装を消して実行しています

# 結果
Traceback (most recent call last):
  File "/Users/user/Documents/Blog/blog_materials/python_operator/at_operator/main.py", line 11, in <module>
    print(f"{a1 @ a2 =}")
             ~~~^~~~
TypeError: unsupported operand type(s) for @: 'Hoge' and 'Hoge'

次に__matmul__を実装した状態で実行してみると以下のようになりました。結果をみると、実装している通り@オペレータによって掛け算が実装できていることが確認できました。

uv run matmul_impl.py  # __matmul__の実装は復元しています

# 結果
a1 @ a2 =200

@=の定義の仕方

ちなみに、他の演算子と同様に__imatmul__のようにiを頭につけることで@=のように代入演算子を定義することができます。先ほどの例をちょっと改造してみました。

imatmul_impl.py
class Hoge:
    def __init__(self, x):
        self.x = x

    def __matmul__(self, other):
        return self.x * other.x

    def __imatmul__(self, other):
        return Hoge(self.x * other.x)


a1 = Hoge(10)
a2 = Hoge(20)
a3 = Hoge(30)
a3 @= a2
print(f"{a1 @ a2 =}")
print(f"{a3.x=}")

先ほどの例に追加して__imatmul__を追加しました。これによりa @= bとするとaにa@bの結果が代入されるようになります。今回の例で言うとa3@a2の結果がa3に代入されるようになります(計算結果としては30*20=600となります)。
それではこちらも実行してみましょう。

uv run imatmul_impl.py

# 結果
a1 @ a2 =200
a3.x=600

こちらも思ったような結果になっていますね。

まとめ

今回はふと気になったのでPythonでの@オペレータについて調べてみました。演算子の挙動を独自で実装するユースケースって意外と少ないかなと思いつつ、その中でも特に@は自前実装している人を見かけたことがなかったので気になって調べてみました。Pythonでは演算子定義をはじめとして様々な特殊関数が利用できるので、ぜひ皆さんも調べてみてください。

Discussion