Search code examples
pythonpython-typingtype-parameterpyrightpep-695

Type hints lost when a decorator is wrapped as a classmethod?


Consider the following code:

from typing import Any, Callable, Coroutine

class Cache[**P, R]:
    @classmethod
    def decorate(cls, **params):
        def decorator(f: Callable[P, Coroutine[Any, Any, R]]) -> Callable[P, Coroutine[Any, Any, R]]:
            return f # in the real world, we instantiate a Cache here
        return decorator

@Cache.decorate()
async def some_function(i: int) -> int:
    return i + 1

cached_function = Cache.decorate()(some_function)

If I ask pyright the type of Cache.decorate before the @classmethod wrapper (inspecting the word decorate in the above code), it returns:

(method) def decorate(
    cls: type[Self@Cache[P@Cache, R@Cache]],
    **params: Unknown
) -> ((f: ((**P@Cache) -> (Coroutine[Any, Any, R@Cache])) -> ((**P@Cache) -> Coroutine[Any, Any, R@Cache]))

That looks to me like it understands that P (the argument types) and R (the return types) are plumbed through correctly.

However, if I ask it to introspect Cache.decorate in the context where it's being used as a decorator, it returns:

(method) def decorate(**params: Unknown) -> ((f: ((...) -> Coroutine[Any, Any, Unknown])) -> ((...) -> Coroutine[Any, Any, Unknown]))

...which is to say, the relationship between input types and output types has been entirely discarded!


Solution

  • decorator depends on two contextual type variables: P and R. In the context of the Cache class, these are not known, but assumed to be known by Pyright at call time.

    However, when Cache.decorate() is called, Pyright does not get enough information to resolve P and R (no explicit type arguments and no arguments), so these two are resolved as Unknown.

    A simple fix is to parametrize Cache explicitly:

    (playground)

    @Cache[[int], int].decorate()
    async def some_function(i: int) -> int:
        return i + 1
    
    reveal_type(some_function)  # (int) -> Coroutine[Any, Any, int]
    

    However, this does not get you to the root of the problem, which is that you are not correctly specifying what you want.

    Cache is generic over P and R, so there must exist a way in which Pyright can infer the types correspond to those parameters. Normally, these are handled by passing arguments to the function creating that class, commonly __new__/__init__ and factory @classmethods.

    # Hypothetical usages
    Cache(return_value, *args, **kwargs)
    Cache.factory(return_value, *args, **kwargs)
    
    # This is also possible, as shown, but rarer
    Cache[[int], int]()
    

    decorate() is a factory method, but it does not create instances of Cache on its own. It takes some arguments, but these do not affect the type of the decorated some_function() nor the type of the to-be-created Cache.

    Instead, decorate() is meant to create decorators that themselves create instances of Cache. It is these decorators that you want to parametrize, because they are the ones receiving the decorated functions and responsible for creating Caches.

    (playground)

    class Cache[**P, R]:
    
        @classmethod
        def decorate(cls, **params: Any):
            def decorator[**P2, R2](f: Callable[P2, Coroutine[Any, Any, R2]]) -> ...:
                #        ^^^^^^^^^^
                return f
            
            return decorator
    
    reveal_type(Cache.decorate())
    # (f: (**P2@decorator) -> Coroutine[Any, Any, R2@decorator]) -> ((**P2@decorator) -> Coroutine[Any, Any, R2@decorator])
    
    reveal_type(some_function)
    # (i: int) -> Coroutine[Any, Any, int]