How to type a class decorator registry pattern?

[cdce8p]

Came across a pattern today that left me wondering how it should be typed instead or if it's even possible to type probably.
It works for Mypy but I'm guessing this is just because class decorators aren't type checked at all python/mypy#3135.

from typing import TypeVar, Callable

_KT = TypeVar("_KT")
_VT = TypeVar("_VT")

class Registry(dict[_KT, _VT]):
    def register(self, name: _KT) -> Callable[[_VT], _VT]:
        def decorator(func: _VT) -> _VT:
            self[name] = func
            return func
        return decorator


class Provider: ...

PROVIDERS: Registry[str, type[Provider]] = Registry()

@PROVIDERS.register("auth")
class AuthProvider(Provider): ...

def a1() -> None:
    reveal_type(AuthProvider)              # this should be type[AuthProvider]
    reveal_type(list(PROVIDERS.values()))  # this should be list[type[Provider]]

    @PROVIDERS.register("error")           # this should be an error
    class Error: ...

[Daverball]

_VT should probably have an upper bound of Callable[..., Any], since register will only work for functions or classes. Or if you actually want to restrict it to types you could also do class Registry(dict[_KT, type[_VT]]): and replace all the other occurrences of _VT with type[_VT].

Class decorators are type checked now, but they can't change the type, which should be fine for your use-case. It looks like your example code already works just fine with the behavior you expect: https://mypy-play.net/?mypy=latest&python=3.12&gist=9843acb42af5acef8f7a982d2c7207ac

[cdce8p]

_VT should probably have an upper bound of Callable[..., Any], since register will only work for functions or classes. Or if you actually want to restrict it to types you could also do class Registry(dict[_KT, type[_VT]]): and replace all the other occurrences of _VT with type[_VT].

👍🏻 Good point. It's actually included in the original source code, just forgot to copy it to the example.

Class decorators are type checked now, but they can't change the type, which should be fine for your use-case. It looks like your example code already works just fine with the behavior you expect: https://mypy-play.net/?mypy=latest&python=3.12&gist=9843acb42af5acef8f7a982d2c7207ac

🤔 Does that apply for all type checkers that class decorators should not be able to change the type?
It seems that's how both Mypy and Pyre interpret it that way. However the example doesn't work with pyright.

pyright playground
pyre playground

/CC: @erictraut

[Daverball]

The pyright output looks correct to me as well, maybe the error message is what is confusing you?

[cdce8p]

The pyright output looks correct to me as well, maybe the error message is what is confusing you?

Whereas both Mypy and pyre infer reveal_type(AuthProvider) as type[AuthProvider], pyright uses the decorator typing and infers type[Provider]. This causes issues for me down the line where AuthProvider is expected, as even with var: AuthProvider, var would only be inferred as Provider, and thus missing the methods that only AuthProvider has.

That's why I was asking if my initial typing was wrong.

[Daverball]

Oh, I see. Yeah I missed that and I also didn't really think hard enough about that part.

What you really want is for the decorator to remain generic, regardless of what you bind the registry to (binding the type var would only change the upper bound). I unfortunately don't think there's a way to express that, it would require some higher kinded TypeVar which could set its upper bound to the value of another TypeVar.

For now I would probably just either create separate non-generic Registry subclasses for each type, where the method itself can remain generic, or give up some amount of type safety and use a different TypeVar in the register method.

Maybe in some hypothetical future of Python we'll be allowed to do things like this with PEP695 syntax, since there's not as much ambiguity about scoping, to allow doing things like changing the upper bound based on another TypeVar in the scope of a generic method within a generic class.

class Registry[_KT, _VT](dict[_KT, _VT]):
    def register[_T: _VT](self, name: _KT) -> Callable[[_T], _T]:
        def decorator(func: _T) -> _T:
            self[name] = func
            return func
        return decorator

[Daverball]

In a future with intersection this might be another potential solution

def register[_T](self, name: _KT) -> Callable[[_T & _VT], _T]:

Although the semantics here are a bit more questionable, since it's unclear what _T should bind to when it's part of an intersection.

[Akuli]

You could cast the decorator to _typeshed.IdentityFunction so it will always return the type passed into it:

from _typeshed import IdentityFunction

...

class Registry(dict[_KT, _VT]):
    def register(self, name: _KT) -> IdentityFunction:
        def decorator(func: _VT) -> _VT:
            self[name] = func
            return func
        return cast(IdentityFunction, decorator)

Pyright playground link

[cdce8p]

IdentifyFunction is a nice trick but, unless I'm missing something, I'd still lose the error detection for this unfortunately

    @PROVIDERS.register("error")           # this should be an error
    class Error: ...

[Daverball]

Yeah, pretty much, you can't have both right now if you want to support pyright (at least if you go with the generic Registry, for a specific Registry you can achieve it using type var upper bounds, but that also means you need a new type var and Registry class definition for each type of registry).

And even if you didn't care about supporting pyright, there's a chance that mypy and other type checkers will respect the return type of the decorator in the future and then it'll break anyways.