Type hint for a generic decorator that can be applied multiple times and checks argument.

register_decorator.py

# pyright: strict
# Code taken from: https://github.com/microsoft/pyright/discussions/7404

from collections.abc import Callable
from typing import Protocol, overload


class RegisterResult[T](Protocol):
    @overload
    def __call__[U, R](self, handler: Callable[[T | U], R]) -> Callable[[T | U], R]: ...

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


def register[T](ev: type[T]) -> RegisterResult[T]: ...


class A: ...


class B: ...


class C: ...


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


handle_a(A())


@register(A)
@register(B)
def handle_ab(ev: A | B): ...

@register(A)
@register(B)
@register(C)
def handle_abc(ev: A | B | C): ...

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())