Skip to content

Commit

Permalink
Support processing of other relations and fields when one is broken (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
flaeppe authored Feb 15, 2024
1 parent 0b8b124 commit 6d07585
Show file tree
Hide file tree
Showing 2 changed files with 127 additions and 110 deletions.
218 changes: 108 additions & 110 deletions mypy_django_plugin/transformers/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from django.db.models import Manager, Model
from django.db.models.fields import DateField, DateTimeField, Field
from django.db.models.fields.reverse_related import ManyToManyRel, OneToOneRel
from django.db.models.fields.reverse_related import ForeignObjectRel, ManyToManyRel, OneToOneRel
from mypy.checker import TypeChecker
from mypy.nodes import (
ARG_STAR2,
Expand Down Expand Up @@ -458,126 +458,124 @@ def reverse_one_to_one_descriptor(self) -> TypeInfo:
def many_to_many_descriptor(self) -> TypeInfo:
return self.lookup_typeinfo_or_incomplete_defn_error(fullnames.MANY_TO_MANY_DESCRIPTOR)

def run_with_model_cls(self, model_cls: Type[Model]) -> None:
# add related managers
for relation in self.django_context.get_model_relations(model_cls):
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
continue
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
return

related_model_cls = self.django_context.get_field_related_model_cls(relation)
related_model_cls = self.django_context.get_field_related_model_cls(relation)
related_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error(related_model_cls)

try:
related_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error(related_model_cls)
except helpers.IncompleteDefnException as exc:
if not self.api.final_iteration:
raise exc
else:
continue
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(related_model_info, [])],
),
)
return

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(related_model_info, [])],
),
)
continue
elif isinstance(relation, ManyToManyRel):
# TODO: 'relation' should be based on `TypeInfo` instead of Django runtime.
to_fullname = helpers.get_class_fullname(relation.remote_field.model)
to_model_info = self.lookup_typeinfo_or_incomplete_defn_error(to_fullname)
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, [])]),
)
return

elif isinstance(relation, ManyToManyRel):
# TODO: 'relation' should be based on `TypeInfo` instead of Django runtime.
to_fullname = helpers.get_class_fullname(relation.remote_field.model)
to_model_info = self.lookup_typeinfo_or_incomplete_defn_error(to_fullname)
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)
related_manager_info = None
try:
related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS)
default_manager = related_model_info.names.get("_default_manager")
if not default_manager:
raise helpers.IncompleteDefnException()
except helpers.IncompleteDefnException as exc:
if not self.api.final_iteration:
raise exc

# If a django model has a Manager class that cannot be
# resolved statically (if it is generated in a way where we
# cannot import it, like `objects = my_manager_factory()`),
# we fallback to the default related manager, so you at
# least get a base level of working type checking.
#
# See https://github.com/typeddjango/django-stubs/pull/993
# for more information on when this error can occur.
fallback_manager = self.get_or_create_manager_with_any_fallback(related_manager=True)
if fallback_manager is not None:
self.add_new_node_to_model_class(
attname,
Instance(
self.many_to_many_descriptor, [Instance(to_model_info, []), Instance(through_model_info, [])]
),
attname, Instance(fallback_manager, [Instance(related_model_info, [])])
)
related_model_fullname = related_model_cls.__module__ + "." + related_model_cls.__name__
self.ctx.api.fail(
(
"Couldn't resolve related manager for relation "
f"{relation.name!r} (from {related_model_fullname}."
f"{relation.field})."
),
self.ctx.cls,
code=MANAGER_MISSING,
)
return

else:
related_manager_info = None
try:
related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(
fullnames.RELATED_MANAGER_CLASS
)
default_manager = related_model_info.names.get("_default_manager")
if not default_manager:
raise helpers.IncompleteDefnException()
except helpers.IncompleteDefnException as exc:
if not self.api.final_iteration:
raise exc

# If a django model has a Manager class that cannot be
# resolved statically (if it is generated in a way where we
# cannot import it, like `objects = my_manager_factory()`),
# we fallback to the default related manager, so you at
# least get a base level of working type checking.
#
# See https://github.com/typeddjango/django-stubs/pull/993
# for more information on when this error can occur.
fallback_manager = self.get_or_create_manager_with_any_fallback(related_manager=True)
if fallback_manager is not None:
self.add_new_node_to_model_class(
attname, Instance(fallback_manager, [Instance(related_model_info, [])])
)
related_model_fullname = related_model_cls.__module__ + "." + related_model_cls.__name__
self.ctx.api.fail(
(
"Couldn't resolve related manager for relation "
f"{relation.name!r} (from {related_model_fullname}."
f"{relation.field})."
),
self.ctx.cls,
code=MANAGER_MISSING,
)
# Check if the related model has a related manager subclassed from the default manager
# TODO: Support other reverse managers than `_default_manager`
default_reverse_manager_info = helpers.get_reverse_manager_info(
self.api, model_info=related_model_info, derived_from="_default_manager"
)
if default_reverse_manager_info:
self.add_new_node_to_model_class(attname, Instance(default_reverse_manager_info, []))
return

continue
# The reverse manager we're looking for doesn't exist. So we
# create it. The (default) reverse manager type is built from a
# RelatedManager and the default manager on the related model
parametrized_related_manager_type = Instance(related_manager_info, [Instance(related_model_info, [])])
default_manager_type = default_manager.type
assert default_manager_type is not None
assert isinstance(default_manager_type, Instance)
# When the default manager isn't custom there's no need to create a new type
# as `RelatedManager` has `models.Manager` as base
if default_manager_type.type.fullname == fullnames.MANAGER_CLASS_FULLNAME:
self.add_new_node_to_model_class(attname, parametrized_related_manager_type)
return

# Check if the related model has a related manager subclassed from the default manager
# TODO: Support other reverse managers than `_default_manager`
default_reverse_manager_info = helpers.get_reverse_manager_info(
self.api, model_info=related_model_info, derived_from="_default_manager"
)
if default_reverse_manager_info:
self.add_new_node_to_model_class(attname, Instance(default_reverse_manager_info, []))
continue
# The reverse manager is based on the related model's manager, so it makes most sense to add the new
# related manager in that module
new_related_manager_info = helpers.add_new_class_for_module(
module=self.api.modules[related_model_info.module_name],
name=f"{related_model_cls.__name__}_RelatedManager",
bases=[parametrized_related_manager_type, default_manager_type],
)
new_related_manager_info.metadata["django"] = {"related_manager_to_model": related_model_info.fullname}
# Stash the new reverse manager type fullname on the related model, so we don't duplicate
# or have to create it again for other reverse relations
helpers.set_reverse_manager_info(
related_model_info,
derived_from="_default_manager",
fullname=new_related_manager_info.fullname,
)
self.add_new_node_to_model_class(attname, Instance(new_related_manager_info, []))

# The reverse manager we're looking for doesn't exist. So we
# create it. The (default) reverse manager type is built from a
# RelatedManager and the default manager on the related model
parametrized_related_manager_type = Instance(related_manager_info, [Instance(related_model_info, [])])
default_manager_type = default_manager.type
assert default_manager_type is not None
assert isinstance(default_manager_type, Instance)
# When the default manager isn't custom there's no need to create a new type
# as `RelatedManager` has `models.Manager` as base
if default_manager_type.type.fullname == fullnames.MANAGER_CLASS_FULLNAME:
self.add_new_node_to_model_class(attname, parametrized_related_manager_type)
continue
def run_with_model_cls(self, model_cls: Type[Model]) -> None:
# add related managers etc.
processing_incomplete = False
for relation in self.django_context.get_model_relations(model_cls):
try:
self.process_relation(relation)
except helpers.IncompleteDefnException:
processing_incomplete = True

# The reverse manager is based on the related model's manager, so it makes most sense to add the new
# related manager in that module
new_related_manager_info = helpers.add_new_class_for_module(
module=self.api.modules[related_model_info.module_name],
name=f"{related_model_cls.__name__}_RelatedManager",
bases=[parametrized_related_manager_type, default_manager_type],
)
new_related_manager_info.metadata["django"] = {"related_manager_to_model": related_model_info.fullname}
# Stash the new reverse manager type fullname on the related model, so we don't duplicate
# or have to create it again for other reverse relations
helpers.set_reverse_manager_info(
related_model_info,
derived_from="_default_manager",
fullname=new_related_manager_info.fullname,
)
self.add_new_node_to_model_class(attname, Instance(new_related_manager_info, []))
if processing_incomplete and not self.api.final_iteration:
raise helpers.IncompleteDefnException


class AddExtraFieldMethods(ModelClassInitializer):
Expand Down
19 changes: 19 additions & 0 deletions tests/typecheck/models/test_related_fields.yml
Original file line number Diff line number Diff line change
Expand Up @@ -123,3 +123,22 @@
on_delete=models.CASCADE,
related_name="model_2s",
)
- case: test_processes_other_relations_when_one_field_is_broken
main: |
from myapp.models import MyModel
reveal_type(MyModel().others) # N: Revealed type is "django.db.models.fields.related_descriptors.RelatedManager[myapp.models.Other]"
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class MyModel(models.Model):
# Plugin doesn't handle lazy reference without app label
broken = models.ManyToManyField("MyModel", related_name="mymodels") # type: ignore[var-annotated]
class Other(models.Model):
field = models.ForeignKey(MyModel, related_name="others", on_delete=models.CASCADE)

0 comments on commit 6d07585

Please sign in to comment.