CODEWITHSUNDEEP
pythonmypypyright
Archaeron
Archaeron

Reputation: 767

Proper typing for class function decorator

What would be the proper way to type the print_before function decorator, so that the wrapped function has the proper type but I can't use the decorator on a class that would not work?

thank you

def print_before(func):
    def func_wrapper(self, *args):
        self.print_hi()
        return func(self, *args)
    return func_wrapper


class PrintThings:
    def print_hi(self):
        print("hi")

    @print_before
    def add_nums(self, a: int, b: int) -> int:
        return a + b


pt = PrintThings()
pt.add_nums(5, 4)


class ShouldNotWork:
    @print_before
    def add_nums(self, a: int, b: int) -> int:
        return a + b


snw = ShouldNotWork()
snw.add_nums(4, 5)

Upvotes: 1

Views: 684

Answers (1)

chepner
chepner

Reputation: 532418

The decorator has a very general type as written. The only thing you know about the wrapper is that it accepts at least one argument of some unknown type, and returns some unknown type. As for the decorator itself, you only know that it returns something with the same type as its argument.

Actually, there is one more thing we know. Whatever the type of self is, it must have a print_hi method. We can represent that using a Protocol.

We also introduce a type variable T, which represents the same class (or a subclass thereof) each time it is used in the signature of print_before.

from typing import Any, Tuple, Protocol, TypeVar

class PrintsHi(Protocol):
    def print_hi(self):
        pass


T = TypeVar('T', bound=PrintsHi)


def print_before(func: Callable[[T, Tuple[Any,...]], Any]) -> Callable[[T, Tuple[Any,...]], Any]:
    def func_wrapper(self: T, *args: Tuple[Any,...]):
        self.print_hi()
        return func(self, *args)
    return func_wrapper

Upvotes: 1

Related Questions