python - 如何输入 float 和/或整数的numpy数组

Question

我不确定如何输入以下内容:

def prop(
    *,
    arr,  # numpy array of floats or/and ints
):
    return 100 * arr / arr.sum()

我试过以下方法:

def prop(
    *,
    arr: npt.NDArray[np._IntType] | npt.NDArray[np._FloatType],  # numpy array of floats or/and ints
) -> npt.NDArray[np._IntType] | npt.NDArray[np._FloatType]:
    return 100 * arr / arr.sum()

并且正在运行 mypy <path>/check_mypy.py --strict 时遇到错误:

check_mypy.py:12: error: Returning Any from function declared to return "Union[ndarray[Any, dtype[np._IntType]], ndarray[Any, dtype[np._FloatType]]]"  [no-any-return]
check_mypy.py:17: error: Need type annotation for "x"  [var-annotated]
check_mypy.py:18: error: Need type annotation for "x1"  [var-annotated]
check_mypy.py:22: error: Need type annotation for "y"  [var-annotated]
check_mypy.py:23: error: Need type annotation for "y1"  [var-annotated]
check_mypy.py:27: error: Need type annotation for "z"  [var-annotated]
check_mypy.py:28: error: Need type annotation for "z1"  [var-annotated]

完整示例:

# check_mypy.py
from __future__ import annotations

import numpy as np
import numpy.typing as npt


def prop(
    *,
    arr: npt.NDArray[np._IntType] | npt.NDArray[np._FloatType],  # numpy array of floats or/and ints
) -> npt.NDArray[np._IntType] | npt.NDArray[np._FloatType]:
    return 100 * arr / arr.sum()


def main() -> int:
    # check 1 - ints
    x = np.array([1, 2, 3])
    x1 = prop(arr=x)
    print(x1)

    # check 2 - mixed
    y = np.array([1.4, 21, 3.2])
    y1 = prop(arr=y)
    print(y1)

    # check 3 - floats
    z = np.array([1.4, 2.1, 3.2])
    z1 = prop(arr=z)
    print(z1)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())

编辑1

If you do x: npt.NDArray[np._IntType] = np.array([1, 2, 3]) does that work?

按照上述我得到以下错误输出:

check_mypy.py:12: error: Returning Any from function declared to return "Union[ndarray[Any, dtype[np._IntType]], ndarray[Any, dtype[np._FloatType]]]"  [no-any-return]
check_mypy.py:17: error: Type variable "numpy._IntType" is unbound  [valid-type]
check_mypy.py:17: note: (Hint: Use "Generic[_IntType]" or "Protocol[_IntType]" base class to bind "_IntType" inside a class)
check_mypy.py:17: note: (Hint: Use "_IntType" in function signature to bind "_IntType" inside a function)
check_mypy.py:18: error: Need type annotation for "x1"  [var-annotated]
check_mypy.py:23: error: Need type annotation for "y"  [var-annotated]
check_mypy.py:24: error: Need type annotation for "y1"  [var-annotated]
check_mypy.py:28: error: Need type annotation for "z"  [var-annotated]
check_mypy.py:29: error: Need type annotation for "z1"  [var-annotated]

编辑2.

将函数更新为:

T = TypeVar('T', np._IntType, np._FloatType)


def prop(
    *,
    arr: npt.NDArray[T],  # numpy array of floats or/and ints
) -> npt.NDArray[T]:
    return 100 * arr / arr.sum()

我得到以下错误输出:

check_mypy.py:15: error: Type variable "numpy._IntType" is unbound  [valid-type]
check_mypy.py:15: note: (Hint: Use "Generic[_IntType]" or "Protocol[_IntType]" base class to bind "_IntType" inside a class)
check_mypy.py:15: note: (Hint: Use "_IntType" in function signature to bind "_IntType" inside a function)
check_mypy.py:15: error: Type variable "numpy._FloatType" is unbound  [valid-type]
check_mypy.py:15: note: (Hint: Use "Generic[_FloatType]" or "Protocol[_FloatType]" base class to bind "_FloatType" inside a class)
check_mypy.py:15: note: (Hint: Use "_FloatType" in function signature to bind "_FloatType" inside a function)
check_mypy.py:22: error: Returning Any from function declared to return "ndarray[Any, dtype[np._IntType?]]"  [no-any-return]
check_mypy.py:22: error: Returning Any from function declared to return "ndarray[Any, dtype[np._FloatType?]]"  [no-any-return]
check_mypy.py:27: error: Type variable "numpy._IntType" is unbound  [valid-type]
check_mypy.py:27: note: (Hint: Use "Generic[_IntType]" or "Protocol[_IntType]" base class to bind "_IntType" inside a class)
check_mypy.py:27: note: (Hint: Use "_IntType" in function signature to bind "_IntType" inside a function)
check_mypy.py:33: error: Need type annotation for "y"  [var-annotated]
check_mypy.py:38: error: Need type annotation for "z"  [var-annotated]

编辑3

If you actually want to allows mixed arrays, then you can drop the type variable and use npt.NDArray[ np._IntType | np._FloatType ] in both cases

将函数更新为:

def prop(
    *,
    arr: npt.NDArray[np._IntType | np._FloatType],  # numpy array of floats or/and ints
) -> npt.NDArray[np._IntType | np._FloatType]:
    return 100 * arr / arr.sum()

出现以下错误:

check_mypy.py:22: error: Returning Any from function declared to return "ndarray[Any, dtype[Union[np._IntType, np._FloatType]]]"  [no-any-return]
check_mypy.py:27: error: Type variable "numpy._IntType" is unbound  [valid-type]
check_mypy.py:27: note: (Hint: Use "Generic[_IntType]" or "Protocol[_IntType]" base class to bind "_IntType" inside a class)
check_mypy.py:27: note: (Hint: Use "_IntType" in function signature to bind "_IntType" inside a function)
check_mypy.py:30: error: Need type annotation for "x1"  [var-annotated]
check_mypy.py:35: error: Need type annotation for "y"  [var-annotated]
check_mypy.py:36: error: Need type annotation for "y1"  [var-annotated]
check_mypy.py:40: error: Need type annotation for "z"  [var-annotated]
check_mypy.py:41: error: Need type annotation for "z1"  [var-annotated]

Answer 1

我真的很好奇这会怎样。我能得到的最好的是以下内容

# check_mypy.py
from __future__ import annotations

# Third party
import numpy as np
import numpy.typing as npt
from typing import cast


def prop(
    *,
    arr: npt.NDArray[np.float64] | npt.NDArray[np.int64] ,  # numpy array of floats or/and ints
) -> npt.NDArray[np.float64]:
    return cast(npt.NDArray[np.float64], 100 * arr / arr.sum())


def main() -> int:
    # check 1 - ints
    x : npt.NDArray[np.float64] = np.array([1, 2, 3])
    x1 = prop(arr=x)
    print(x1)

    # check 2 - mixed
    y : npt.NDArray[np.float64]  = np.array([1.4, 21, 3.2])
    y1 = prop(arr=y)
    print(y1)

    # check 3 - floats
    z: npt.NDArray[np.int64] = np.array([1.4, 2.1, 3.2])
    z1 = prop(arr=z)
    print(z1)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())

所以没有混合类型，因为这没有意义。 Numpy 数组始终是单一类型，即使您在定义数组时添加了其他类型也是如此。

使用 np.float64/np.int64 实际上是它最终起作用的原因。不确定这是否适合您。在这里您可能想要定义一个 Union 或所有 float 的类型 var，但又不确定。重要的是它是 NP.generic 的子类。

还需要对所有定义的数组进行类型注释，因为 mypy 无法自行确定类型。

最后，返回固定的 npt.NDArray[np.float64] 作为你正在计算的总是那种类型。