CODEWITHSUNDEEP
pythonpython-typing
rider45
rider45

Reputation: 119

How to avoid losing type hinting of decorated function

I noticed that when wrapping a function or method that have some type hinting then the wrapped method loses the type hinting informations when I am coding using Visual studio code.

For example with this code:

from typing import Callable
import functools

def decorate(function: Callable):
    @functools.wraps(function)
    def wrapper(object: "A", *args, **kwargs):
        return function(object, *args, **kwargs)
    return wrapper

class A:
    @decorate
    def g(self, count: int) -> str:
        return f"hello {count}"

a = A()

print(a.g(2))

When I am hovering within visual studio code over the name g then I lose the type hinting informations. Would you know a way to prevent this?

Sincerely

Upvotes: 7

Views: 2608

Answers (1)

Daniil Fajnberg
Daniil Fajnberg

Reputation: 18608

The best you can do with Python 3.8 (and 3.9) is the following:

from __future__ import annotations
from functools import wraps
from typing import Any, Callable, TypeVar


T = TypeVar("T")


def decorate(function: Callable[..., T]) -> Callable[..., T]:
    @wraps(function)
    def wrapper(obj: A, *args: Any, **kwargs: Any) -> T:
        return function(obj, *args, **kwargs)
    return wrapper


class A:
    @decorate
    def g(self, count: int) -> str:
        return f"hello {count}"

This will preserve the return type information, but no details about the parameter types. The @wraps decorator should at least keep the signature intact. If the decorator is supposed to be universal for methods of A, this is as good as it gets IMO.

If you want it to be more specific, you can always restrict the function type to Callable[[A, B, C], T], but then the decorator will not be as universal anymore.

If you upgrade to Python 3.10, you will have access to ParamSpec. Then you can do the following:

from __future__ import annotations
from functools import wraps
from typing import Callable, ParamSpec, TypeVar


P = ParamSpec("P")
T = TypeVar("T")


def decorate(function: Callable[P, T]) -> Callable[P, T]:
    @wraps(function)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        return function(*args, **kwargs)
    return wrapper


class A:
    @decorate
    def g(self, count: int) -> str:
        return f"hello {count}"

This actually preserves (as the name implies) all the parameter specifications.

Hope this helps.

Upvotes: 11

Related Questions