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

Question

我想编写具有详细类型的函数，以便在使用任何类型检查器时，例如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] 或类似的东西。

Answer 1

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

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

Answer 2

这是你想要的吗？

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.

编译导入

Answer 3

您的问题特定于 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:

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

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

python

function

types

pep8