NumPy ndarray dtype 的类型提示？

Question

我想要一个函数来包含 NumPy ndarray 的类型提示及其 dtype。

例如，对于列表，可以执行以下操作...

def foo(bar: List[int]):
   ...

...为了给出 bar 必须是 list 由 int 组成的类型提示。

不幸的是，此语法会引发 NumPy 异常 ndarray:

def foo(bar: np.ndarray[np.bool]):
   ...

> np.ndarray[np.bool]) (...) TypeError: 'type' object is not subscriptable

是否可以为 np.ndarray 提供 dtype 特定的类型提示？

Answer 1

据我所知，尚无法在函数签名的 numpy 数组类型提示中指定 dtype。计划在未来的某个时候实施。有关当前开发状态的更多详细信息，请参阅 numpy GitHub issue #7370 and numpy-stubs GitHub。

Answer 2

你可以查看 nptyping:

from nptyping import NDArray, Bool

def foo(bar: NDArray[Bool]):
   ...

或者您可以只使用字符串作为类型提示：

def foo(bar: 'np.ndarray[np.bool]'):
   ...

Answer 3

检查 data-science-types 包。

pip install data-science-types

MyPy 现在可以访问 Numpy、Pandas 和 Matplotlib 存根。允许以下场景：

# program.py

import numpy as np
import pandas as pd

arr1: np.ndarray[np.int64] = np.array([3, 7, 39, -3])  # OK
arr2: np.ndarray[np.int32] = np.array([3, 7, 39, -3])  # Type error

df: pd.DataFrame = pd.DataFrame({'col1': [1,2,3], 'col2': [4,5,6]}) # OK
df1: pd.DataFrame = pd.Series([1,2,3]) # error: Incompatible types in assignment (expression has type "Series[int]", variable has type "DataFrame")

正常使用 mypy。

$ mypy program.py

与function-parameters

的用法

def f(df: pd.DataFrame):
    return df.head()

if __name__ == "__main__":
    x = pd.DataFrame({'col1': [1, 2, 3, 4, 5, 6]})
    print(f(x))

$ mypy program.py
> Success: no issues found in 1 source file

Answer 4

类型文档的一种非正式解决方案如下：

from typing import TypeVar, Generic, Tuple, Union, Optional
import numpy as np

Shape = TypeVar("Shape")
DType = TypeVar("DType")


class Array(np.ndarray, Generic[Shape, DType]):
    """
    Use this to type-annotate numpy arrays, e.g.

        def transform_image(image: Array['H,W,3', np.uint8], ...):
            ...

    """
    pass


def func(arr: Array['N,2', int]):
    return arr*2


print(func(arr = np.array([(1, 2), (3, 4)])))

我们一直在我的公司使用它并制作了一个 MyPy 检查器，它实际上检查形状是否有效（我们应该在某个时候发布）。

唯一的问题是它不会让 PyCharm 开心（即你仍然会看到讨厌的警告线）：

NumPy ndarray dtype 的类型提示？

Type hint for NumPy ndarray dtype?

python

numpy

type-hinting

python-3.x

python-typing