Refactor support for annotate to utilise mypy.types.Instance.extra_attrs

[flaeppe]

I figured out that the implementation in mypy for functools.partial, in python/mypy#16939, added some pieces that we can utilise to get better support for .annotate.

Namely that python/mypy#16939 added support for writing Instance.extra_attrs to cache.

There are 2 main pieces at play in this implementation:

1. Annotated attributes are written to mypy.types.Instance.extra_attrs. This allows .annotate calls to bring in temporary attributes on same model instances.
2. The approach in 1. does not help out when comparing instance A´1 with A´2 for annotated attributes. Essentially nothing in .extra_attrs is considered when comparing instances, there could be completely different attributes in there and mypy still considers the instances equal/compatible. To be able to resolve that, the plugin creates a subclass for each model that also has django_stubs_ext.Annotations as base with a generic parameter. This subclass will be replace any encounter of WithAnnotations[<Model>, <TypedDict>], where <TypedDict> populates the generic parameter of the model subclass. This allows for falling back on mypy for type compatibility depending on annotated attributes.

The description in 2. is something like this in code

class MyModel(models.Model):
    username = models.CharField(max_length=100)

class MyModel_AnnotatedWith(models.Model, django_stubs_ext.Annotations[django_stubs_ext.annotations._Annotations]):
    ...

Then with a use case example

class AnnotatedAttrs(TypedDict):
    foo: int

def func(instance: django_stubs_ext.WithAnnotations[MyModel, AnnotatedAttrs]) -> int:
    return instance.foo

func(MyModel.objects.annotate(foo=Value("")).get())  # OK
func(MyModel.objects.annotate(bar=Value("")).get())  # Error

reveal_type(func)  # Revealed type is "def (instance: MyModel@AnnotatedWith[TypedDict('AnnotatedAttrs', {'foo': builtins.int})]) -> builtins.int"

The main downside here is the subclass and that it needs a name, which is set to {Model name}@AnnotatedWith

---

Most of this is backwards compatible with our old WithAnnotations except for the fact that;

• WithAnnotations[MyModel] is no longer supported, it didn't really make sense to ask for an annotation of nothing. Though this enabled support for accepting any attribute in the previous version, but IMO # type: ignore will make do for that.

It might be more compatibility issues but our test suite is kind of lacking cases for using .annotate.

Let me know what you think, any feedback is greatly appreciated here.

Related issues

There are a whole bunch of issues reported in regards to annotate: https://github.com/typeddjango/django-stubs/issues?q=is%3Aissue+is%3Aopen+annotate

I'm not going to hunt down and include which repro cases are resolved with this PR, but was thinking to do that subsequently to close off resolved issues. Though I will include closing one (major) cache issue we've had forever, since that should no longer be any issue with this implementation.

• Closes WithAnnotations Crash in 1.9.0 #760

[flaeppe]

I don't know if what I try to describe here is a mypy problem, I tried to reproduce this issue in mypy playground but wasn't able to.

This is as far as I got, but without getting any error: https://mypy-play.net/?mypy=latest&python=3.12&gist=e4745bfc8e2f42b02911adec2214e6e5

from typing import Any, Generic, Mapping, TypeVar, overload, TypedDict

_Annotations = TypeVar("_Annotations", covariant=True, bound=Mapping[str, Any])

class Annotations(Generic[_Annotations]): ...

T = TypeVar("T")

class Field(Generic[T]):
    @overload
    def __get__(self, instance: None, owner: Any) -> None: ...
    @overload
    def __get__(self, instance: Model, owner: Any) -> T: ...
    def __get__(self, instance: Model | None, owner: Any) -> T | None: ...

class Model:
    field = Field[str]()

class Foo(TypedDict):
    foo: Any

class A(Model, Annotations[_Annotations]):
    ...
    
a = A[Foo]()
reveal_type(a)
reveal_type(a.field)

[sobolevn]

Thank you! Very good idea!

[sobolevn]

I don't think that we should remove this feature. Because I see that a lot of users discuss this and use this. Maybe we can treat WithAnnotations[MyModel] as WithAnnotations[MyModel, Any] implicitly? Via type var defaults or something?

This way we can continue to support this code (maybe with a deprecation warning, why not). But breaking it - does not seem right to me :(

[sobolevn]

But, I agree that we should not have done this in the first place.

[flaeppe]

Yeah, I suppose, gonna have to see if I can get it working. The support for any/arbitrary attributes is a bit of work I think

[sobolevn]

We only need to support Any, no other arguments / types.

[flaeppe]

Yes, the Any type argument is fine. But that doesn't add support for accessing an arbitrary attribute, like the previous implementation supported.

Essentially you'll still get something like:

def func(m: WithAnnotations[MyModel]) -> str:
    return m.asdf  # E: MyModel@AnnotatedWith[Any] has no attribute "asdf"

I don't have any solution that removes the "has no attribute" error.

[flaeppe]

There's 1 problem that I see with that, unless I'm missing something, which is that Model@AnnotatedWith[Any] is an Instance and not a TypeInfo. But we need to add the __getattr__ method to the TypeInfo though if we do, all instances of Model@AnnotatedWith gets the __getattr__ method

[sobolevn]

Won't .extra_attrs do the job for single instance __getattr__? 🤔
I am not sure, just asking.

[flaeppe]

__getattr__ will work for the arbitrary attributes. But we still want/need to be able to match the model type on e.g. a function argument.

This whole thing is kind of twofold.

1. Subclassing the model and the generic TypedDict argument for proper matching for WithAnnotations as a function argument
2. extra_attrs for accessing declared annotated attributes e.g. inside a function body

__getattr__ would help out with attributes accessed inside the body. But I can't find a way to add it and keep 1. working and not affect all other instances. I've also tried adding the __getattr__ method to extra_attrs with no luck: ExtraAttrs(attrs={"__getattr__": CallableType(...)}, ...)

There's also the TypeInfo.fallback_to_any which is the same as subclassing Any or our AnyAttrAllowed. But same issue here, it affects the type and not instances of said type.

[sobolevn]

Thanks a lot for your research! Not really what I hoped for :(

In this case, let's do the cleanest solution and break WithAnnotations[M] usage in a way that we restrict known attributes.

[flaeppe]

I was also hoping to find a way that kept backwards compatibility. At least I'm hoping breaking this will make it a bit easier going forward.

Do you mean restrict known attributes like emitting "has no attribute" errors?

Then, calling wise, if we replace

WithAnnotations[MyModel] -> MyModel@AnnotatedWith[Any]

We can require that .annotate has been called, but we don't care with what. Just like before. I'm thinking this is the implementation we want.

Else if we replace like (this is what we do now)

WithAnnotations[MyModel] -> MyModel

This wouldn't distinguish between having called .annotate or not. Meaning that there's no requirement of having called .annotate but having done it is also accepted.

[sobolevn]

Let's go for it!

[flaeppe]

Agreed, let's do it. 🚀