トップ > Python > numpy.typing.NDArrayの整理

numpy.typing.NDArrayの整理

Python

概要

numpy version=1.20からnumpy.typingが提供されています。

numpy.org

型アノテーションを記述する際に、numpy.ndarrayで指定するのと、numpy.typing.NDArrayで指定するのは何が違うのかを整理します。

numpy.typing.NDArrayとは何か

Typing (numpy.typing) — NumPy v1.26 Manualによると、

A generic version of np.ndarray[Any, np.dtype[+ScalarType]].

と記載されています。

numpy==1.26.2におけるnumpy.typing.NDArrayの実装を確認してみます。

import numpy as np
import numpy.typing as npt

npt.NDArray

numpy.ndarray[typing.Any, numpy.dtype[+_ScalarType_co]]

コード上では以下のように実装されています。

from numpy import (
    ndarray,
    dtype,
    generic,
    ...
)
...
_ScalarType_co = TypeVar("_ScalarType_co", bound=generic, covariant=True)
...
NDArray = ndarray[Any, dtype[_ScalarType_co]]

numpy.typing.NDArrayはndarray[Any, dtype[_ScalarType_co]]の型エイリアスということが分かります。

また、_ScalarType_coはTypeVar("_ScalarType_co", bound=generic, covariant=True)という型変数です。

  • bound=genericと指定しています。Scalars — NumPy v1.26 Manualによると、numpy.genericはnumpyにおけるスカラー型の基底クラスとのことです。_ScalarType_coはnumpy.genericの部分型（要はスカラー型）を受け取ります。
  • covariant=Trueと指定しています。AがBの部分型の場合、dtype[A]はdtype[B]の部分型と見なされます。

ndarray[Any, dtype[_ScalarType_co]]という実装について、Anyとdtype[_ScalarType_co]]はそれぞれ何に対応しているのかを確認していきます。

コード上では、ndarrayは以下のように定義されています。

_DType_co = TypeVar("_DType_co", covariant=True, bound=dtype[Any])
...
# TODO: Set the `bound` to something more suitable once we
# have proper shape support
_ShapeType = TypeVar("_ShapeType", bound=Any)
...
class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
    __hash__: ClassVar[None]
...

_ShapeTypeはndarrayのshape情報が与えられます。 ただ、現状の_ShapeTypeの実装ではbound=Anyとなっているので、Anyの部分型（要はなんでも）を受け取れてしまいます。 TODOコメントにもあるように、shapeチェックができるようになると良いですね。

…というわけで、型アノテーションにおいて2つを以下のように整理しました。

  • ndarrayではShapeTypeとDtypeの2つを指定できる。ただし、現状ShapeTypeには任意の型が渡せてしまう。
  • NDArrayはndarray[Any, dtype[_ScalarType_co]]と等価。

個人的には、型アノテーションがシンプルになることからNDArrayを使うのが良いんじゃないかなと思います1。

NDArrayを利用した型アノテーション

実際にNDArrayを使って型アノテーションしてみましょう。

例えば以下のように型アノテーションが可能です。

import numpy as np
import numpy.typing as npt


def add_one(xs: npt.NDArray) -> npt.NDArray:
    return xs + 1


xs = np.array([1, 2, 3])
add_one(xs)

$ mypy exp.py

Success: no issues found in 1 source file

ただし上記ではNDArrayのdtype部分を指定していないため、mypy --strictもしくはmypy --disallow-any-genericsでは以下のように怒られます。

$ mypy --disallow-any-generics exp.py 

exp.py:5: error: Missing type parameters for generic type "npt.NDArray"  [type-arg]
Found 1 error in 1 file (checked 1 source file)

以下のようにdtypeを指定してあげればOKです。

def add_one(xs: npt.NDArray[np.int32]) -> npt.NDArray[np.int32]:
    return xs + 1

誤ってNDArray[np.float32]型の値を渡すと、ちゃんとmypyが怒ってくれます。

xs: npt.NDArray[np.float32] = np.array([1, 2, 3])
add_one(xs)

$ mypy --disallow-any-generics exp.py
exp.py:10: error: Argument 1 to "add_one" has incompatible type "ndarray[Any, dtype[floating[_32Bit]]]"; expected "ndarray[Any, dtype[signedinteger[_32Bit]]]"  [arg-type]
Found 1 error in 1 file (checked 1 source file)

まとめ

numpy.typing.NDArrayの実装および使い方について見てきました。

numpy.typingは他にもArrayLikeやDTypeLikeなどの型を提供しており、うまく使ってあげることで素敵な型アノテーションができそうです。

  1. nptypingではShape型を提供しており、ShapeTypeを指定した型アノテーションができるようです。メンテナンスは止まっていそうですが、、↩