Python Type annotations for decorators

Question

I have a decorator that checks an instance attribute self.enabled and returns 0 if it is not enabled. Otherwise it returns the return value of the method which is an int, the index of the a unique string in the passed in list.

def check_if_enabled(func: Callable[..., int]) -> Callable[..., int]:
    @wraps(func)
    def wrapper(cls, list_of_strings):
        if not cls.enabled:
            return 0
        return func(cls, list_of_strings)
    return wrapper

I'd like to make the Type Annotation more specific but I am not sure how to.

Callable[..., int] is obviously what I want to change, and I want to make it so the Callable takes two arguments, an instance of a class and a list of strings. Is this possible?

Answer 1

For a start, please remember to provide a full working example whenever possible. Documenting your problem with more code and less words makes it easier to understand what you're trying to achieve, and less likely for us to misunderstand.

I presume that you want to decorate a method alike the following:

class Example(object):
    @check_if_enabled
    def func(self,  # implicit Example
             list_of_strings: List[str]
             ) -> int:
        pass

---

The first to do is to establish how the Callable for it would look like. You pass the list of function parameters as the first argument, and the return value as the second. Note that the first parameter is self, i.e. the class instance. Therefore, it would either be:

Callable[[Example, List[str]], int]

or if you don't want to limit it to one specific class:

Callable[[object, List[str]], int]

Note that you will probably want to use this prototype before declaring Example. In order to permit that, you need to pass "Example", i.e. the class name as string. This is treated as a forward declaration for the class.

---

So your code annotated would be:

def check_if_enabled(func: Callable[["Example", List[str]], int]
                     ) -> Callable[["Example", List[str]], int]:
    @wraps(func)
    def wrapper(cls: "Example",
                list_of_strings: List[str]
                ) -> int:
        if not cls.enabled:
            return 0
        return func(cls, list_of_strings)
    return wrapper

or as a complete working example:

from functools import wraps
from typing import Callable, List


def check_if_enabled(func: Callable[["Example", List[str]], int]
                     ) -> Callable[["Example", List[str]], int]:
    @wraps(func)
    def wrapper(cls: "Example",
                list_of_strings: List[str]
                ) -> int:
        if not cls.enabled:
            return 0
        return func(cls, list_of_strings)
    return wrapper


class Example(object):
    def __init__(self, enabled: bool) -> None:
        self.enabled = enabled

    @check_if_enabled
    def func(self,  # implicit Example
             list_of_strings: List[str]
             ) -> int:
        print('yes, it is')
        return 10


ex = Example(True)
print(ex.func(['1', '2', '3']))

ex = Example(False)
print(ex.func(['1', '2', '3']))

which produces:

yes, it is
10
0

---

As a general note, you probably want your decorators to be generic rather than fit for only one specific method. For example, the above decorator could be extended to fit any method returning an int by using *args, **kwargs and relaxed typing:

def check_if_enabled(func: Callable[..., int]
                     ) -> Callable[..., int]:
    @wraps(func)
    def wrapper(cls: "Example",
                *args: Any,
                **kwargs: Any
                ) -> int:
        if not cls.enabled:
            return 0
        return func(cls, *args, **kwargs)
    return wrapper

Going even further, you could allow a generic return value and pass it as a parameter to the decorator:

def check_if_enabled(return_if_disabled: Any
                     ) -> Callable[[Callable[..., int]],
                                   Callable[..., int]]:
    def decorator(func: Callable[..., int]
                  ) -> Callable[..., int]:
        @wraps(func)
        def wrapper(cls: "Example",
                    *args: Any,
                    **kwargs: Any
                    ) -> Any:
            if not cls.enabled:
                return return_if_disabled
            return func(cls, *args, **kwargs)
        return wrapper
    return decorator


class Example(object):
    def __init__(self, enabled: bool) -> None:
        self.enabled = enabled

    @check_if_enabled(return_if_disabled=0)
    def func(self,  # implicit Example
             list_of_strings: List[str]
             ) -> int:
        print('yes, it is')
        return -1

Though of course this kills static typing entirely.

---

If in doubt, I'd suggest using mypy. It can suggest the correct type if you get it wrong.

Finally, I don't think @wraps does anything here. I've preserved it because it was present in the pasted code.

Hope this helps.