如何在 Python 中为外部库的函数定义创建自定义输入类型?

How to create custom input types for function defenition for external libraries in Python?

我想编写具有详细类型的函数,以便在使用任何类型检查器时,例如Pylance 在 VScode 你可以清楚地看到输入类型。我设法使用 int 等默认输入类型取得了一些结果,然后将其扩展为 numpy.ndarray。但这是我 运行 遇到问题的地方。示例函数的代码如下:

import numpy as np

def get_distance(arr1: np.ndarray, arr2: np.ndarray) -> int:
    """Returns Euclidian distance between 2 points in 3D space

    Args:
        arr1 (np.ndarray[float]): array 1 format [[x,y,z],...]
        arr2 (np.ndarray[float]): array 2 format [[x,y,z],...]

    Returns:
        [int]: Distance between 2 points
    """
    ...

return some_distance

当我定义它然后在 Pylance 下寻找定义时,我设法得到 np.ndarray[Unknown, Unknown]:

(function) get_distance: (arr1: ndarray[Unknown, Unknown], arr2: ndarray[Unknown, Unknown]) -> int 

Returns Euclidian distance between 2 points in 3D space

Args:
    arr1 (np.ndarray[float]): point array 1
    arr2 (np.ndarray[float]): point array 2

Returns:
    [int]: Distance between 2 points

我的问题是:是否有任何方法可以定义输入类型,以便即使在使用 numpy 或任何其他外部库时也能显示正确的输入类型?例如 np.ndarray[float] 或类似的东西。

我不确定我是否理解您的问题,但如果您希望显示 np.ndarray[float],这正是您需要传递的内容,而不仅仅是 np.ndarray,如下所示:

def get_distance(arr1: np.ndarray[float], arr2: np.ndarray[float]) -> int:

这是你想要的吗?

from numpy.typing import NDArray
import numpy as np

def get_distance(
  arr1: NDArray[np.float64], 
  arr2: NDArray[np.float64]
) -> int:

    # ...

    return some_distance

请注意,您需要使用 numpy 的 typing 子模块。

对我来说这个方法很管用。

但请再次注意,此方法无效:arr: np.ndarray[float]

正如您在我的带有 LSP-pyright 的编辑器中看到的那样

这是带下划线的 red (error)

但是这个方法很好:

再次注意:你可以使用任何你喜欢的浮动

def get_distance(arr1: NDArray[np.float32] ...
def get_distance(arr1: NDArray[np.float64] ...
def get_distance(arr1: NDArray[np.float128] ...
# also integer
def get_distance(arr1: NDArray[np.int32],

但是,最重要的是,您使用的数据类型应该在 numpy 模块中定义,并且您必须使用来自 numpy 的 typing,而不是来自标准的 python。为什么?因为numpy有自己的类似C的数据类型,因为numpy是用C写的,在python.

编译导入

您的问题特定于 ndarray -> 它的类型有 2 个参数 - 形状和数据类型。 TypeError: 'type' object is not subscriptable 是因为您将浮点数作为形状参数传递,它应该是类似列表的。 (尽管我个人在 Python 3.10 - TypeError: Too few arguments for numpy.ndarray

上得到了更明显的错误

这会起作用:

from typing import Any

def get_distance(arr1: np.ndarray[Any, float], arr2: np.ndarray[Any, float]) -> int: