如何输入hint函数以与numpy兼容

kym*_*kym 7 python numpy mypy python-typing

源代码example.py

from typing import Union, Any
import numpy as np

Number = Union[int, float, np.floating[Any]]


def add_one(num: Number) -> Number:
    return num + 1


inputs = [1, 2, 3]
outputs = [add_one(n) for n in inputs]

avg = np.mean(outputs)
Run Code Online (Sandbox Code Playgroud)

运行 mypy:

mypy example.py
src/example.py:14: error: Argument 1 to "mean" has incompatible type "List[Union[float, floating[Any]]]"; expected "Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]"
Found 1 error in 1 file (checked 1 source file)
Run Code Online (Sandbox Code Playgroud)

我可以更改np.floating[Any]解决 numpy 问题的所有内容,但随后我必须将基元转换为np.float32(...)

mypy example.py
src/example.py:14: error: Argument 1 to "mean" has incompatible type "List[Union[float, floating[Any]]]"; expected "Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]"
Found 1 error in 1 file (checked 1 source file)
Run Code Online (Sandbox Code Playgroud)

是否有正确的方法来键入提示add_one函数,以便其输出与 numpy 函数兼容,例如np.mean 不破坏与 python 基元类型的兼容性?最终目标是能够像这样使用它:

from typing import Any
import numpy as np


def add_one(num: np.floating[Any]) -> np.floating[Any]:
    return num + 1


inputs = [1, 2, 3]
outputs = [add_one(np.float32(n)) for n in inputs]

avg = np.mean(outputs)
Run Code Online (Sandbox Code Playgroud)

小智 0

仅使用Number = Union[int, float]不会引发任何 mypy 错误。