Typing of method decorators not working properly [question]

[lonelyenvoy]

Hi. It seems that the typing of decorators is problematic with functions inside classes.

Suppose we want to see how many times a function has been called, we can write a simple decorator:

def counter(func):
    def wrapper(*args, **kwargs):
        wrapper.count += 1
        return func(*args, **kwargs)
    wrapper.count = 0
    return wrapper

and use it like:

@counter
def greet(name: str) -> str:
    return f'hello, {name}'

print(greet('John'))
print(f'greet() called {greet.count} times')

To type this counter decorator, we can:

T = TypeVar('T', bound=Callable[..., Any])

class CountedFunction(Protocol[T]):
    __call__: T
    count: int

def counter(func: T) -> CountedFunction[T]: ...

This is fine with functions outside classes. However, it is not with a class function method:

class A:
    @counter
    def greet(self, name: str) -> str:
        return f'hello, {name}'

a = A()
print(a.greet('John'))  # error
print(f'greet() called {a.greet.count} times')

We get static errors from mypy, although the code can be run without any runtime errors:

test.py:37: error: Too few arguments for "greet" of "A"
test.py:37: error: Argument 1 to "greet" of "A" has incompatible type "str"; expected "A"

Full code (click to expand)

from typing import Any, Protocol, TypeVar, Callable, cast

T = TypeVar('T', bound=Callable[..., Any])


class CountedFunction(Protocol[T]):
    __call__: T
    count: int


def counter(func: T) -> CountedFunction[T]:
    def wrapper(*args, **kwargs):
        wrapper.count += 1
        return func(*args, **kwargs)
    wrapper.count = 0  # type: ignore
    return cast(CountedFunction[T], wrapper)


# usage with a normal function
@counter
def greet(name: str) -> str:
    return f'hello, {name}'


print(greet('John'))  # OK
print(f'greet() called {greet.count} times')


# usage with a class function
class A:
    @counter
    def greet(self, name: str) -> str:
        return f'hello, {name}'


a = A()
print(a.greet('John'))  # error
print(f'greet() called {a.greet.count} times')

[erictraut]

I think mypy is doing the right thing by emitting an error here. Pyright also emits an error in the same place.

You are replacing the method greet with a class variable of the same name. That variable has type CountedFunction[T], but it's no longer a method. When evaluating the type of the expression a.greet, the binding process that would normally bind the object a to the self parameter of the method is skipped because greet isn't a method. The evaluated type of a.greet is therefore a function with two parameters: self and name, but you are providing only one argument.

In short, you're confounding the type analyzer by indicating that the decorated type is a class instance when it's really a method. The annotated types are not consistent with your implementation. A consistent implement would look like the following. This does generate a runtime exception in the way that mypy indicates it will.

class CountedFunction(Generic[T]):
    def __init__(self, func: T):
        self.count = 0
        self.func = func

    def __call__(self, *args, **kwargs):
        self.count += 1
        return self.func(*args, **kwargs)

def counter(func):
    return CountedFunction(func)

...
print(greet("John"))  # OK
...
print(a.greet("John"))  # TypeError: A.greet() missing 1 required positional argument: 'name'

I don't see a good way to properly annotate the implemented code. Perhaps someone else has some ideas.

[lonelyenvoy]

Anyone has an idea to annotate the counter? Thanks.