return 函数值的类型注解 returns numpy 单值

Question

如果我想给下面函数的return值加上类型注解，应该怎么办？

import numpy as np

def myfunc(x: np.ndarray):
    return x.sum()

return函数值的类型不同：

np.int8 如果 x.dtype 是 dtype('int8')
np.int16 如果 x.dtype 是 dtype('int16')
np.int32 如果 x.dtype 是 dtype('int32')
np.float16 如果 x.dtype 是 dtype('float16')
np.float32 如果 x.dtype 是 dtype('float32') ...等等

因为我想避免使用 -> Any。

Answer 1

您可以使用 numpy 类型模块中的 NDArray 并定义绑定到 numpy 标量类型的 TypeVar。您可以将此绑定类型传递给 NDArray 类型并使用相同类型注释函数的 return 类型。

import numpy as np
from typing import TypeVar
from numpy.typing import NDArray

ScalarType = TypeVar("ScalarType", bound=np.generic, covariant=True)

def myfunc(x: NDArray[ScalarType]) -> ScalarType:
    return x.sum()

return 函数值的类型注解 returns numpy 单值

Type annotation of return value of a function which returns numpy single value

python

numpy

mypy