设计支持float与NDArray参数的类型标注装饰器遇类型错误求助

你的代码触发类型错误，是因为原有的装饰器类型标注没让Pylance正确识别被装饰函数同时支持float输入返回float、**NDArray[float64]输入返回NDArray[float64]**这两种行为。原TypeVar的用法无法建立输入输出的类型关联，导致类型检查器认为func只能接收数组。

下面是修正后的代码：

import numpy as np
from numpy import float64
from numpy.typing import NDArray
from collections.abc import Callable
from typing import TypeAlias, TypeVar, overload

T = TypeVar('T', float, NDArray[float64])
PreWrapFunc: TypeAlias = Callable[[NDArray[float64]], NDArray[float64]]

# 用重载明确装饰器的两种返回签名
@overload
def my_decorator(method: PreWrapFunc) -> Callable[[float], float]: ...

@overload
def my_decorator(method: PreWrapFunc) -> Callable[[NDArray[float64]], NDArray[float64]]: ...

# 实际实现的类型标注
def my_decorator(method: PreWrapFunc) -> Callable[[T], T]:
    def wrapper(arg: T) -> T:
        if isinstance(arg, float):
            result = method(np.array([arg]))
            return result[0]
        else:
            return method(arg)
    return wrapper

@my_decorator
def func(a: NDArray[float64]) -> NDArray[float64]:
    return a * 2

func(1.0)  # 类型检查通过，返回float
func(np.array([1.0,]))  # 类型检查通过，返回NDArray[float64]

关键修改说明： #

• 新增@overload装饰器，给类型检查器明确两个分支：装饰后的函数既能接收float返回float，也能接收数组返回数组。
• 底层实现保留Callable[[T], T]，让重载与实际逻辑对应，确保类型推断一致性。
• Pylance现在能通过重载识别func的合法调用签名，不会再报类型不兼容的错误。

内容的提问来源于stack exchange，提问作者Nukesub