Support related_query_name in Model._meta.get_fields (typeddjango#2821) (#41)

Co-authored-by: Thibaut Decombe <[email protected]>

Author: asottile-sentry

mypy_django_plugin/lib/helpers.py

@@ -1,5 +1,5 @@
 from collections.abc import Iterable, Iterator
-from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple, TypedDict, cast
 
 from django.db.models.base import Model
 from django.db.models.fields import Field
@@ -54,7 +54,7 @@
     get_proper_type,
 )
 from mypy.types import Type as MypyType
-from typing_extensions import TypedDict
+from typing_extensions import Self
 
 from mypy_django_plugin.lib import fullnames
 
@@ -209,6 +209,21 @@ class DjangoModel(NamedTuple):
     def info(self) -> TypeInfo:
         return self.typ.type
 
+    @classmethod
+    def from_model_type(cls, model_type: Instance, django_context: "DjangoContext") -> Self | None:
+        model_info = model_type.type
+        is_annotated = is_annotated_model(model_info)
+
+        model_cls = (
+            django_context.get_model_class_by_fullname(model_info.bases[0].type.fullname)
+            if is_annotated
+            else django_context.get_model_class_by_fullname(model_info.fullname)
+        )
+        if model_cls is None:
+            return None
+
+        return cls(cls=model_cls, typ=model_type, is_annotated=is_annotated)
+
 
 def extract_model_type_from_queryset(queryset_type: Instance, api: TypeChecker) -> Instance | None:
     """Extract the django model `Instance` associated to a queryset `Instance`"""
@@ -242,18 +257,7 @@ def get_model_info_from_qs_ctx(
     if not (isinstance(ctx.type, Instance) and (model_type := extract_model_type_from_queryset(ctx.type, api))):
         return None
 
-    model_info = model_type.type
-    is_annotated = is_annotated_model(model_info)
-
-    model_cls = (
-        django_context.get_model_class_by_fullname(model_info.bases[0].type.fullname)
-        if is_annotated
-        else django_context.get_model_class_by_fullname(model_info.fullname)
-    )
-    if model_cls is None:
-        return None
-
-    return DjangoModel(cls=model_cls, typ=model_type, is_annotated=is_annotated)
+    return DjangoModel.from_model_type(model_type, django_context)
 
 
 def _get_class_init_type(call: CallExpr) -> CallableType | None:

mypy_django_plugin/transformers/meta.py

@@ -1,9 +1,11 @@
+from django.core.exceptions import FieldDoesNotExist
 from mypy.plugin import MethodContext
 from mypy.types import AnyType, Instance, TypeOfAny, get_proper_type
 from mypy.types import Type as MypyType
 
 from mypy_django_plugin.django.context import DjangoContext, get_field_type_from_model_type_info
 from mypy_django_plugin.lib import helpers
+from mypy_django_plugin.lib.helpers import DjangoModel
 
 
 def return_proper_field_type_from_get_field(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
@@ -20,5 +22,15 @@ def return_proper_field_type_from_get_field(ctx: MethodContext, django_context:
     if field_type is not None:
         return field_type
 
-    ctx.api.fail(f"{model_type.type.name} has no field named {field_name!r}", ctx.context)
-    return AnyType(TypeOfAny.from_error)
+    if (django_model := DjangoModel.from_model_type(model_type, django_context)) is None:
+        return ctx.default_return_type
+
+    try:
+        field = django_model.cls._meta.get_field(field_name)
+        if field_info := helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), field.__class__):
+            return Instance(field_info, [])
+    except FieldDoesNotExist as e:
+        ctx.api.fail(str(e), ctx.context)
+        return AnyType(TypeOfAny.from_error)
+
+    return ctx.default_return_type

pyproject.toml

@@ -130,7 +130,8 @@ ignore = [
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "_typeshed.Self".msg = "Use typing_extensions.Self (PEP 673) instead. If you type a metaclass, add a noqa"
-"typing.assert_type".msg = "Use typing_extensions.assert_type instead."
+"typing.assert_type".msg = "Only available in Python 3.11 and above. Use `typing_extensions.assert_type` instead."
+"typing.Self".msg = "Only available in Python 3.11 and above. Use `typing_extensions.Self` instead."
 
 [tool.ruff.lint.isort]
 known-first-party = ["django_stubs_ext", "mypy_django_plugin"]

tests/typecheck/models/test_meta_options.yml

@@ -44,11 +44,11 @@
 -   case: get_field_with_abstract_inheritance
     main: |
         from typing_extensions import reveal_type
-        from myapp.models import AbstractModel
-        class MyModel(AbstractModel):
-            pass
+        from myapp.models import MyModel
 
-        MyModel._meta.get_field('field')  # E: MyModel has no field named 'field'  [misc]
+        MyModel._meta.get_field('non_existant')  # E: MyModel has no field named 'non_existant'  [misc]
+
+        reveal_type(MyModel._meta.get_field('field'))  # N: Revealed type is "django.contrib.postgres.fields.array.ArrayField[typing.Sequence[builtins.float | builtins.int | builtins.str] | django.db.models.expressions.Combinable, builtins.list[builtins.int]]"
 
         field: str
         reveal_type(MyModel._meta.get_field(field))  # N: Revealed type is "django.db.models.fields.Field[Any, Any] | django.db.models.fields.reverse_related.ForeignObjectRel | django.contrib.contenttypes.fields.GenericForeignKey"
@@ -65,6 +65,34 @@
                     class Meta(TypedModelMeta):
                         abstract = True
 
+                class MyModel(AbstractModel):
+                    field = ArrayField(models.IntegerField(), default=[])
+
+-   case: get_field_reverse_fk_with_related_query_name
+    main: |
+        from typing_extensions import reveal_type
+        from myapp.models import ModelA
+
+        reveal_type(ModelA._meta.get_field("model_b")) # N: Revealed type is "django.db.models.fields.reverse_related.ManyToOneRel"
+        reveal_type(ModelA.modelb_set.field) # N: Revealed type is "django.db.models.fields.related.ForeignKey[myapp.models.ModelB, myapp.models.ModelB]"
+
+        reveal_type(ModelA._meta.get_field("model_b_bis")) # N: Revealed type is "django.db.models.fields.reverse_related.ManyToOneRel"
+        reveal_type(ModelA.model_b_bis.field) # N: Revealed type is "django.db.models.fields.related.ForeignKey[myapp.models.ModelB, myapp.models.ModelB]"
+
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+
+                class ModelA(models.Model): ...
+
+                class ModelB(models.Model):
+                    model_a = models.ForeignKey(ModelA, on_delete=models.CASCADE, related_query_name="model_b")
+                    model_a_bis = models.ForeignKey(ModelA, on_delete=models.CASCADE, related_name="model_b_bis")
+
 -   case: base_model_meta_incompatible_types
     main: |
         from django.db import models