Pyright: Better type inference for decorators.

[giyyapan]

Is your feature request related to a problem? Please describe.
Decorated functions will lose all of its signatures if used with decorators. I know some decorators might change the signature of the original function, but most of them are just using *args and **kwargs to proxy all the parameters.

Please see this example:

from functools import wraps

def some_decorator_generator():
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper
    return decorator

def some_decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper

@some_decorator_generator()
def fn1(a: int):
    pass

@some_decorator
def fn2(a: int):
    pass

fn1() # parameter a is not provided, but no error reported
fn2() # parameter a is not provided, but no error reported

reveal_type(fn2)

Describe the solution you'd like
Not sure if it's easy, but I hope a function decorated by decorators that only proxy all the parameters can keep its original signatures. this will make decorators a lot safer to use.

Thanks!

[erictraut]

Pyright is doing the right thing in this case. It is inferring as much as it can given the information it is provided within your code plus the stdlib type stubs.

The wraps decorator in functools.pyi is declared as follows:

_AnyCallable = Callable[..., Any]

def wraps(wrapped: _AnyCallable, assigned: Sequence[str] = ..., updated: Sequence[str] = ...) -> Callable[[_AnyCallable], _AnyCallable]: ...

Note that _AnyCallable is not a TypeVar, so there's no generic type matching here. _AnyCallable is simply a type alias for Callable[..., Any]. Because of this definition, the @wraps(func) always results in a decorated function type of Callable[..., Any]. Pyright has nothing to infer in this case because a return type is provided. Inference is used only in cases where a type is omitted and Pyright needs to infer it from other information.

I don't know enough about functools or the wraps decorator method to say whether this is a bug in the type stub. Should it be using a TypeVar that is bound to an _AnyCallable? If you think this is a bug, please report it in the typeshed repo.

In any case, there's a simple workaround that you can apply to your sample. Define a TypeVar that is bound to a function type, and use it in the declaration of decorator and some_decorator.

_TFunc = TypeVar("_TFunc", bound=Callable[..., Any])

def some_decorator_generator():
    def decorator(func: _TFunc) -> _TFunc:
       ...

def some_decorator(func: _TFunc) -> _TFunc:
 ...

[giyyapan]

Thank you so much for your clear explanation and brilliant suggestion! That's exactly what I need. :heart:

[moi90]

If I do this:

def decorator(f: _TFunc) -> _TFunc:
    def wrapper(*args, **kwargs):
        print("wrapper", *args, **kwargs)
        return f(*args, **kwargs)

    return wrapper

I get the following error (at the return wrapper line):

Expression of type "(*args: Unknown, **kwargs: Unknown) -> Any" cannot be assigned to return type "_TFunc"
  Type "(*args: Unknown, **kwargs: Unknown) -> Any" cannot be assigned to type "_TFunc" Pylance(reportGeneralTypeIssues)

Did I do something wrong? Can this error be removed?

[erictraut]

In Python 3.9, there's not a great solution to this problem. The best workaround I can offer is the following:

    return cast(_TFunc, wrapper)

In Python 3.10, there's a new facility called a ParamSpec. It's described in PEP 612. You can use it in older versions of Python through the typing_extensions package.

Here's how it would look with the new capability:

from typing import Callable, TypeVar
from typing_extensions import ParamSpec

_P = ParamSpec("_P")
_R = TypeVar("_R")


def decorator(f: Callable[_P, _R]) -> Callable[_P, _R]:
    def wrapper(*args: _P.args, **kwargs: _P.kwargs):
        print("wrapper", *args, **kwargs)
        return f(*args, **kwargs)

    return wrapper

[moi90]

Thanks for the explanation and the hint to PEP 612! It's great to see how the language keeps evolving.