How can I write a decorator that preserves the docstring of an overloaded function correctly? (and why does one of my approaches not work)

[Glinte]

Pyright Playground

I know about functools.wraps, but it does nothing here so I have omitted it.

I think the decorator function loses type information because P and T are two separate type variables and the relationship between them and the original function kind of get lost when pyright trace into multiple overloaded definitions, so it cannot assume the function returned by decorator is still the original function, and so it does not show the docstring?

So I see that tenacity uses something similar to decorator_with_type_var, and it works, the docstring is kept. But I hit a new error that I don't know how to solve. I expect that because wrapped only has Any annotations that it is identical to Callable[..., Any] and thus is compatible with any WrappedFn, but my guess is that it is not actually compatible because of some intricate variance rules?

Raw Code:

from collections.abc import Awaitable
import functools
from typing import Any, Callable, ParamSpec, TypeVar, overload


T = TypeVar("T")
P = ParamSpec("P")
WrappedFn = TypeVar("WrappedFn", bound=Callable[..., Any])

def decorator(fn: Callable[P, T]) -> Callable[P, T]:
    def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
        return fn(*args, **kwargs)

    return wrapped


def decorator_with_type_var(fn: WrappedFn) -> WrappedFn:
	# The error become even more cursed if you add @functools.wraps(fn) here
    def wrapped(*args: Any, **kwargs: Any) -> Any:
        return fn(*args, **kwargs)

    return wrapped
#          ^^^^^^^
# Type "(...) -> Any" is not assignable to return type "WrappedFn@decorator_with_type_var"
#   Type "(...) -> Any" is not assignable to type "WrappedFn@decorator_with_type_var"  (reportReturnType)

@overload
def f(x: int) -> int:
    """Docstring for f"""

@overload
def f(x: str) -> str:
    """Docstring for f"""

def f(x: int | str) -> int | str:
    """Docstring for f"""
    return x


g1 = decorator(f)  # Has overload information, no docstring
g2 = decorator_with_type_var(f)  # Has overload information and docstring

[erictraut]

Pyright preserves the docstring of the decorated function if it's not overloaded. It should do the same in the case that it's overloaded. That appears to be a bug. Could you please file a bug report to that effect so it is tracked? Thanks!

The diagnostic you're seeing here also appears to be a bug. The type (...) -> Any is the gradual form of a callable. It's the equivalent of Any, and it should be assignable to any callable (and vice versa), so wrapped should be assignable to the return type. Please file a separate bug for that issue. You can use a cast call or a # pyright: ignore comment to work around this issue in the meantime.

[Glinte]

#10701 #10702

[Glinte]

#10701 is now microsoft/pylance-release#7600