需要装饰器函数接受与绑定的`TypeVar`匹配的参数而不缩小到该类型

Question

如果我这样定义我的装饰器：

T = TypeVar('T', bound=Event)

def register1(evtype: Type[T]) -> Callable[[Callable[[T], None]], Callable[[T], None]]:
    def decorator(handler):
        # register handler for event type
        return handler
    return decorator

当我把它用在错误的函数上时，会出现一个合适的错误：

class A(Event):
    pass

class B(Event):
    pass

@register1(A) # Argument of type "(ev: B) -> None" cannot be assigned to parameter of type "(A) -> None"
def handler1_1(ev: B):
    pass

不过，如果我多次使用这个装饰器，它就不管用了：

@register1(A) # Argument of type "(B) -> None" cannot be assigned to parameter of type "(A) -> None"
@register1(B)
def handler1_3(ev: A|B):
    pass

我希望这些装饰器能够组合出一个允许或必需的参数类型的Union。

我觉得ParamSpec可以解决这个问题，但我该如何使用ParamSpec，既不覆盖参数类型，又能确保参数类型与装饰器中的类型匹配呢？

使用ParamSpec并不会导致任何类型错误：

P = ParamSpec("P")

def register2(evtype: Type[T]) -> Callable[[Callable[P, None]], Callable[P, None]]:
    def decorator(handler):
        # ...
        return handler
    return decorator

@register2(A) # This should be an error
def handler2_1(ev: B):
    pass

如果我再添加一个TypeVar并使用Union，那么对于双重装饰和甚至三重装饰的函数都是有效的，但对于单重装饰的函数就不行了。

T2 = TypeVar('T2')

def register3(evtype: Type[T]) -> Callable[[Callable[[Union[T,T2]], None]], Callable[[Union[T,T2]], None]]:
    def decorator(handler):
        # ...
        return handler
    return decorator

# Expected error:
@register3(A) # Argument of type "(ev: B) -> None" cannot be assigned to parameter of type "(A | T2@register3) -> None"
def handler3_1(ev: B):
    pass

# Wrong error:
@register3(A) # Argument of type "(ev: A) -> None" cannot be assigned to parameter of type "(A | T2@register3) -> None"
def handler3_2(ev: A):
    pass

# Works fine
@register3(A)
@register3(B)
def handler3_3(ev: A|B):
    pass

在写这个问题的过程中，我逐渐接近了解决方案。 我会在回答中提供我自己的解决方案。

不过，我也想知道是否还有更好的解决方法。

Answer 1

这只是一个部分解决方案。它只对 register 这一侧有效，但对调用这一侧的类型检查并不准确。

通过在装饰器的参数中添加一个情况，处理当被装饰的函数只接受一个参数时，使用 Union，我不再收到 pyright 的意外错误了：

def register4(evtype: Type[T]) -> Callable[[Union[Callable[[T|T2], None],Callable[[T], None]]], Callable[[T|T2], None]]:
    def decorator(handler):
        # ...
        return handler
    return decorator

#Expected errors
@register4(A) # Argument of type "(ev: B) -> None" cannot be assigned to parameter of type "((A | T2@register4) -> None) | ((A) -> None)"
def handler4_1(ev: B):
    pass

@register4(A)
def handler4_2(ev: A):
    pass

@register4(A)
@register4(B)
#@register4(C)
def handler4_3(ev: A|B|C):
    pass

正如评论中提到的：

handler4_2(B())

这并不会导致错误，尽管应该会有。

我尝试通过将 Union 拆分成 @overload 声明来修复这个问题，但那并没有奏效：

@overload
def register4(evtype: Type[T]) -> Callable[[Callable[[T | T2], None]],Callable[[T | T2], None]]: ...


@overload
# Overload 2 for "register2" will never be used because its parameters overlap overload 1
def register4(evtype: Type[T]) -> Callable[[Callable[[T], None]], Callable[[T], None]]: ...

我认为它忽略了第二个重叠，因为它认为这两个重叠，但我们已经看到，在 register3 中这两种情况的解释是不同的。而且如果我交换这两个声明，示例的行为也会改变。所以这可能是 pyright 的一个bug。

这是在 pyright 1.1.352 的情况下。

Answer 2

来自 https://github.com/microsoft/pyright/discussions/7404 --

from __future__ import annotations

from typing import Any, Callable, Protocol, TypeVar, overload

T_co = TypeVar("T_co", covariant=True)

T0 = TypeVar("T0")
T1 = TypeVar('T1')

class RegisterResult(Protocol[T_co]):
    @overload
    def __call__(self, handler: Callable[[T_co | T1], None]) -> Callable[[T_co | T1], None]: ...

    @overload
    def __call__(self, handler: Callable[[T_co], None]) -> Callable[[T_co], None]: ...

def register(evtype: type[T0]) -> RegisterResult[T0]:
    def decorator(handler: Any) -> Any:
        return handler
    
    return decorator

class A: ...
class B: ...
class C: ...

@register(A)
def handle_a(ev: A): ...

handle_a(A())

@register(A)
@register(B)
# ... Should support infinite amount of @register calls
def handle_ab(ev: A|B): ... 

handle_ab(A())
handle_ab(B())

#Expected error cases because of wrong types:
@register(A)
def handle_b(ev: B): ... 
handle_a(B())
handle_ab(C())

请注意，上面的代码在最新的Pyright版本（v1.1.353）中是可以正常工作的，但未来的版本可能会因为Pyright处理重载函数的兼容性方式不同而导致代码无法正常运行。而且我检查过，它在最新的Mypy版本（v1.9.0）中并不能完全正常工作。