Numpyの型ヒント備忘録(NDArray)

Numpy の型ヒント・型チェックについて、ちょっとハマったので備忘録として残します

今回使っているのは nptyping という numpy 用の型付けのためのライブラリです。リポジトリ: GitHub - ramonhagenaars/nptyping

ベクトルのとき

下の二つは要素が int でなくても True となります。(Any は標準ライブラリの typing からインポートします)

vec = np.array([1, 2, 3], int)
isinstance(vec, NDArray[3, int]) # True
isinstance(vec, NDArray[(3,), int]) # True
isinstance(vec, NDArray[(3, ...), int]) # True
isinstance(vec, NDArray[(3,), Any]) # True
isinstance(vec, NDArray[3]) # True

行列のとき

一番上は $(any)\times3$ 行列。下の二つは int じゃなくても True になります。(Any は標準ライブラリの typing からインポートします)

mat = np.array([[1, 2, 3], [4, 5, 6]], int)
isinstance(mat, NDArray[(Any, 3), int]) # True
isinstance(mat, NDArray[(2, 3), int]) # True
isinstance(mat, NDArray[2, 3]) # True
isinstance(mat, NDArray[(2, 3), Any]) # True