From 870369651565dffc1d3b6223b7ae6d91b2d3654b Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Fri, 28 Jun 2024 09:37:28 +0200 Subject: [PATCH] Make plugin handle explicitly declared reverse descriptors for FKs (#2230) * Make plugin handle explicitly declared reverse descriptors for FKs If a reverse descriptor was explicitly declared on a model, the plugin skipped to create a more specific subclass for its related manager. * fixup! Make plugin handle explicitly declared reverse descriptors for FKs --- mypy_django_plugin/transformers/models.py | 52 ++++++++++--------- .../models/fields/test_related_descriptors.py | 18 +++++++ tests/typecheck/fields/test_related.yml | 30 +++++++++++ 3 files changed, 75 insertions(+), 25 deletions(-) create mode 100644 tests/assert_type/db/models/fields/test_related_descriptors.py diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 91fa28fa4..0888fff58 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -463,34 +463,43 @@ def many_to_many_descriptor(self) -> TypeInfo: def process_relation(self, relation: ForeignObjectRel) -> None: attname = relation.get_accessor_name() - if attname is None or attname in self.model_classdef.info.names: - # No reverse accessor or already declared. Note that this would also leave any - # explicitly declared(i.e. non-inferred) reverse accessors alone + if attname is None: + # No reverse accessor. return to_model_cls = self.django_context.get_field_related_model_cls(relation) to_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error(to_model_cls) + reverse_lookup_declared = attname in self.model_classdef.info.names if isinstance(relation, OneToOneRel): - self.add_new_node_to_model_class( - attname, - Instance( - self.reverse_one_to_one_descriptor, - [Instance(self.model_classdef.info, []), Instance(to_model_info, [])], - ), - ) + if not reverse_lookup_declared: + self.add_new_node_to_model_class( + attname, + Instance( + self.reverse_one_to_one_descriptor, + [Instance(self.model_classdef.info, []), Instance(to_model_info, [])], + ), + ) return - elif isinstance(relation, ManyToManyRel): - # TODO: 'relation' should be based on `TypeInfo` instead of Django runtime. - assert relation.through is not None - through_fullname = helpers.get_class_fullname(relation.through) - through_model_info = self.lookup_typeinfo_or_incomplete_defn_error(through_fullname) + if not reverse_lookup_declared: + # TODO: 'relation' should be based on `TypeInfo` instead of Django runtime. + assert relation.through is not None + through_fullname = helpers.get_class_fullname(relation.through) + through_model_info = self.lookup_typeinfo_or_incomplete_defn_error(through_fullname) + self.add_new_node_to_model_class( + attname, + Instance( + self.many_to_many_descriptor, [Instance(to_model_info, []), Instance(through_model_info, [])] + ), + is_classvar=True, + ) + return + elif not reverse_lookup_declared: + # ManyToOneRel self.add_new_node_to_model_class( - attname, - Instance(self.many_to_many_descriptor, [Instance(to_model_info, []), Instance(through_model_info, [])]), + attname, Instance(self.reverse_many_to_one_descriptor, [Instance(to_model_info, [])]), is_classvar=True ) - return related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS) # TODO: Support other reverse managers than `_default_manager` @@ -525,10 +534,6 @@ def process_relation(self, relation: ForeignObjectRel) -> None: self.ctx.cls, code=MANAGER_MISSING, ) - - self.add_new_node_to_model_class( - attname, Instance(self.reverse_many_to_one_descriptor, [Instance(to_model_info, [])]) - ) return # Create a reverse manager subclassed from the default manager of the related @@ -548,9 +553,6 @@ def process_relation(self, relation: ForeignObjectRel) -> None: derived_from="_default_manager", fullname=new_related_manager_info.fullname, ) - self.add_new_node_to_model_class( - attname, Instance(self.reverse_many_to_one_descriptor, [Instance(to_model_info, [])]) - ) def run_with_model_cls(self, model_cls: Type[Model]) -> None: # add related managers etc. diff --git a/tests/assert_type/db/models/fields/test_related_descriptors.py b/tests/assert_type/db/models/fields/test_related_descriptors.py new file mode 100644 index 000000000..61ded3d82 --- /dev/null +++ b/tests/assert_type/db/models/fields/test_related_descriptors.py @@ -0,0 +1,18 @@ +from typing import ClassVar + +from django.db import models +from django.db.models.fields.related_descriptors import RelatedManager, ReverseManyToOneDescriptor +from typing_extensions import assert_type + + +class Other(models.Model): + explicit_descriptor: ClassVar[ReverseManyToOneDescriptor["MyModel"]] + + +class MyModel(models.Model): + rel = models.ForeignKey[Other, Other](Other, on_delete=models.CASCADE, related_name="explicit_descriptor") + + +# Pyright doesn't allow "runtime" usage of @type_check_only 'RelatedManager' but we're +# only type checking these files so it should be fine. +assert_type(Other().explicit_descriptor, RelatedManager[MyModel]) # pyright: ignore[reportGeneralTypeIssues] diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index 4ef7101d5..843bf3a93 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -1035,6 +1035,36 @@ class Book(PrintedGood): name = models.CharField() +- case: test_explicit_reverse_many_to_one_descriptor + main: | + from myapp.models import Other + reveal_type(Other.explicit_descriptor) # N: Revealed type is "django.db.models.fields.related_descriptors.ReverseManyToOneDescriptor[myapp.models.MyModel]" + reveal_type(Other().explicit_descriptor) # N: Revealed type is "myapp.models.MyModel_RelatedManager" + reveal_type(Other().explicit_descriptor.custom_method()) # N: Revealed type is "builtins.int" + installed_apps: + - myapp + monkeypatch: true + files: + - path: myapp/__init__.py + - path: myapp/models/__init__.py + content: | + from typing import ClassVar + from django.db import models + from django.db.models.fields.related_descriptors import ReverseManyToOneDescriptor + + class Other(models.Model): + explicit_descriptor: ClassVar[ReverseManyToOneDescriptor["MyModel"]] + + class CustomManager(models.Manager["MyModel"]): + def custom_method(self) -> int: ... + + class MyModel(models.Model): + rel = models.ForeignKey( + Other, on_delete=models.CASCADE, related_name="explicit_descriptor" + ) + + objects = CustomManager() + - case: test_reverse_one_to_one_descriptor main: | from myapp.models import MyModel, Other