Refactor support for annotate to utilise mypy.types.Instance.extra_attrs (#2319)

Author: flaeppe

README.md

@@ -211,11 +211,11 @@ class MyManager(models.Manager["MyModel"]):
 
 ### How do I annotate cases where I called QuerySet.annotate?
 
-Django-stubs provides a special type, `django_stubs_ext.WithAnnotations[Model]`, which indicates that the `Model` has
-been annotated, meaning it allows getting/setting extra attributes on the model instance.
+Django-stubs provides a special type, `django_stubs_ext.WithAnnotations[Model, <Annotations>]`, which indicates that
+the `Model` has been annotated, meaning it requires extra attributes on the model instance.
 
-Optionally, you can provide a `TypedDict` of these attributes,
-e.g. `WithAnnotations[MyModel, MyTypedDict]`, to specify which annotated attributes are present.
+You should provide a `TypedDict` of these attributes, e.g. `WithAnnotations[MyModel, MyTypedDict]`, to specify which
+annotated attributes are present.
 
 Currently, the mypy plugin can recognize that specific names were passed to `QuerySet.annotate` and
 include them in the type, but does not record the types of these attributes.
@@ -235,21 +235,11 @@ class MyModel(models.Model):
     username = models.CharField(max_length=100)
 
 
-def func(m: WithAnnotations[MyModel]) -> str:
-    return m.asdf  # OK, since the model is annotated as allowing any attribute
-
-
-func(MyModel.objects.annotate(foo=Value("")).get(id=1))  # OK
-func(
-    MyModel.objects.get(id=1)
-)  # Error, since this model will not allow access to any attribute
-
-
 class MyTypedDict(TypedDict):
     foo: str
 
 
-def func2(m: WithAnnotations[MyModel, MyTypedDict]) -> str:
+def func(m: WithAnnotations[MyModel, MyTypedDict]) -> str:
     print(m.bar)  # Error, since field "bar" is not in MyModel or MyTypedDict.
     return m.foo  # OK, since we said field "foo" was allowed
 

ext/django_stubs_ext/annotations.py

@@ -17,7 +17,7 @@ class Annotations(Generic[_Annotations]):
 
 WithAnnotations = Annotated[_T, Annotations[_Annotations]]
 """Alias to make it easy to annotate the model `_T` as having annotations
-`_Annotations` (a `TypedDict` or `Any` if not provided).
+`_Annotations` (a `TypedDict`).
 
-Use as `WithAnnotations[MyModel]` or `WithAnnotations[MyModel, MyTypedDict]`.
+Use as `WithAnnotations[MyModel, MyTypedDict]`.
 """

mypy_django_plugin/django/context.py

@@ -36,7 +36,6 @@
 
 from mypy_django_plugin.exceptions import UnregisteredModelError
 from mypy_django_plugin.lib import fullnames, helpers
-from mypy_django_plugin.lib.fullnames import WITH_ANNOTATIONS_FULLNAME
 
 # This import fails when `psycopg2` is not installed, avoid crashing the plugin.
 try:
@@ -142,15 +141,6 @@ def model_modules(self) -> Dict[str, Dict[str, Type[Model]]]:
 
     def get_model_class_by_fullname(self, fullname: str) -> Optional[Type[Model]]:
         """Returns None if Model is abstract"""
-        annotated_prefix = WITH_ANNOTATIONS_FULLNAME + "["
-        if fullname.startswith(annotated_prefix):
-            # For our "annotated models", extract the original model fullname
-            fullname = fullname[len(annotated_prefix) :].rstrip("]")
-            if "," in fullname:
-                # Remove second type arg, which might be present
-                fullname = fullname[: fullname.index(",")]
-            fullname = fullname.replace("__", ".")
-
         module, _, model_cls_name = fullname.rpartition(".")
         return self.model_modules.get(module, {}).get(model_cls_name)
 

mypy_django_plugin/lib/fullnames.py

@@ -19,7 +19,7 @@
 MANAGER_CLASS_FULLNAME = "django.db.models.manager.Manager"
 RELATED_MANAGER_CLASS = "django.db.models.fields.related_descriptors.RelatedManager"
 
-WITH_ANNOTATIONS_FULLNAME = "django_stubs_ext.WithAnnotations"
+WITH_ANNOTATIONS_FULLNAME = "django_stubs_ext.annotations.WithAnnotations"
 ANNOTATIONS_FULLNAME = "django_stubs_ext.annotations.Annotations"
 
 BASEFORM_CLASS_FULLNAME = "django.forms.forms.BaseForm"

mypy_django_plugin/lib/helpers.py

@@ -41,14 +41,14 @@
 from typing_extensions import TypedDict
 
 from mypy_django_plugin.lib import fullnames
-from mypy_django_plugin.lib.fullnames import WITH_ANNOTATIONS_FULLNAME
 
 if TYPE_CHECKING:
     from mypy_django_plugin.django.context import DjangoContext
 
 
 class DjangoTypeMetadata(TypedDict, total=False):
     is_abstract_model: bool
+    is_annotated_model: bool
     from_queryset_manager: str
     reverse_managers: Dict[str, str]
     baseform_bases: Dict[str, int]
@@ -104,6 +104,14 @@ def get_manager_to_model(manager: TypeInfo) -> Optional[str]:
     return get_django_metadata(manager).get("manager_to_model")
 
 
+def mark_as_annotated_model(model: TypeInfo) -> None:
+    get_django_metadata(model)["is_annotated_model"] = True
+
+
+def is_annotated_model(model: TypeInfo) -> bool:
+    return get_django_metadata(model).get("is_annotated_model", False)
+
+
 class IncompleteDefnException(Exception):
     pass
 
@@ -285,10 +293,6 @@ def get_nested_meta_node_for_current_class(info: TypeInfo) -> Optional[TypeInfo]
     return None
 
 
-def is_annotated_model_fullname(model_cls_fullname: str) -> bool:
-    return model_cls_fullname.startswith(WITH_ANNOTATIONS_FULLNAME + "[")
-
-
 def create_type_info(name: str, module: str, bases: List[Instance]) -> TypeInfo:
     # make new class expression
     classdef = ClassDef(name, Block([]))
@@ -382,9 +386,12 @@ def convert_any_to_type(typ: MypyType, referred_to_type: MypyType) -> MypyType:
 
 
 def make_typeddict(
-    api: CheckerPluginInterface, fields: "OrderedDict[str, MypyType]", required_keys: Set[str]
+    api: Union[SemanticAnalyzer, CheckerPluginInterface], fields: Dict[str, MypyType], required_keys: Set[str]
 ) -> TypedDictType:
-    fallback_type = api.named_generic_type("typing._TypedDict", [])
+    if isinstance(api, CheckerPluginInterface):
+        fallback_type = api.named_generic_type("typing._TypedDict", [])
+    else:
+        fallback_type = api.named_type("typing._TypedDict", [])
     typed_dict_type = TypedDictType(fields, required_keys=required_keys, fallback=fallback_type)
     return typed_dict_type
 

mypy_django_plugin/main.py

@@ -296,7 +296,7 @@ def get_type_analyze_hook(self, fullname: str) -> Optional[Callable[[AnalyzeType
             "typing_extensions.Annotated",
             "django_stubs_ext.annotations.WithAnnotations",
         ):
-            return partial(handle_annotated_type, django_context=self.django_context)
+            return partial(handle_annotated_type, fullname=fullname)
         else:
             return None
 

mypy_django_plugin/transformers/models.py

@@ -22,19 +22,29 @@
     TypeInfo,
     Var,
 )
-from mypy.plugin import AnalyzeTypeContext, AttributeContext, CheckerPluginInterface, ClassDefContext
+from mypy.plugin import AnalyzeTypeContext, AttributeContext, ClassDefContext
 from mypy.plugins import common
 from mypy.semanal import SemanticAnalyzer
 from mypy.typeanal import TypeAnalyser
-from mypy.types import AnyType, Instance, ProperType, TypedDictType, TypeOfAny, TypeType, TypeVarType, get_proper_type
+from mypy.types import (
+    AnyType,
+    ExtraAttrs,
+    Instance,
+    ProperType,
+    TypedDictType,
+    TypeOfAny,
+    TypeType,
+    TypeVarType,
+    get_proper_type,
+)
 from mypy.types import Type as MypyType
-from mypy.typevars import fill_typevars
+from mypy.typevars import fill_typevars, fill_typevars_with_any
 
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.errorcodes import MANAGER_MISSING
 from mypy_django_plugin.exceptions import UnregisteredModelError
 from mypy_django_plugin.lib import fullnames, helpers
-from mypy_django_plugin.lib.fullnames import ANNOTATIONS_FULLNAME, ANY_ATTR_ALLOWED_CLASS_FULLNAME, MODEL_CLASS_FULLNAME
+from mypy_django_plugin.lib.fullnames import ANNOTATIONS_FULLNAME, MODEL_CLASS_FULLNAME
 from mypy_django_plugin.transformers.fields import FieldDescriptorTypes, get_field_descriptor_types
 from mypy_django_plugin.transformers.managers import (
     MANAGER_METHODS_RETURNING_QUERYSET,
@@ -200,6 +210,50 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
         raise NotImplementedError(f"Implement this in subclass {self.__class__.__name__}")
 
 
+class AddAnnotateUtilities(ModelClassInitializer):
+    """
+    Creates a model subclass that will be used when the model's manager/queryset calls
+    'annotate' to hold on to attributes that Django adds to a model instance.
+
+    Example:
+
+        class MyModel(models.Model):
+            ...
+
+        class MyModel@AnnotatedWith(MyModel, django_stubs_ext.Annotations[_Annotations]):
+            ...
+    """
+
+    def run(self) -> None:
+        annotations = self.lookup_typeinfo_or_incomplete_defn_error("django_stubs_ext.Annotations")
+        object_does_not_exist = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.OBJECT_DOES_NOT_EXIST)
+        multiple_objects_returned = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.MULTIPLE_OBJECTS_RETURNED)
+        annotated_model_name = self.model_classdef.info.name + "@AnnotatedWith"
+        annotated_model = self.lookup_typeinfo(self.model_classdef.info.module_name + "." + annotated_model_name)
+        if annotated_model is None:
+            model_type = fill_typevars_with_any(self.model_classdef.info)
+            assert isinstance(model_type, Instance)
+            annotations_type = fill_typevars(annotations)
+            assert isinstance(annotations_type, Instance)
+            annotated_model = self.add_new_class_for_current_module(
+                annotated_model_name, bases=[model_type, annotations_type]
+            )
+            annotated_model.defn.type_vars = annotations.defn.type_vars
+            annotated_model.add_type_vars()
+            helpers.mark_as_annotated_model(annotated_model)
+            if self.is_model_abstract:
+                # Below are abstract attributes, and in a stub file mypy requires
+                # explicit ABCMeta if not all abstract attributes are implemented i.e.
+                # class is kept abstract. So we add the attributes to get mypy off our
+                # back
+                helpers.add_new_sym_for_info(
+                    annotated_model, "DoesNotExist", TypeType(Instance(object_does_not_exist, []))
+                )
+                helpers.add_new_sym_for_info(
+                    annotated_model, "MultipleObjectsReturned", TypeType(Instance(multiple_objects_returned, []))
+                )
+
+
 class InjectAnyAsBaseForNestedMeta(ModelClassInitializer):
     """
     Replaces
@@ -1034,6 +1088,7 @@ def run(self) -> None:
 
 def process_model_class(ctx: ClassDefContext, django_context: DjangoContext) -> None:
     initializers = [
+        AddAnnotateUtilities,
         InjectAnyAsBaseForNestedMeta,
         AddDefaultPrimaryKey,
         AddPrimaryKeyAlias,
@@ -1059,77 +1114,72 @@ def set_auth_user_model_boolean_fields(ctx: AttributeContext, django_context: Dj
     return Instance(boolinfo, [])
 
 
-def handle_annotated_type(ctx: AnalyzeTypeContext, django_context: DjangoContext) -> MypyType:
+def handle_annotated_type(ctx: AnalyzeTypeContext, fullname: str) -> MypyType:
+    """
+    Replaces the 'WithAnnotations' type with a type that can represent an annotated
+    model.
+    """
+    is_with_annotations = fullname == fullnames.WITH_ANNOTATIONS_FULLNAME
     args = ctx.type.args
+    if not args:
+        return AnyType(TypeOfAny.from_omitted_generics) if is_with_annotations else ctx.type
     type_arg = ctx.api.analyze_type(args[0])
     if not isinstance(type_arg, Instance) or not type_arg.type.has_base(MODEL_CLASS_FULLNAME):
         return type_arg
 
     fields_dict = None
     if len(args) > 1:
         second_arg_type = get_proper_type(ctx.api.analyze_type(args[1]))
-        if isinstance(second_arg_type, TypedDictType):
+        if isinstance(second_arg_type, TypedDictType) and is_with_annotations:
             fields_dict = second_arg_type
         elif isinstance(second_arg_type, Instance) and second_arg_type.type.fullname == ANNOTATIONS_FULLNAME:
             annotations_type_arg = get_proper_type(second_arg_type.args[0])
             if isinstance(annotations_type_arg, TypedDictType):
                 fields_dict = annotations_type_arg
             elif not isinstance(annotations_type_arg, AnyType):
                 ctx.api.fail("Only TypedDicts are supported as type arguments to Annotations", ctx.context)
+            elif annotations_type_arg.type_of_any == TypeOfAny.from_omitted_generics:
+                ctx.api.fail("Missing required TypedDict parameter for generic type Annotations", ctx.context)
+
+    if fields_dict is None:
+        return type_arg
 
     assert isinstance(ctx.api, TypeAnalyser)
     assert isinstance(ctx.api.api, SemanticAnalyzer)
-    return get_or_create_annotated_type(ctx.api.api, type_arg, fields_dict=fields_dict)
+    return get_annotated_type(ctx.api.api, type_arg, fields_dict=fields_dict)
 
 
-def get_or_create_annotated_type(
-    api: Union[SemanticAnalyzer, CheckerPluginInterface], model_type: Instance, fields_dict: Optional[TypedDictType]
+def get_annotated_type(
+    api: Union[SemanticAnalyzer, TypeChecker], model_type: Instance, fields_dict: TypedDictType
 ) -> ProperType:
     """
-
-    Get or create the type for a model for which you getting/setting any attr is allowed.
-
-    The generated type is an subclass of the model and django._AnyAttrAllowed.
-    The generated type is placed in the django_stubs_ext module, with the name WithAnnotations[ModelName].
-    If the user wanted to annotate their code using this type, then this is the annotation they would use.
-    This is a bit of a hack to make a pretty type for error messages and which would make sense for users.
+    Get a model type that can be used to represent an annotated model
     """
-    model_module_name = "django_stubs_ext"
-
-    if helpers.is_annotated_model_fullname(model_type.type.fullname):
-        # If it's already a generated class, we want to use the original model as a base
-        model_type = model_type.type.bases[0]
-
-    if fields_dict is not None:
-        type_name = f"WithAnnotations[{model_type.type.fullname.replace('.', '__')}, {fields_dict}]"
+    if model_type.extra_attrs:
+        extra_attrs = ExtraAttrs(
+            attrs={**model_type.extra_attrs.attrs, **(fields_dict.items if fields_dict is not None else {})},
+            immutable=model_type.extra_attrs.immutable.copy(),
+            mod_name=None,
+        )
     else:
-        type_name = f"WithAnnotations[{model_type.type.fullname.replace('.', '__')}]"
-
-    annotated_typeinfo = helpers.lookup_fully_qualified_typeinfo(
-        cast(TypeChecker, api), model_module_name + "." + type_name
-    )
-    if annotated_typeinfo is None:
-        model_module_file = api.modules.get(model_module_name)  # type: ignore[union-attr]
-        if model_module_file is None:
-            return AnyType(TypeOfAny.from_error)
-
-        if isinstance(api, SemanticAnalyzer):
-            annotated_model_type = api.named_type_or_none(ANY_ATTR_ALLOWED_CLASS_FULLNAME, [])
-            assert annotated_model_type is not None
-        else:
-            annotated_model_type = api.named_generic_type(ANY_ATTR_ALLOWED_CLASS_FULLNAME, [])
-
-        annotated_typeinfo = helpers.add_new_class_for_module(
-            model_module_file,
-            type_name,
-            bases=[model_type] if fields_dict is not None else [model_type, annotated_model_type],
-            fields=fields_dict.items if fields_dict is not None else None,
-            no_serialize=True,
+        extra_attrs = ExtraAttrs(
+            attrs=fields_dict.items if fields_dict is not None else {},
+            immutable=None,
+            mod_name=None,
         )
-        if fields_dict is not None:
-            # To allow structural subtyping, make it a Protocol
-            annotated_typeinfo.is_protocol = True
-            # Save for later to easily find which field types were annotated
-            annotated_typeinfo.metadata["annotated_field_types"] = fields_dict.items
-    annotated_type = Instance(annotated_typeinfo, [])
-    return annotated_type
+
+    annotated_model: Optional[TypeInfo]
+    if helpers.is_annotated_model(model_type.type):
+        annotated_model = model_type.type
+        if model_type.args and isinstance(model_type.args[0], TypedDictType):
+            fields_dict = helpers.make_typeddict(
+                api,
+                fields={**model_type.args[0].items, **fields_dict.items},
+                required_keys={*model_type.args[0].required_keys, *fields_dict.required_keys},
+            )
+    else:
+        annotated_model = helpers.lookup_fully_qualified_typeinfo(api, model_type.type.fullname + "@AnnotatedWith")
+
+    if annotated_model is None:
+        return model_type
+    return Instance(annotated_model, [fields_dict], extra_attrs=extra_attrs)

mypy_django_plugin/transformers/orm_lookups.py

@@ -5,7 +5,6 @@
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.exceptions import UnregisteredModelError
 from mypy_django_plugin.lib import fullnames, helpers
-from mypy_django_plugin.lib.helpers import is_annotated_model_fullname
 
 
 def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
@@ -33,13 +32,10 @@ def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext)
             provided_type = resolve_combinable_type(provided_type, django_context)
 
         lookup_type: MypyType
-        if is_annotated_model_fullname(model_cls_fullname):
-            lookup_type = AnyType(TypeOfAny.implementation_artifact)
-        else:
-            try:
-                lookup_type = django_context.resolve_lookup_expected_type(ctx, model_cls, lookup_kwarg)
-            except UnregisteredModelError:
-                lookup_type = AnyType(TypeOfAny.from_error)
+        try:
+            lookup_type = django_context.resolve_lookup_expected_type(ctx, model_cls, lookup_kwarg)
+        except UnregisteredModelError:
+            lookup_type = AnyType(TypeOfAny.from_error)
         # Managers as provided_type is not supported yet
         if isinstance(provided_type, Instance) and helpers.has_any_of_bases(
             provided_type.type, (fullnames.MANAGER_CLASS_FULLNAME, fullnames.QUERYSET_CLASS_FULLNAME)