How to type hint a function, added to class by class decorator in Python

Question

I have a class decorator, which adds a few functions and fields to decorated class.

@mydecorator
@dataclass
class A:
    a: str = ""

Added (via setattr()) is a .save() function and a set of info for dataclass fields as a separate dict.

I'd like VScode and mypy to properly recognize that, so that when I use:

a=A()
a.save()

or a.my_fields_dict those 2 are properly recognized.

Is there any way to do that? Maybe modify class A type annotations at runtime?

Answer 1

TL;DR

What you are trying to do is not possible with the current type system.

---

1. Intersection types

If the attributes and methods you are adding to the class via your decorator are static (in the sense that they are not just known at runtime), then what you are describing is effectively the extension of any given class T by mixing in a protocol P. That protocol defines the method save and so on.

To annotate this you would need an intersection of T & P. It would look something like this:

from typing import Protocol, TypeVar


T = TypeVar("T")


class P(Protocol):
    @staticmethod
    def bar() -> str: ...


def dec(cls: type[T]) -> type[Intersection[T, P]]:
    setattr(cls, "bar", lambda: "x")
    return cls  # type: ignore[return-value]


@dec
class A:
    @staticmethod
    def foo() -> int:
        return 1

You might notice that the import of Intersection is conspicuously missing. That is because despite being one of the most requested features for the Python type system, it is still missing as of today. There is currently no way to express this concept in Python typing.

---

2. Class decorator problems

The only workaround right now is a custom implementation alongside a corresponding plugin for the type checker(s) of your choice. I just stumbled across the typing-protocol-intersection package, which does just that for mypy.

If you install that and add plugins = typing_protocol_intersection.mypy_plugin to your mypy configuration, you could write your code like this:

from typing import Protocol, TypeVar

from typing_protocol_intersection import ProtocolIntersection


T = TypeVar("T")


class P(Protocol):
    @staticmethod
    def bar() -> str: ...


def dec(cls: type[T]) -> type[ProtocolIntersection[T, P]]:
    setattr(cls, "bar", lambda: "x")
    return cls  # type: ignore[return-value]


@dec
class A:
    @staticmethod
    def foo() -> int:
        return 1

But here we run into the next problem. Testing this with reveal_type(A.bar()) via mypy will yield the following:

error: "Type[A]" has no attribute "bar"  [attr-defined]
note: Revealed type is "Any"

Yet if we do this instead:

class A:
    @staticmethod
    def foo() -> int:
        return 1


B = dec(A)

reveal_type(B.bar())

we get no complaints from mypy and note: Revealed type is "builtins.str". Even though what we did before was equivalent!

This is not a bug of the plugin, but of the mypy internals. It is another long-standing issue, that mypy does not handle class decorators correctly.

A person in that issue thread even mentioned your use case in conjunction with the desired intersection type.

---

DIY

In other words, you'll just have to wait until those two holes are patched. Or you can hope that at least the decorator issue by mypy is fixed soon-ish and write your own VSCode plugin for intersection types in the meantime. Maybe you can get together with the person behind that mypy plugin I mentioned above.