Mypy 1.10 reports error when functools.wraps() is used on a generic function

Question

TLDR;

I have a decorator that:

• changes the function signature
• the wrapped function uses some generic type arguments
• Other than the signature I would like to use funtools.wraps to preserve the rest of the information.

Is there any way to achieve that without mypy complaining?

---

More context

A minimal working example would look like this:

from functools import wraps
from typing import Callable, TypeVar


B = TypeVar('B', bound=str)

def str_as_int_wrapper(func: Callable[[int], int]) -> Callable[[B], B]:
    WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__qualname__', '__doc__',)
    WRAPPER_UPDATES = ('__dict__', '__annotations__')
    
    @wraps(func, assigned=WRAPPER_ASSIGNMENTS, updated=WRAPPER_UPDATES)
    def _wrapped_func(val: B) -> B:
        num = int(val)
        result = func(num)
        return val.__class__(result)
    
    return _wrapped_func

@str_as_int_wrapper
def add_one(val: int) -> int:
    return val + 1

This seems to work alright, but mypy (version 1.10.0) does not like it. Instead, it complains with

test.py:17: error: Incompatible return value type (got "_Wrapped[[int], int, [Never], Never]", expected "Callable[[B], B]")  [return-value]
test.py:17: note: "_Wrapped[[int], int, [Never], Never].__call__" has type "Callable[[Arg(Never, 'val')], Never]"

If I either remove the @wraps decorator or replace the B type annotations by str, the error disappears.

Question

Am I missing something? Is this some already reported bug or limitation from mypy (couldn't find anything)? Should it be reported?

Thanks!

Answer 1

This is, in some sense, a regression, though it's more of a limitation. Your code should have worked as-is, and Mypy 1.9 does pass your code, which means the behaviour change was added in 1.10.

According to this similar issue (which was reported as a bug three months ago, but hasn't been triaged), the cause is this PR, in which the definitions of wraps and related symbols in Mypy's copy of typeshed were changed from (simplified):

class IdentityFunction(Protocol):
    def __call__(self, x: _T, /) -> _T: ...

_AnyCallable: TypeAlias = Callable[..., object]

def wraps(
    wrapped: _AnyCallable,  # ...
) -> IdentityFunction: ...

...to (also simplified):

class _Wrapped(Generic[_PWrapped, _RWrapped, _PWrapper, _RWrapper]):
    __wrapped__: Callable[_PWrapped, _RWrapped]
    def __call__(self, *args: _PWrapper.args, **kwargs: _PWrapper.kwargs) -> _RWrapper: ...

class _Wrapper(Generic[_PWrapped, _RWrapped]):
    def __call__(self, f: Callable[_PWrapper, _RWrapper]) -> _Wrapped[_PWrapped, _RWrapped, _PWrapper, _RWrapper]: ...

def wraps(
    wrapped: Callable[_PWrapped, _RWrapped],  # ...
) -> _Wrapper[_PWrapped, _RWrapped]: ...

Apparently, this was done in an attempt to fix this typeshed issue.

Solution

If you are not interested in the explanation, apply these general solutions and you are done:

return _wrapped_func  # type: ignore[return-value]

return cast(Callable[[B], B], _wrapped_func)

if TYPE_CHECKING:
    _T = TypeVar('_T')

    class IdentityFunction(Protocol):
       def __call__(self, x: _T, /) -> _T: ...

    _AnyCallable = Callable[..., object]

    def wraps(
        wrapped: _AnyCallable,
        # Default values omitted here for brevity
        assigned: Sequence[str] = (...),  
        updated: Sequence[str] = (...),
    ) -> IdentityFunction: ...

Explanation

Originally, wraps(func) (where the type of func is entirely irrelevant) returned IdentityFunction, a non-generic Protocol whose __call__ is generic over _T:

class IdentityFunction(Protocol):
    def __call__(self, x: _T, /) -> _T: ...

Thus, in the following (implicit) call wraps(func)(_wrapped_func), _wrapped_func retained its original type: a generic function. Mypy got this correctly, and it still does.

_{(playgrounds: 1.9, 1.11)}

def str_as_int_wrapper(func: Callable[[int], int]) -> Callable[[B], B]:
    @wraps(func, ...)
    def _wrapped_func(val: B) -> B:
        ...
    
    reveal_type(_wrapped_func)  # def [B <: builtins.str] (val: B`-1) -> B`-1
    return _wrapped_func

After the change, wraps(func) now returns _Wrapper[<P_of_func>, <R_of_func>], whose __call__ returns _Wrapped[<P_of_func>, <R_of_func>, <P_of_wrapped_func>, <R_of_wrapped_func>], a different type than the original type of _wrapped_func.

This is where Mypy took the wrong step, just as it would before 1.10 if the typeshed copy were not modified. As noted by STerliakov:

The problem originates from trying to resolve B too early: it's bound to _Wrapped['s] generic argument and resolved, B does not survive after that (hence Never in the output).