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.

[SamB]

In short, you're confounding the type analyzer by indicating that the decorated type is a class instance when it's really a method.

Wait a minute ... that's not what they're indicating. They defined a protocol called CountedFunction[T], which (outside a stub file) does create a class, but this class cannot be instantiated, and references to it in type signatures actually refer to objects having the structure it describes.

I believe the problem is that this does not indicate that the object in question will be an ordinary function, that is, will have runtime type types.FunctionType, so the type checker doesn't know that the class statement will turn it into a method at runtime.

Worse: it seems impossible to express the concept "function implementing protocol Foo" at all.

The following almost works:

def counter[**P, R](func: Callable[P, R]) -> Callable[P, R]:
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        wrapper.count += 1
        return func(*args, **kwargs)
    wrapper.count = 0
    return wrapper

... but probably shouldn't, as it doesn't say anything more about the type() of the return value than OP's signature does.

(By "almost", I mean that if we do this, BasedPyright and MyPy both complain a lot about .count because that's not part of the API of Callable.)