Cython0.27で導入された型ヒント・変数アノテーション対応を試す


この記事は[Python Advent Calendar](https://qiita.com/advent-calendar/2017/python_python)の1日目の記事です。

## 背景 ##

Cythonのバージョン0.27からPython 3.5で追加された型ヒント、3.6で追加された変数アノテーションを、CythonとしてもPythonとしても実行できるコード(Pure Python Mode)を記述する際の型情報として受け付けることができるようになりました。経緯などはこの[issue](https://github.com/cython/cython/issues/1672)を読むと分かります。

型ヒントは[PEP484](https://www.python.org/dev/peps/pep-0484/)で、変数アノテーションは[PEP526](https://www.python.org/dev/peps/pep-0526/)で提案されたものですが、要は下記のように関数の引数、返り値、ローカル変数などに型情報をアノテーションとして付与できるというものです。

```python
def add_one(x: float) -> float:
    one : float = 1.0
    return x + one
```

一方、Cythonでは型情報を与えて高速化する際に、下記のように型情報や返り値の型を前置するc言語などと同様の記法を採用していました。

```cython
cdef double c_add_one(double x):
    cdef double one = 1.0
    return x + one
```

Cython 0.27からはこれを下記のような形で記述することができるようになります。

```cython
import cython

@cython.cfunc
def cy_add_one(x: cython.double) -> cython.double:
    one : cython.double = 1.0
    return x + one
```

上記は後述するPure Python Modeとしての記述であるため、Cythonコードとして扱われた際には型情報が付与されるが、Pythonコードとして扱われた際には通常の型ヒント、変数アノテーションと同様の扱いをされるようになります。

## CythonのPure Python Mode ##

Cythonにはコード自体はPythonの文法の範囲内で記述し、最適化したい対象に対してデコレータを付与することなどによりCython上での最適化を図ることができる[Pure Python Mode](https://cython.readthedocs.io/en/latest/src/tutorial/pure.html)という記述方法があります。

例えば、オライリーのCython本8章にある下記コード例をCython対応にすることを考えます。

```python
# 高速化の対象となる関数(引数fをaからbまで積分)
def integrate(a, b, f, N=2000):
    dx = (a - b) / N
    s = 0.0
    for i in range(N):
        s += f(a + i * dx)
    return s * dx
```

このコードに対し、従来のcythonでは下記のような形で型情報を記述していました。

```cython
cdef double typed_integrate(double a, double b, f, int N=2000):
    cdef:
        int i
        double dx = (b - a) / N
        double s = 0.0
    
    for i in range(N):
        s += f(a + i * dx)
    return s * dx
```

しかし先にも述べたように、このコードはcdefキーワードを用いるなどCythonに特化しており、Pythonのコードとしては実行できません。
CythonのPure Python Modeではこのコードを次のように修正することでPythonコードとしても実行できるようにします。

```cython
%%cython -a
import cython

@cython.cfunc
@cython.returns(cython.double)
@cython.locals(a=cython.double, b=cython.double, N=cython.int, i=cython.int, dx=cython.double, s=cython.double)
def pure_integrate(a, b, f, N=2000):
    dx = (b - a) / N
    s = 0.0
    for i in range(N):
        s += f(a + i * dx)
    return s * dx
```

1行目の`%%cython -a`はJupyter Notebookのコードセルなどで使えるマジックコメントであり、`%%cython`を付与することでcythonコードをその場でコンパイルすることが出来ますが、上記コードはPythonコードとしても実行できるため`%%cython`の部分をコメントアウトしても実行することができます。なおマジックコメントの`-a`は最適化のための情報を出力するための引数です。

上記コードをぱっと見で思うこととして、デコレータだらけになっており見た目がよろしくないということが挙げられます。しかし！Cython0.27からは冒頭で述べたように、返り値の型を指定するための`@cython.returns`と引数やローカル変数の型を指定するための`@cython.locals`を省略して、次のように書くことができるわけです。


```cython
import cython

@cython.cfunc
def pure_integrate(a: cython.double, b: cython.double, f, N : cython.int = 2000) -> cython.double:
    i  : cython.int
    dx : cython.double = (b - a) / N
    s  : cython.double = 0.0
    for i in range(N):
        s += f(a + i*dx)
    return s * dx
```

cython.intやcython.doubleという型を書く部分が気にならなくもないですが、これでデコレータが減った分、大分スッキリした感がありますね。なお上記のように@cython.cfuncでcdefでの関数宣言扱いにした上でdefで本体を宣言する場合は方ヒントが使えますが、従来のようにcdefで関数宣言自体を書いた場合は型ヒントは使えないようでした。

## pyannotate ##

話は変わりますが、先日dropboxから[pyannotate](https://github.com/dropbox/pyannotate)という型ヒント情報の付与されていないコードに対して、自動的に推論した型ヒント情報のコメントを挿入するというツールがリリースされました。型ヒントコメントはPEP484にも記述のある、Python 2.7など型ヒントなどをサポートしていないバージョンに対して型ヒントを下記のようにコメントとして記述するものです。

```python
def add_one(x):
    # type: (float) -> float
    """add one to x"""
    one = 1.0
    return x + one
```

pyannotateではこの型ヒントコメントを自動で推論、挿入できます。
githubに記載のある[例](https://github.com/dropbox/pyannotate/tree/master/example)では型情報の付与されていないgcd関数に対しコメントを挿入する例が示されています。

```python:gcd.py
def main():
    print(gcd(15, 10))
    print(gcd(45, 12))

def gcd(a, b):
    while b:
        a, b = b, a%b
    return a
```

というコードに対し、このコードを読み込み、型情報を収集するためのdriver.pyを記述します。

```python:driver.py
from gcd import main
from pyannotate_runtime import collect_types

if __name__ == '__main__':
    collect_types.init_types_collection()
    with collect_types.collect():
        main()
    collect_types.dump_stats('type_info.json')
```

これを`python driver.py`として実行すると、実行したディレクトリ配下に`type_info.json`という下記内容が示されたjsonファイルが吐き出されます。

```json:type_info.json
[
    {
        "path": "gcd.py",
        "line": 1,
        "func_name": "main",
        "type_comments": [
            "() -> None"
        ],
        "samples": 1
    },
    {
        "path": "gcd.py",
        "line": 5,
        "func_name": "gcd",
        "type_comments": [
            "(int, int) -> int"
        ],
        "samples": 2
    }
]
```

type_info.jsonが存在する状態で、`pyannotate -w gcd.py`を実行するとgcd.pyが上書きされた上で、

```python:gcd.py
Refactored gcd.py
--- gcd.py        (original)
+++ gcd.py        (refactored)
@@ -1,8 +1,10 @@
 def main():
+    # type: () -> None
     print(gcd(15, 10))
     print(gcd(45, 12))
 
 def gcd(a, b):
+    # type: (int, int) -> int
     while b:
         a, b = b, a%b
     return a
Files that were modified:
gcd.py
```

とdiffが標準出力に出力されます。`-w`を外すとdiffのみを見るdry-runを実行することもできます。

上記は結局型ヒントのコメントなので直接Cythonの最適化とは関係ないのですが、githubレポジトリのReadmeの内容を見ると、TODOとして"Python 3 code generation."と書かれています。もしpyannotateによってこの先Python3コードが型情報付きで吐き出されるようになれば、先に述べたCythonの型情報として用いることで、より手軽に高速化が図れるようになるのではないか、そんな希望を抱かせますね。

## 速度比較 ##

そんな訳で先に例に挙げたコードの素のPythonバージョンと、Cythonで型情報を付与して高速化を図ったバージョンの速度比較を行いたいと思います。

まず素のPythonです。

```python
def integrate(a, b, f, N=2000):
    dx = (a - b) / N
    s = 0.0
    for i in range(N):
        s += f(a + i * dx)
    return s * dx
```

これを次のようにJupyter notebook上で実行すると、

```python
%%timeit f = lambda x: x + 1.0
integrate(0.0, 10.0, f)

# => 373 µs ± 19.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
```

ということで約373usです。

次にCythonで型を付与した場合を考えます。型としてはcython.intやcython.doubleを与えます。

```cython
%%cython -a

import cython

def pure_integrate(a: cython.double, b: cython.double, f, N : cython.int = 2000) -> cython.double:
    i : cython.int
    dx : cython.double = (b - a) / N
    s : cython.double = 0.0
    for i in range(N):
        s += f(a + i*dx)
    return s * dx
```

これを次のように実行すると、

```python
%%timeit f = lambda x: x + 1.0
pure_integrate(0.0, 10.0, f)

# => 214 µs ± 11.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
```

214usとなって約150usほど高速化されていることが分かります。

最後にcython.intなどではなく、単なるintやfloatで型ヒントが付与された場合を想定します。これはpyannotateが吐いた型情報が使える場合を想定しています。

```cython
%%cython -a

def pytype_integrate(a: float, b: float, f, N : int = 2000) -> float:
    i : int
    dx : float = (b - a) / N
    s : float = 0.0
    for i in range(N):
        s += f(a + i*dx)
    return s * dx
```

同様に下記のように実行すると、

```python
%%timeit f = lambda x: x + 1.0
pytype_integrate(0.0, 10.0, f)

# => 355 µs ± 9.46 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
```

！！！！！特に高速化されていない！！！！！

### 理由と思われるもの ###

[Pure Python Modeのドキュメント](https://cython.readthedocs.io/en/latest/src/tutorial/pure.html)に載っているコードを下記に引用すると、

```python
def func():
    # Cython types are evaluated as for cdef declarations
    x : cython.int               # cdef int x
    y : cython.double = 0.57721  # cdef double y = 0.57721
    z : cython.float  = 0.57721  # cdef float z  = 0.57721

    # Python types shadow Cython types for compatibility reasons
    a : float = 0.54321          # cdef double a = 0.54321
    b : int = 5                  # cdef object b = 5
    c : long = 6                 # cdef object c = 6
```

`Python types shadow Cython types for compatibility reasons`とあるので、intやlongはobject扱いになってしまうようです。
実際に`%cython -a`で吐き出されたcythonコンパイル結果のうち、関数の引数の型宣言に当たる部分を見ると

```c
static PyObject *__pyx_pw_46_cython_magic_a41350ee0647c8a4bbe2de265f286c79_1pytype_integrate(PyObject *__pyx_self, PyObject *__pyx_args, PyObject *__pyx_kwds) {
  double __pyx_v_a;
  double __pyx_v_b;
  PyObject *__pyx_v_f = 0;
  PyObject *__pyx_v_N = 0;
  PyObject *__pyx_r = 0;
  __Pyx_RefNannyDeclarations
  __Pyx_RefNannySetupContext("pytype_integrate (wrapper)", 0);
  {
```

となっており、引数のaとbはcレベルでのdoubleになっていますが、引数のNはintではなくPyObject*になっていることが分かります。一方、引数の型を`cython.int`で与えていた場合はcのintになっていました。

## まとめ ##

* CythonとしてもPythonとしても実行できるコードを書く際に、Cython0.27から型ヒントや変数アノテーションが使えるようになり、これによりデコレータを使わない簡潔な記述ができるようになった
* pyannotateの使い方を簡単に紹介し、将来的に型情報が自動で付与できる未来に対する希望を浮かべたが、速度比較を実施したところ、自動で付与されるであろう型では高速化が困難であるという結果が得られた


## おまけ ##

記事の趣旨からは外れますが、実際にcythonで高速化する場合は引数を関数ポインタとし、引数に与える関数自体をcレベルで書いてしまった方が速くなります。


```cython
%%cython -a
import cython

cdef double add_one(double x):
    cdef double one = 1.0
    return x + one

cdef double c_integrate(double a, double b, double (*f)(double), int N=2000):
    cdef:
        int i
        double dx = (b - a) / N
        double s = 0.0
    for i in range(N):
        s += f(a + i*dx)
    return s * dx

# cdef宣言された関数はpythonから直接呼べないのでpythonから呼ぶためのラッパー
def c_integrate_wrapper(double a, double b, int N=2000):
    return c_integrate(a, b, add_one, N)
```

これを実行すると、

```python
%%timeit
c_integrate_wrapper(0.0, 10.0)

# => 2.41 µs ± 51.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
```

素のpythonで373usだった処理が2.41usに減っており相当な高速化です。

この高速化と同様な記述を今回の型ヒントを用いてできないかとあがいた結果も下記に載せておきます。cy027_integrateの引数の関数ポインタにcy027_add_oneが渡せないので訳がわからないことになっています。

```cython
%%cython -a

import cython

ctypedef cython.double (*p_d2d_func)(cython.double x)

@cython.cfunc
def cy027_add_one(x: cython.double) -> cython.double:
    return x + 1.0

cdef double add_one(double x):
    return x + 1.0

@cython.cfunc
@cython.cdivision(True)
def cy027_integrate(a: cython.double, b: cython.double, f: p_d2d_func, N : cython.int = 2000) -> cython.double:
    i : cython.int
    dx : cython.double = (b - a) / N
    s : cython.double = 0.0
    for i in range(N):
        s += f(a + i*dx)
    return s * dx

def cy027_integrate_wrapper(a: cython.double, b: cython.double, N : int = 2000):
    # 下記はerrorになる(Cannot assign type 'double (double) except? -1' to 'p_d2d_func')
    # return cy027_integrate(a, b, cy027_add_one, N)
    return cy027_integrate(a, b, add_one, N)
```

## 参考 ##

* 同一テーマの記事 https://qiita.com/natsukium/items/99cffa5407aaf32a0b20
* 0.27での採用に至るまでの議論の[issue](https://github.com/cython/cython/issues/1672)
* What's new in Cython 0.27 http://blog.behnel.de/posts/whats-new-in-cython-027.html
* PEP 484 https://www.python.org/dev/peps/pep-0484/
* PEP 526 https://www.python.org/dev/peps/pep-0526/

## 次回 ##

明日は[@sin_tanaka](https://qiita.com/sin_tanaka)さんです。