Type hints and tooltips for decorators/wrappers/proxy objects

[stereobutter]

I often use decorators/wrappers/proxy objects to extend classes from third party libraries (instead of inheritance). I find this is usually much easier to implement than mucking with super() and taking care to uphold the LSP.

Say I've got the following class in a third party library that I want to extend:

class Obj:
    def foo(self, question: str) -> int:
        """Return the answer to the Ultimate Question of Life, the Universe, and Everything"""
        return 42

and I write the following wrapper

class Wrapper:
    def __init__(self, obj: Obj):
        self._obj = obj

    def foo(self, *args, **kwargs):
        print("Hello from Wrapper")
        return self._obj.foo(*args, **kwargs)

The problem with this is that one looses both the type hints and the doc string this way. Unfortunately functool.wraps also does currently not work to restore these to us so that VS Code is able to display them.

I saw there are multiple issues about this already, in particular #3939, microsoft/pylance-release#442 and python/typeshed#6670. But since I don't want to wait around for typeshed to fix this eventually, I've hacked away and ended up with this

from typing import TypeVar, Sequence, Callable, cast, Any
from typing_extensions import ParamSpec, Concatenate, TypeAlias,
import functools

W = TypeVar("W")
P = ParamSpec("P")
R = TypeVar("R")


def wraps_method(
    wrapped: Callable[Concatenate[Any, P], R],
    assigned: Sequence[str] = ...,
    updated: Sequence[str] = ...,
) -> Callable[[Callable[Concatenate[Any, P], R]], Callable[Concatenate[W, P], R]]:
    return cast(
        Callable[[Callable[Concatenate[Any, P], R]], Callable[Concatenate[W, P], R]],
        functools.wraps,
    )

that when used with the wrapper class

class Wrapper:
    def __init__(self, obj: Obj):
        self._obj = obj

    @wraps_method(Obj.foo)
    def foo(self, *args, **kwargs):
        return self._obj.foo(*args, **kwargs)

tricks VS Code into displaying the original type hints and doc string ✨ 🚀

I though this might be of some use for others but I've also got a couple of questions/issues:

• Is this a reliable workaround that I can expect to work in future versions of pyright/pylance? How does pyright/pylance know that it should display the doc string by looking at the type hints?
• Is this use case also addressed and resolved by Updated functools.wraps and functools.update_wrapper to use ParamSpec… python/typeshed#6670?
• Can I ignore pyright warning me about TypeVar "W" appears only once in generic function signaturePylance[reportInvalidTypeVarUse]?
• When I tried to tidy wraps_method up by replacing the type definitions by more compact type aliases everything stopped working. Is this a bug or did I miss something there?

  WrappedT: TypeAlias = Callable[Concatenate[Any, P], R]
  WrapperT: TypeAlias = Callable[Concatenate[W, P], R]
  
  
  def wraps_method(
      wrapped: WrappedT,
      assigned: Sequence[str] = ...,
      updated: Sequence[str] = ...,
  ) -> Callable[[WrappedT], WrapperT]:
      return functools.wraps

  

[erictraut]

Yes, this is a reliable workaround. Pyright captures the docstring from function/method captured by the ParamSpec and uses it in the resulting (wrapped) callable.

Yes, this use case should be resolved by the functools.wraps change in typeshed. That PR has been delayed because of limitations in mypy and other type checkers. I'm hopeful that it will eventually be accepted.

The error about W is legitimate. You're using it in a manner where it has no meaning. It might as well be replaced by Any. To have meaning, W should appear at least once in one of the input parameter type annotations.

The type aliases you're defining are generic type aliases. For example, WrappedT has two type parameters. The first is a ParamSpec and the second is a regular TypeVar. When you use the type alias, you must provide corresponding type arguments, but you've failed to do so. The type expression WrappedT is equivalent to WrappedT[..., Any], which is probably not what you intended. If you replace WrappedT with WrappedT[P, R], it will work as you expect. If you enable the reportMissingTypeArgument check, pyright will tell you about these missing type arguments.

[stereobutter]

Thanks for your answer 🙇; especially your last paragraph cleared up a gap in my knowledge of how type aliases work that I previously didn't even know about. For anyone that might find their way here what I ended up using is

P = ParamSpec("P")
R = TypeVar("R")
MethodT: TypeAlias = Callable[Concatenate[Any, P], R]

def wraps_method(
    wrapped: MethodT[P, R],
    assigned: Sequence[str] = ...,
    updated: Sequence[str] = ...,
) -> Callable[[MethodT[P, R]], MethodT[P, R]]:
    return cast(Callable[[MethodT[P, R]], MethodT[P, R]], functools.wraps)

[erictraut]

Yes, generic type aliases are commonly misunderstood in Python. That's partly what is motiving the new syntax proposed in PEP 695.

[stereobutter]

@erictraut I though I was happy with my workaround but some of the methods I want to wrap use overload which apparently doesn't work with ParamSpec (see python/mypy#13540). Any ideas?

[stereobutter]

This seems to do the trick

def wraps(
    wrapped: T,
    assigned: Sequence[str] = ...,
    updated: Sequence[str] = ...,
) -> Callable[[T], T]:
    return cast(Callable[[Any], T], functools.wraps)

although it requires the use of # type: ignore for the definition of the wrapped method

class Obj:
    @overload
    def foo(self, question: str) -> int:
        ...

    @overload
    def foo(self, question: int) -> str:
        ...

    def foo(self, question):
        if isinstance(question, str):
            return int(question)
        else:
            return str(question)


class Wrapper:
    def __init__(self, obj: Obj):
        self._obj = obj

    @wraps(Obj.foo)  # type: ignore
    def foo(self, *args, **kwargs):
        print("hello from Wrapper")
        return self._obj.foo(*args, **kwargs)
    

One has to live with the fact though that the lie is exposed by looking at