mypyc_ipython: mypycを気軽に動かせるIPythonマジックコマンドの紹介と実装解説

最近mypyc(型ヒントのついたPythonのコードからC拡張を生成するコンパイラ)のコードを読んでいたのですが、Cythonの %%cython マジックコマンドみたいに手軽に確認したいなということでマジックコマンドを実装してみました。CythonやPythonとサクッと性能を比較したいときにぜひ使ってみてください。

github.com

ここでは実装時のメモと、mypyc・cython・(pure) pythonのマイクロベンチマークの結果を残しておきます。

実装について

IPythonマジックコマンドの実装自体は↓のページに書かれています。

Defining custom magics — IPython 7.25.0 documentation

ポイントとしては、 %loadext <module name> を呼び出したときに、そのモジュールに定義されている load_ipython_extension() 関数が呼び出されます。そこに次のように関数を定義しておくと読み込まれます。

def load_ipython_extension(ip):
    """Load the extension in IPython via %load_ext mypyc_ipython."""
    from ._magic import MypycMagics

    ip.register_magics(MypycMagics)

MypycMagics の実装のポイントとして、次のような手順でコンパイルから読み込みが実行されます。

%%mypyc コードセルの中身をファイルに書き出す。
1. 書き出し先は、 $IPYTHONDIR/mypyc にしています。
2. Cythonにも似たようなマジックコマンドがあるのですが、そちらは $IPYTHONDIR/cython に書き出していたのでそれを真似しました。
3. コード保存時のファイル名は、こんな感じのhashを計算して使用していて、全く同じコードが再度定義された場合はコンパイルが省略されます。それだと困る場合には --force オプションを指定します。このあたりもCythonのマジックコマンドと完全に同じです。
mypycが提供する mypyc.build.mypycify(paths: List[str]) -> Extension 関数を使ってsetuptoolsのExtensionオブジェクトを取得。C拡張のコードを生成・コンパイル。
生成された .so モジュールを読み込む。
1. CythonだとPython2対応もあるため、 imp.load_dynamic() が使われていますが、 mypyc_ipython はPython3以降のみをサポートするため、importlibを使っています。
2. imp.load_dynamic() に相当する処理はこんな感じで実装できます。
module内のattributesをすべて読み込み
1. https://github.com/c-bata/mypyc_ipython/blob/cefc1ce28559194ea4de6d2f686385d84b0a970e/mypyc_ipython/_magic.py#L38-L50
2. これまでgunicornなどに送ったpatch などで似たような処理を実装したことはあったのですが、 __all__ をチェックするのを省いていたことに気づきました。たしかに見ない理由はあまりないので、今からgunicornの方にも修正patch送ってもいいかも。

あとは通常のPythonの関数と同じようにIPython上で実行できます。

マイクロベンチマーク

fibonacci数を計算するコードを2種類用意してみました。

再帰で実装

In [1]: %load_ext mypyc_ipython

In [2]: %%mypyc
   ...: def my_fibonacci(n: int) -> int:
   ...:     if n <= 2:
   ...:         return 1
   ...:     else:
   ...:         return my_fibonacci(n-1) + my_fibonacci(n-2)
   ...:

In [3]: my_fibonacci(10)
Out[3]: 55

In [4]: def py_fibonacci(n: int) -> int:
   ...:     if n <= 2:
   ...:         return 1
   ...:     else:
   ...:         return py_fibonacci(n-1) + py_fibonacci(n-2)
   ...:

In [5]: py_fibonacci(10)
Out[5]: 55

In [6]: %load_ext cython

In [7]: %%cython
   ...: cpdef int cy_fibonacci(int n):
   ...:     if n <= 2:
   ...:         return 1
   ...:     else:
   ...:         return cy_fibonacci(n-1) + cy_fibonacci(n-2)
   ...:

In [8]: cy_fibonacci(10)
Out[8]: 55

In [9]: %timeit py_fibonacci(10)
10.3 µs ± 30.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [10]: %timeit my_fibonacci(10)
848 ns ± 5.82 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [11]: %timeit cy_fibonacci(10)
142 ns ± 1.18 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

In [12]:

今回のマイクロベンチマークではPureなPythonに比べ、mypycはおよそ1桁速くなっています。Cythonはさらに1/8の実行時間になりました。 (cythonのほうで int を使いましたが、Pythonのintとは意味が違うので、long long とかで比較するほうがよかったかもしれません。)

forループで実装

In [1]: %load_ext mypyc_ipython

In [2]: %load_ext cython

In [3]: %%mypyc
   ...: 
   ...: def mypyc_fib(n: int) -> float:
   ...:     i: int
   ...:     a: float = 0.0
   ...:     b: float = 1.0
   ...:     for i in range(n):
   ...:         a, b = a + b, a
   ...:     return a
   ...: 

In [4]: def py_fib(n: int) -> float:
   ...:     i: int
   ...:     a: float = 0.0
   ...:     b: float = 1.0
   ...:     for i in range(n):
   ...:         a, b = a + b, a
   ...:     return a
   ...: 

In [5]: %%cython
   ...: cpdef cython_fib(int n):
   ...:     cdef int i
   ...:     cdef double a = 0.0, b=1.0
   ...:     for i in range(n):
   ...:         a, b = a + b, a
   ...:     return a
   ...: 
   ...: 


In [6]: timeit py_fib(10)
627 ns ± 3.81 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [7]: timeit mypyc_fib(10)
891 ns ± 26.5 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [8]: timeit cython_fib(10)
44.5 ns ± 0.092 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)