Skip to content

Commit

Permalink
Make plugin handle explicitly declared reverse descriptors for FKs (#…
Browse files Browse the repository at this point in the history
…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
  • Loading branch information
flaeppe authored Jun 28, 2024
1 parent 590f033 commit 8703696
Show file tree
Hide file tree
Showing 3 changed files with 75 additions and 25 deletions.
52 changes: 27 additions & 25 deletions mypy_django_plugin/transformers/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down
18 changes: 18 additions & 0 deletions tests/assert_type/db/models/fields/test_related_descriptors.py
Original file line number Diff line number Diff line change
@@ -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]
30 changes: 30 additions & 0 deletions tests/typecheck/fields/test_related.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 8703696

Please sign in to comment.