对于具有“@”（矩阵乘法）的对象的 Python 类型提示

Question

我有一个函数fun()接受 NumPy ArrayLike和“矩阵”，并返回 numpy 数组。

from numpy.typing import ArrayLike
import numpy as np

def fun(A, x: ArrayLike) -> np.ndarray:
    return (A @ x) ** 2 - 27.0

type对于有业务的实体，正确的是什么@？请注意，fun()也可以接受scipy.sparse；也许更多。

Answer 1

您可以使用typing.Protocol来断言该类型实现了__matmul__.

class SupportsMatrixMultiplication(typing.Protocol):
    def __matmul__(self, x):
        ...


def fun(A: SupportsMatrixMultiplication, x: ArrayLike) -> np.ndarray:
    return (A @ x) ** 2 - 27.0

我相信，x如果您想要的不仅仅是@作为运算符提供支持，您可以通过提供类型提示和返回类型提示来进一步完善这一点。