Search code examples
pythonmypytyping

Mypy doesn't recognize class decorator


Problem

Suppose I want to implement a class decorator that adds some attributes and functions to an existing class.

In particular, let's say I have a protocol called HasNumber, and I need a decorator can_add that adds the missing methods to convert HasNumber class to CanAdd.

class HasNumber(Protocol):
    num: int

class CanAdd(HasNumber):
    def add(self, num: int) -> int: ...

Implementation

I implement the decorator as follows:

_HasNumberT = TypeVar("_HasNumberT", bound=HasNumber)


def can_add(cls: Type[_HasNumberT]) -> Type[CanAdd]:
    def add(self: _HasNumberT, num: int) -> int:
        return self.num + num

    setattr(cls, "add", add)

    return cast(Type[CanAdd], cls)


@can_add
class Foo:
    num: int = 12

Error

The code works just fine when I run it, but mypy is unhappy about it for some reason.

It gives the error "Foo" has no attribute "add" [attr-defined], as if it doesn't take the return value (annotated as Type[CanAdd]) of the can_add decorator into account.

foo = Foo()
print(foo.add(4))  # "Foo" has no attribute "add"  [attr-defined]
reveal_type(foo) # note: Revealed type is "test.Foo"

Question

In this issue, someone demonstrated a way of annotating this with Intersection. However, is there a way to achieve it without Intersection? (Supposing that I don't care about other attributes in Foo except the ones defined in the protocols)

Or, is it a limitation of mypy itself?

Related posts that don't solve my problem:


Solution

  • cast tells mypy that cls (with or without an add attribute) is safe to use as the return value for can_add. It does not guarantee that the protocol holds.

    As a result, mypy cannot tell whether Foo has been given an add attribute, only that it's OK to use the can_add decorator. The fact that can_add has a side effect of defining the add attribute isn't visible to mypy.

    You can, however, replace the decorator with direct inheritance, something like

    class HasNumber(Protocol):
        num: int
    
    
    _HasNumberT = TypeVar("_HasNumberT", bound=HasNumber)
    
    class Adder(HasNumber):
        def add(self, num: int) -> int:
            return self.num + num
    
    class Foo(Adder):
        num: int = 12
    
    foo = Foo()
    print(foo.add(4))