Add plugin support for apps.get_model()

django-stubs/apps/registry.pyi

@@ -24,7 +24,7 @@ class Apps:
     def get_app_config(self, app_label: str) -> AppConfig: ...
     # it's not possible to support it in plugin properly now
     def get_models(self, include_auto_created: bool = ..., include_swapped: bool = ...) -> list[type[Model]]: ...
-    def get_model(self, app_label: str, model_name: str | None = ..., require_ready: bool = ...) -> type[Any]: ...
+    def get_model(self, app_label: str, model_name: str | None = ..., require_ready: bool = ...) -> type[Model]: ...
     def register_model(self, app_label: str, model: type[Model]) -> None: ...
     def is_installed(self, app_name: str) -> bool: ...
     def get_containing_app_config(self, object_name: str) -> AppConfig | None: ...

mypy_django_plugin/django/context.py

@@ -3,7 +3,21 @@
 from collections import defaultdict
 from contextlib import contextmanager
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, Literal, Optional, Sequence, Set, Tuple, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    Literal,
+    Mapping,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
 
 from django.core.exceptions import FieldDoesNotExist, FieldError
 from django.db import models
@@ -270,6 +284,14 @@ def all_registered_model_classes(self) -> Set[Type[models.Model]]:
     def all_registered_model_class_fullnames(self) -> Set[str]:
         return {helpers.get_class_fullname(cls) for cls in self.all_registered_model_classes}
 
+    @cached_property
+    def model_class_fullnames_by_label_lower(self) -> Mapping[str, str]:
+        return {
+            klass._meta.label_lower: helpers.get_class_fullname(klass)
+            for klass in self.all_registered_model_classes
+            if klass is not models.Model
+        }
+
     def get_field_nullability(self, field: Union["Field[Any, Any]", ForeignObjectRel], method: Optional[str]) -> bool:
         if method in ("values", "values_list"):
             return field.null

mypy_django_plugin/lib/fullnames.py

@@ -1,3 +1,4 @@
+APPS_FULLNAME = "django.apps.registry.Apps"
 ABSTRACT_USER_MODEL_FULLNAME = "django.contrib.auth.models.AbstractUser"
 PERMISSION_MIXIN_CLASS_FULLNAME = "django.contrib.auth.models.PermissionsMixin"
 MODEL_METACLASS_FULLNAME = "django.db.models.base.ModelBase"

mypy_django_plugin/lib/helpers.py

@@ -12,6 +12,7 @@
     MDEF,
     Block,
     ClassDef,
+    Context,
     Expression,
     MemberExpr,
     MypyFile,
@@ -385,3 +386,36 @@ def add_new_manager_base(api: SemanticAnalyzerPluginInterface, fullname: str) ->
     if sym is not None and isinstance(sym.node, TypeInfo):
         bases = get_django_metadata_bases(sym.node, "manager_bases")
         bases[fullname] = 1
+
+
+def resolve_lazy_reference(
+    reference: str, *, api: Union[TypeChecker, SemanticAnalyzer], django_context: "DjangoContext", ctx: Context
+) -> Optional[TypeInfo]:
+    """
+    Attempts to resolve a lazy reference(e.g. "<app_label>.<object_name>") to a
+    'TypeInfo' instance.
+    """
+    if "." not in reference:
+        # <object_name> -- needs prefix of <app_label>. We can't implicitly solve
+        # what app label this should be, yet.
+        api.fail("Could not resolve lazy reference without an app label", ctx)
+        api.note(
+            ("Try to use a reference explicitly prefixed with app label:" f' "<app_label>.{reference}" instead'),
+            ctx,
+        )
+        return None
+
+    # Reference conforms to the structure of a lazy reference: '<app_label>.<object_name>'
+    fullname = django_context.model_class_fullnames_by_label_lower.get(reference.lower())
+    if fullname is not None:
+        model_info = lookup_fully_qualified_typeinfo(api, fullname)
+        if model_info is not None:
+            return model_info
+        elif isinstance(api, SemanticAnalyzer):
+            if not api.final_iteration:
+                # Getting this far, where Django matched the reference but we still can't
+                # find it, we want to defer
+                api.defer()
+    else:
+        api.fail("Could not match lazy reference with any model", ctx)
+    return None

mypy_django_plugin/main.py

@@ -23,7 +23,7 @@
 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.transformers import fields, forms, init_create, meta, querysets, request, settings
+from mypy_django_plugin.transformers import apps, fields, forms, init_create, meta, querysets, request, settings
 from mypy_django_plugin.transformers.functional import resolve_str_promise_attribute
 from mypy_django_plugin.transformers.managers import (
     create_new_manager_class_from_as_manager_method,
@@ -121,6 +121,10 @@ def get_additional_deps(self, file: MypyFile) -> List[Tuple[int, str, int]]:
                     return []
                 return [self._new_dependency(auth_user_module), self._new_dependency("django_stubs_ext")]
 
+        if file.fullname == "django.apps":
+            # Preload all registered models.py so that we can resolve lazy references
+            # passed to 'apps.get_model()'. e.g. 'apps.get_model("myapp.MyModel")'
+            return [self._new_dependency(models) for models in self.django_context.model_modules.keys()]
         # ensure that all mentioned to='someapp.SomeModel' are loaded with corresponding related Fields
         defined_model_classes = self.django_context.model_modules.get(file.fullname)
         if not defined_model_classes:
@@ -219,6 +223,10 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
                 mypy_django_plugin.transformers.orm_lookups.typecheck_queryset_filter,
                 django_context=self.django_context,
             )
+        elif method_name == "get_model":
+            info = self._get_typeinfo_or_none(class_fullname)
+            if info and info.has_base(fullnames.APPS_FULLNAME):
+                return partial(apps.resolve_model_for_get_model, django_context=self.django_context)
 
         return None
 

mypy_django_plugin/transformers/apps.py

@@ -0,0 +1,41 @@
+from typing import Optional
+
+from mypy.nodes import StrExpr, TypeInfo
+from mypy.plugin import MethodContext
+from mypy.types import Instance, TypeType
+from mypy.types import Type as MypyType
+
+from mypy_django_plugin.django.context import DjangoContext
+from mypy_django_plugin.lib import helpers
+
+
+def resolve_model_for_get_model(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
+    """
+    Attempts to refine the return type of an 'apps.get_model()' call
+    """
+    if not ctx.args:
+        return ctx.default_return_type
+
+    model_info: Optional[TypeInfo] = None
+    # An 'apps.get_model("...")' call
+    if ctx.args[0] and not ctx.args[1]:
+        expr = ctx.args[0][0]
+        if isinstance(expr, StrExpr):
+            model_info = helpers.resolve_lazy_reference(
+                expr.value, api=helpers.get_typechecker_api(ctx), django_context=django_context, ctx=expr
+            )
+    # An 'apps.get_model("...", "...")' call
+    elif ctx.args[0] and ctx.args[1]:
+        app_label = ctx.args[0][0]
+        model_name = ctx.args[1][0]
+        if isinstance(app_label, StrExpr) and isinstance(model_name, StrExpr):
+            model_info = helpers.resolve_lazy_reference(
+                f"{app_label.value}.{model_name.value}",
+                api=helpers.get_typechecker_api(ctx),
+                django_context=django_context,
+                ctx=model_name,
+            )
+
+    if model_info is None:
+        return ctx.default_return_type
+    return TypeType(Instance(model_info, []))

tests/typecheck/apps/test_config.yml

@@ -27,3 +27,57 @@
       reveal_type(BarConfig("bar", None).default_auto_field)  # N: Revealed type is "builtins.str"
       reveal_type(BazConfig("baz", None).default_auto_field)  # N: Revealed type is "builtins.str"
       reveal_type(FooBarConfig("baz", None).default_auto_field)  # N: Revealed type is "builtins.str"
+
+-   case: test_get_model
+    main: |
+        from django.apps import apps
+        reveal_type(apps.get_model("app1.First"))
+        reveal_type(apps.get_model("app1.first"))
+        reveal_type(apps.get_model("app1", "First"))
+        reveal_type(apps.get_model(app_label="app1", model_name="First"))
+        reveal_type(apps.get_model(app_label="app1", model_name="first"))
+        reveal_type(apps.get_model(model_name="First", app_label="app1"))
+
+        reveal_type(apps.get_model("app2.Second"))
+        reveal_type(apps.get_model("app2", "Second"))
+
+        reveal_type(apps.get_model("app1.Nonexisting"))
+        reveal_type(apps.get_model("app2", "Unknown"))
+
+        reveal_type(apps.get_model("sites.Site"))  # Note that sites is not installed
+        reveal_type(apps.get_model("contenttypes.ContentType"))
+    out: |
+        main:2: note: Revealed type is "Type[app1.models.First]"
+        main:3: note: Revealed type is "Type[app1.models.First]"
+        main:4: note: Revealed type is "Type[app1.models.First]"
+        main:5: note: Revealed type is "Type[app1.models.First]"
+        main:6: note: Revealed type is "Type[app1.models.First]"
+        main:7: note: Revealed type is "Type[app1.models.First]"
+        main:9: note: Revealed type is "Type[app2.models.Second]"
+        main:10: note: Revealed type is "Type[app2.models.Second]"
+        main:12: note: Revealed type is "Type[django.db.models.base.Model]"
+        main:12: error: Could not match lazy reference with any model
+        main:13: note: Revealed type is "Type[django.db.models.base.Model]"
+        main:13: error: Could not match lazy reference with any model
+        main:15: note: Revealed type is "Type[django.db.models.base.Model]"
+        main:15: error: Could not match lazy reference with any model
+        main:16: note: Revealed type is "Type[django.contrib.contenttypes.models.ContentType]"
+    installed_apps:
+        - app1
+        - app2
+    files:
+        -   path: app1/__init__.py
+        -   path: app1/models.py
+            content: |
+                from django.db import models
+
+                class First(models.Model):
+                    field = models.IntegerField()
+
+        -   path: app2/__init__.py
+        -   path: app2/models.py
+            content: |
+                from django.db import models
+
+                class Second(models.Model):
+                    field = models.CharField()