From 96e236e2f56b59e04cbb649d94d1b8c9caeaf2fa Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Thu, 21 Sep 2023 08:33:52 +0200 Subject: [PATCH 1/3] Add plugin support for `apps.get_model()` Attempts to resolve the provided lazy reference to a model type. As such it can now return a specific model depending on what it was called with --- django-stubs/apps/registry.pyi | 2 +- mypy_django_plugin/django/context.py | 24 ++++++++++++- mypy_django_plugin/lib/fullnames.py | 1 + mypy_django_plugin/lib/helpers.py | 34 ++++++++++++++++++ mypy_django_plugin/main.py | 10 +++++- mypy_django_plugin/transformers/apps.py | 41 +++++++++++++++++++++ tests/typecheck/apps/test_config.yml | 48 +++++++++++++++++++++++++ 7 files changed, 157 insertions(+), 3 deletions(-) create mode 100644 mypy_django_plugin/transformers/apps.py diff --git a/django-stubs/apps/registry.pyi b/django-stubs/apps/registry.pyi index 528704e31..d7a3963e4 100644 --- a/django-stubs/apps/registry.pyi +++ b/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: ... diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py index 18c8834c8..44e76a7a6 100644 --- a/mypy_django_plugin/django/context.py +++ b/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 diff --git a/mypy_django_plugin/lib/fullnames.py b/mypy_django_plugin/lib/fullnames.py index 33a2e1181..8a8b01dd8 100644 --- a/mypy_django_plugin/lib/fullnames.py +++ b/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" diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index 9f5289970..87608dd4d 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/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. ".") to a + 'TypeInfo' instance. + """ + if "." not in reference: + # -- needs prefix of . 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' ".{reference}" instead'), + ctx, + ) + return None + + # Reference conforms to the structure of a lazy reference: '.' + 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 diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index 454a6cb13..ed4071308 100644 --- a/mypy_django_plugin/main.py +++ b/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 diff --git a/mypy_django_plugin/transformers/apps.py b/mypy_django_plugin/transformers/apps.py new file mode 100644 index 000000000..b92075729 --- /dev/null +++ b/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, [])) diff --git a/tests/typecheck/apps/test_config.yml b/tests/typecheck/apps/test_config.yml index cbf2bac43..66be74b31 100644 --- a/tests/typecheck/apps/test_config.yml +++ b/tests/typecheck/apps/test_config.yml @@ -27,3 +27,51 @@ 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")) + 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 + 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() From 80df1b655bab22b1b6180cda45740f216c534f28 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Thu, 21 Sep 2023 09:02:24 +0200 Subject: [PATCH 2/3] fixup! Add plugin support for `apps.get_model()` --- tests/typecheck/apps/test_config.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/typecheck/apps/test_config.yml b/tests/typecheck/apps/test_config.yml index 66be74b31..4c368d1b7 100644 --- a/tests/typecheck/apps/test_config.yml +++ b/tests/typecheck/apps/test_config.yml @@ -43,6 +43,9 @@ 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]" @@ -56,6 +59,9 @@ 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 From e44cd6cf505d44bec4bd65e0866f1e915275bc31 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Thu, 21 Sep 2023 20:42:19 +0200 Subject: [PATCH 3/3] fixup! fixup! Add plugin support for `apps.get_model()` --- mypy_django_plugin/lib/helpers.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index 87608dd4d..85e354b93 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/mypy_django_plugin/lib/helpers.py @@ -398,11 +398,6 @@ def resolve_lazy_reference( if "." not in reference: # -- needs prefix of . 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' ".{reference}" instead'), - ctx, - ) return None # Reference conforms to the structure of a lazy reference: '.'