Chainer に型ヒントをつける

この記事は Chainer/CuPy Advent Calendar 2018 の六日目です。前回は @marshi さんの「Chainerで確率分布が扱えるようになりました。chainer.distributionsの紹介」でした。

最近、Chainer に型ヒント (type hint) をつけるために色々と試行錯誤をしている。これによって、Chainer を使ったソフトウェアの型検査ができるようにしたり、PyCharm などの IDE でのオートコンプリートが機能するようにしたいと考えている。ここでは、既存の Python ライブラリに型を付けるにあたってのノウハウ的な話を共有したい。

型ヒントとは？

言うまでもなく Python は動的型付け (dynamic typing) なので、近年まで静的な型検査 (static type checking) には対応していなかった。そのため、例えば、ある関数に渡した変数が正しい型であるかは実行してみるまで分からない。また、コード中に出てきた変数の型が分からないので、IDE などでオートコンプリートがうまく働かないという問題があった。

そこで、Python 3.5 では型ヒントと呼ばれる機能が導入された。具体的には、関数アノテーションという記法を使って以下のように書ける:

>>> import typing
>>> def f(a: int, b: int) -> int:
...     return a + b
... 
>>> typing.get_type_hints(f)
{'a': <class 'int'>, 'b': <class 'int'>, 'return': <class 'int'>}

Python インタプリタは、この情報を __annotations__ 属性に格納する以外には何もしない。つまり、この機能は実行時に型検査を行うためのものではない。実際に型検査を行うには、mypy などのサードパーティの静的型検査器を使う。型検査器は、与えられた Python コードを検査して型の整合性を確認する。

mypy-lang.org

あるいは、関数アノテーションではなくコメントで書いてもよい（型コメント）。型検査器は # type: で始まるコメントを認識して、関連付けられた式の型を判断する。Python インタプリタ自体は型コメントを認識しないので、上記のように get_type_hints() などで実行時に型情報を取得することはできないが、関数アノテーション記法のない Python 2 系でも動作するコードになるというメリットがある。Chainer は Python 2 もサポートしているので、今回はこちらの記法を使うことになる。

>>> def f(a, b):
...     # type: (int, int) -> int
...     return a + b
>>> typing.get_type_hints(f)
{}

ユーザは、組み込み型以外にも自前で型を定義できる。単純なエイリアスだけでなく、ジェネリクスのような高度な型にも対応している。型ヒントに使う型の実体は単なる type クラスのインスタンスなので、新しい型を定義する際はコード中に直接書けばよい。

from typing import Generic, TypeVar

UserName = str  # alias

T = TypeVar('T')
class Node(Generic[T]):  # generic type
    ...

面白いのは、これ自体は普通の Python コードである点。例えば、角カッコ ([]) は配列アクセスにも使われる合法的な文法なので Python 3.5 以前でも問題なく実行できる。定義した型は、通常の実行時にはインスタンス化されるだけで使われない。

その他の詳細については、型ヒントの仕様である PEP 484（日本語訳）を参照してほしい。仕様とは言っても分かりやすく書かれているので、静的型付き言語のユーザならば馴染みのある話が多いし簡単に読めると思う。

mypy を既存のプロジェクトに組み込む

mypy は単体でもコマンドとして実行できるが、テストや静的コード解析の一種だということを考えれば、pytest や flake8 のような他のツールと同様に CI を回したい。既存のコードベースへの組み込む際の設定方法については、このドキュメントが参考になった。

注意点としては、mypy は Python 3.4 以降でないとインストールすらできないという点。Chainer のように Python 2.7 を含む各バージョンのテストを回しているライブラリでは、下記のようにチェックしてやる必要がある。

if python -c "import sys; assert sys.version_info >= (3, 4)"; then
    pip install mypy
    mypy chainer
fi

直和型 (Union Type)

既存の動的型付き言語のライブラリに型ヒントをつける際の課題として、ある関数の引数に渡される値が様々な型を取りうる、という点が挙げられる。例えば、以下の関数 double は、数字を渡しても文字列 (str) を渡しても内部的に int に変換して二倍にする。

>>> double(2)
4
>>> double("2")
4

@overload アノテーションで型ごとにシグネチャを宣言してもよいが、これだと複雑な組み合わせへの対応が大変になる。このようなケースを扱うため、typing モジュールは Union という特別な型を提供している。型検査器は、Union に含まれていればどの型が渡されても正しい式として扱う。

>>> from typing import Union
>>> def double(a: Union[int, str]) -> int:
...     if isinstance(a, int):
...         return a * 2
...     else:
...         return int(a) * 2
... 
>>> typing.get_type_hints(double)
{'a': typing.Union[int, str], 'return': <class 'int'>}

実際の例としては、Chainer では ndarray を受け取る関数が多数存在するが、実際には numpy と cupy の両方の場合がありえる。また、今回 ChainerX の導入によって chainerx.ndarray も仲間に加わった。Union を使うことで、このような型もエンコードできる。

NdArray = Union[numpy.ndarray, cuda.ndarray, chainerx.ndarray]

また、同様に Union[..., None] の簡易表現として Optional[...] が使える。これは、以下のように None が渡されることを許容する関数のシグネチャを書く時に使う。Optional といっても（他の言語のような）モナドではないのでそういう使い方はできないが、Optional でない変数に None を渡すとエラーを出してくれるので便利。

def run(config: Optional[Config]=None):
    if config is None:
        ...

ダックタイピングにプロトコルで立ち向かう

Chainer で配列の初期化処理を実装している Initializer は function クラスではないが、__call__ メソッドを実装しているので、あたかも関数であるかのように呼び出すことができる。また、Chainer では initializer を引数に取る関数に、ndarray を引数に取る普通の関数を渡しても動作する。なぜなら、内部的にはダックタイピングで単に initializer(array) (= initializer.__call__(array)) を呼んでいるだけだからだ。

この挙動は通常は便利なのだが、今回のように型をつけようとすると問題になる。なぜなら、Initializer は関数ではないので generate_array(initializer: Callable[[NdArray], None], ...) とすると型が合わないからだ。

この問題は Protocol で解決できる。これは type hint とは別の仕様 (PEP 544) で導入が検討されているもので、まだ draft の段階で正式な仕様ではないが mypy などでは既に使うことができる。

プロトコルは、簡単に言うと「同じ変数やメソッドを持つクラス同士は同じ型として扱う」というものだ。いわゆる構造的部分型で、「静的なダックタイピング」という説明をされることが多い。例えば、Initializer のプロトコルを定義するとこのようになる:

class AbstractInitializer(Protocol):
    dtype = None  # type: Optional[DTypeSpec]

    def __call__(self, array: NdArray) -> None:
        ...

def generate_array(initializer: AbstractInitializer, ...):
    ...

このように定義すれば、generate_array は Initializer と function の両方を受け取ることができる。また、これは実行時には単なるアノテーションなので、これらをラップする共通クラスを作る場合と違って性能上のオーバヘッドもない。

循環参照の解決

若干ハマった点として循環参照の問題がある。例えば、このような chainer.backend.Device 型を含む型を chainer.types モジュールに定義する。

from typing import Tuple, Union  # NOQA
from chainer import backend  # NOQA

DeviceSpec = Union[backend.Device, str, Tuple[str, int], ...]

そして、定義した型を使って backend モジュールの関数に型ヒントを書く。

from chainer import types  # NOQA
from chainer._backend import Device  # NOQA

def get_device(device_spec):
    # type: (types.DeviceSpec) -> Device
    ...

ところが、こうすると import 時に循環参照を起こしてエラーになってしまう。

>>> from chainer import backend
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/okapies/chainer/chainer/backend.py", line 1, in <module>
    from chainer import types  # NOQA
  File "/home/okapies/chainer/chainer/types.py", line 4, in <module>
    DeviceSpec = Union[backend.Device, str, Tuple[str, int]]
AttributeError: module 'chainer.backend' has no attribute 'Device'

これを防ぐため、types モジュールに以下のようなおまじないを追加して、型検査器が Python コードを読み込んだ時のみ chainer.backend が import されるようにする。

try:
    from typing import TYPE_CHECKING  # NOQA
except ImportError:
    # typing.TYPE_CHECKING doesn't exist before Python 3.5.2
    TYPE_CHECKING = False

# import chainer modules only for type checkers to avoid circular import
if TYPE_CHECKING:
    from chainer import backend  # NOQA

DeviceSpec = Union['backend.Device', str, Tuple[str, int], ...]

TYPE_CHECKING は Python 3.5.2 から追加された変数で、Python インタプリタでは常に False として扱われる。一方、mypy などの型検査器はこれを True として扱うので、型チェックをする場合のみ import 文が動くようにできる。さらに、DeviceSpec に定義する型を文字列として記述することで、実行時に import されていないシンボルを読み込むことによる文法エラーを回避する（型検査器は文字列を type クラスに変換してくれる）。

まとめ

というわけでボチボチと手を動かしていが、歴史のあるプロダクトなだけに API 本数も多く、まだまだ先は長そう。気長にやっていこうと思うので、応援して頂けると幸い。

github.com