pythonpycharmpython-decoratorspython-typing

How to write type hint for decorated implemented by a class


Here's a classical example of a decorator implemented by a class:

class Decorator:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        self.func(*args, **kwargs)

How to make __call__ have the same signature and type hints as func has?

I've tried the following code:

from typing import Callable, TypeVar, ParamSpec, Generic

PT = ParamSpec('PT')
RT = TypeVar('RT')


class Decorator(Generic[PT, RT]):
    def __init__(self, func: Callable[PT, RT]) -> None:
        self.func = func

    def __call__(self, *args: PT.args, **kwargs: PT.kwargs) -> RT:
        return self.func(*args, **kwargs)


@Decorator
def add(x: int, y: int) -> int:
    return x + y

But I failed to get the correct argument list of add in PyCharm.

Is it PyCharm's fault?

PyCharm screenshot


Solution

  • You want to explicitly replace the function with a callable with a specific signature, instead of replacing the function with an instance of a callable class that works the same but does not share the signature.

    This shows the difference:

    from typing import Callable, TypeVar, ParamSpec, Generic
    
    PT = ParamSpec('PT')
    RT = TypeVar('RT')
    
    
    class Decorator(Generic[PT, RT]):
        def __init__(self, func: Callable[PT, RT]) -> None:
            self.func = func
    
        def __call__(self, *args: PT.args, **kwargs: PT.kwargs) -> RT:
            return self.func(*args, **kwargs)
    
    
    def decorator(func: Callable[PT, RT]) -> Callable[PT, RT]:
        return Decorator(func)
    
    
    @decorator
    def add_func_decorated(x: int, y: int) -> int:
        return x + y
    
    
    @Decorator
    def add_class_decorated(x: int, y: int) -> int:
        return x + y
    
    
    print(add_func_decorated(1, 2))
    print(add_class_decorated(1, 2))
    

    You will will find that both type hints and type checking work correctly for add_func_decorated, in PyCharm as well. While, for add_class_decorated type checking works, but the type hint displays as you showed in your question.