Skip to content

Commit c693e2a

Browse files
authored
Make ReverseManyToOneDescriptor generic over a model (#2227)
This allows us to remove some plugin code, as we now can fall back to the default related manager on the descriptor instead of having to implement the fallback behaviour manually.
1 parent 23530b9 commit c693e2a

File tree

10 files changed

+155
-103
lines changed

10 files changed

+155
-103
lines changed

django-stubs/db/models/fields/related_descriptors.pyi

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ class ReverseOneToOneDescriptor(Generic[_From, _To]):
7373
def __set__(self, instance: _From, value: _To | None) -> None: ...
7474
def __reduce__(self) -> tuple[Callable[..., Any], tuple[type[_To], str]]: ...
7575

76-
class ReverseManyToOneDescriptor:
76+
class ReverseManyToOneDescriptor(Generic[_To]):
7777
"""
7878
In the example::
7979
@@ -84,29 +84,29 @@ class ReverseManyToOneDescriptor:
8484
"""
8585

8686
rel: ManyToOneRel
87-
field: ForeignKey
87+
field: ForeignKey[_To, _To]
8888
def __init__(self, rel: ManyToOneRel) -> None: ...
8989
@cached_property
90-
def related_manager_cls(self) -> type[RelatedManager[Any]]: ...
90+
def related_manager_cls(self) -> type[RelatedManager[_To]]: ...
9191
@overload
9292
def __get__(self, instance: None, cls: Any = ...) -> Self: ...
9393
@overload
94-
def __get__(self, instance: Model, cls: Any = ...) -> RelatedManager[Any]: ...
94+
def __get__(self, instance: Model, cls: Any = ...) -> RelatedManager[_To]: ...
9595
def __set__(self, instance: Any, value: Any) -> NoReturn: ...
9696

9797
# Fake class, Django defines 'RelatedManager' inside a function body
9898
@type_check_only
99-
class RelatedManager(Manager[_M], Generic[_M]):
99+
class RelatedManager(Manager[_To], Generic[_To]):
100100
related_val: tuple[int, ...]
101-
def add(self, *objs: _M | int, bulk: bool = ...) -> None: ...
102-
async def aadd(self, *objs: _M | int, bulk: bool = ...) -> None: ...
103-
def remove(self, *objs: _M | int, bulk: bool = ...) -> None: ...
104-
async def aremove(self, *objs: _M | int, bulk: bool = ...) -> None: ...
105-
def set(self, objs: QuerySet[_M] | Iterable[_M | int], *, bulk: bool = ..., clear: bool = ...) -> None: ...
106-
async def aset(self, objs: QuerySet[_M] | Iterable[_M | int], *, bulk: bool = ..., clear: bool = ...) -> None: ...
101+
def add(self, *objs: _To | int, bulk: bool = ...) -> None: ...
102+
async def aadd(self, *objs: _To | int, bulk: bool = ...) -> None: ...
103+
def remove(self, *objs: _To | int, bulk: bool = ...) -> None: ...
104+
async def aremove(self, *objs: _To | int, bulk: bool = ...) -> None: ...
105+
def set(self, objs: QuerySet[_To] | Iterable[_To | int], *, bulk: bool = ..., clear: bool = ...) -> None: ...
106+
async def aset(self, objs: QuerySet[_To] | Iterable[_To | int], *, bulk: bool = ..., clear: bool = ...) -> None: ...
107107
def clear(self) -> None: ...
108108
async def aclear(self) -> None: ...
109-
def __call__(self, *, manager: str) -> RelatedManager[_M]: ...
109+
def __call__(self, *, manager: str) -> RelatedManager[_To]: ...
110110

111111
def create_reverse_many_to_one_manager(
112112
superclass: type[BaseManager[_M]], rel: ManyToOneRel

ext/django_stubs_ext/patch.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from django.db.models.expressions import ExpressionWrapper
1313
from django.db.models.fields import Field
1414
from django.db.models.fields.related import ForeignKey
15+
from django.db.models.fields.related_descriptors import ReverseManyToOneDescriptor
1516
from django.db.models.lookups import Lookup
1617
from django.db.models.manager import BaseManager
1718
from django.db.models.query import QuerySet
@@ -71,6 +72,7 @@ def __repr__(self) -> str:
7172
MPGeneric(Lookup),
7273
MPGeneric(BaseConnectionHandler),
7374
MPGeneric(ExpressionWrapper),
75+
MPGeneric(ReverseManyToOneDescriptor),
7476
# These types do have native `__class_getitem__` method since django 3.1:
7577
MPGeneric(QuerySet, (3, 1)),
7678
MPGeneric(BaseManager, (3, 1)),

mypy_django_plugin/lib/fullnames.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
}
3535

3636
REVERSE_ONE_TO_ONE_DESCRIPTOR = "django.db.models.fields.related_descriptors.ReverseOneToOneDescriptor"
37+
REVERSE_MANY_TO_ONE_DESCRIPTOR = "django.db.models.fields.related_descriptors.ReverseManyToOneDescriptor"
3738
MANY_TO_MANY_DESCRIPTOR = "django.db.models.fields.related_descriptors.ManyToManyDescriptor"
3839
MANY_RELATED_MANAGER = "django.db.models.fields.related_descriptors.ManyRelatedManager"
3940
RELATED_FIELDS_CLASSES = frozenset(

mypy_django_plugin/main.py

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,17 @@
2424
from mypy_django_plugin.django.context import DjangoContext
2525
from mypy_django_plugin.exceptions import UnregisteredModelError
2626
from mypy_django_plugin.lib import fullnames, helpers
27-
from mypy_django_plugin.transformers import fields, forms, init_create, manytomany, meta, querysets, request, settings
27+
from mypy_django_plugin.transformers import (
28+
fields,
29+
forms,
30+
init_create,
31+
manytomany,
32+
manytoone,
33+
meta,
34+
querysets,
35+
request,
36+
settings,
37+
)
2838
from mypy_django_plugin.transformers.functional import resolve_str_promise_attribute
2939
from mypy_django_plugin.transformers.managers import (
3040
create_new_manager_class_from_as_manager_method,
@@ -190,11 +200,13 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
190200
if info and info.has_base(fullnames.FORM_MIXIN_CLASS_FULLNAME):
191201
return forms.extract_proper_type_for_get_form
192202

193-
elif method_name == "__get__" and class_fullname in {
194-
fullnames.MANYTOMANY_FIELD_FULLNAME,
195-
fullnames.MANY_TO_MANY_DESCRIPTOR,
196-
}:
197-
return manytomany.refine_many_to_many_related_manager
203+
elif method_name == "__get__":
204+
hooks = {
205+
fullnames.MANYTOMANY_FIELD_FULLNAME: manytomany.refine_many_to_many_related_manager,
206+
fullnames.MANY_TO_MANY_DESCRIPTOR: manytomany.refine_many_to_many_related_manager,
207+
fullnames.REVERSE_MANY_TO_ONE_DESCRIPTOR: manytoone.refine_many_to_one_related_manager,
208+
}
209+
return hooks.get(class_fullname)
198210

199211
manager_classes = self._get_current_manager_bases()
200212

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
from typing import Optional
2+
3+
from mypy.plugin import MethodContext
4+
from mypy.types import Instance
5+
from mypy.types import Type as MypyType
6+
7+
from mypy_django_plugin.lib import fullnames, helpers
8+
9+
10+
def get_model_of_related_manager(ctx: MethodContext) -> Optional[Instance]:
11+
"""
12+
Returns the type parameter (_To) instance of a `RelatedManager` instance when it's a
13+
model. Otherwise `None` is returned.
14+
15+
For example: if given a `RelatedManager[A]` where `A` is a model then `A` is
16+
returned.
17+
"""
18+
if (
19+
isinstance(ctx.default_return_type, Instance)
20+
and ctx.default_return_type.type.fullname == fullnames.RELATED_MANAGER_CLASS
21+
):
22+
# This is a call to '__get__' overload with a model instance of
23+
# 'ReverseManyToOneDescriptor'. Returning a 'RelatedManager'. Which we want to,
24+
# just like Django, build from the default manager of the related model.
25+
related_manager = ctx.default_return_type
26+
# Require first type argument of 'RelatedManager' to be a model
27+
if (
28+
len(related_manager.args) >= 1
29+
and isinstance(related_manager.args[0], Instance)
30+
and helpers.is_model_type(related_manager.args[0].type)
31+
):
32+
return related_manager.args[0]
33+
34+
return None
35+
36+
37+
def refine_many_to_one_related_manager(ctx: MethodContext) -> MypyType:
38+
"""
39+
Updates the 'RelatedManager' returned by e.g. 'ReverseManyToOneDescriptor' to be a subclass
40+
of 'RelatedManager' and the related model's default manager.
41+
"""
42+
to_model_instance = get_model_of_related_manager(ctx)
43+
if to_model_instance is None:
44+
return ctx.default_return_type
45+
46+
checker = helpers.get_typechecker_api(ctx)
47+
related_manager_info = helpers.get_reverse_manager_info(
48+
checker, to_model_instance.type, derived_from="_default_manager"
49+
)
50+
if related_manager_info is None:
51+
return ctx.default_return_type
52+
return Instance(related_manager_info, [])

mypy_django_plugin/transformers/models.py

Lines changed: 55 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ def get_generated_manager_info(self, manager_fullname: str, base_manager_fullnam
114114
# Not a generated manager
115115
return None
116116

117-
def get_or_create_manager_with_any_fallback(self, related_manager: bool = False) -> Optional[TypeInfo]:
117+
def get_or_create_manager_with_any_fallback(self) -> Optional[TypeInfo]:
118118
"""
119119
Create a Manager subclass with fallback to Any for unknown attributes
120120
and methods. This is used for unresolved managers, where we don't know
@@ -123,7 +123,7 @@ def get_or_create_manager_with_any_fallback(self, related_manager: bool = False)
123123
The created class is reused if multiple unknown managers are encountered.
124124
"""
125125

126-
name = "UnknownManager" if not related_manager else "UnknownRelatedManager"
126+
name = "UnknownManager"
127127

128128
# Check if we've already created a fallback manager class for this
129129
# module, and if so reuse that.
@@ -134,9 +134,7 @@ def get_or_create_manager_with_any_fallback(self, related_manager: bool = False)
134134
fallback_queryset = self.get_or_create_queryset_with_any_fallback()
135135
if fallback_queryset is None:
136136
return None
137-
base_manager_fullname = (
138-
fullnames.MANAGER_CLASS_FULLNAME if not related_manager else fullnames.RELATED_MANAGER_CLASS
139-
)
137+
base_manager_fullname = fullnames.MANAGER_CLASS_FULLNAME
140138
base_manager_info = self.lookup_typeinfo(base_manager_fullname)
141139
if base_manager_info is None:
142140
return None
@@ -455,6 +453,10 @@ class AddReverseLookups(ModelClassInitializer):
455453
def reverse_one_to_one_descriptor(self) -> TypeInfo:
456454
return self.lookup_typeinfo_or_incomplete_defn_error(fullnames.REVERSE_ONE_TO_ONE_DESCRIPTOR)
457455

456+
@cached_property
457+
def reverse_many_to_one_descriptor(self) -> TypeInfo:
458+
return self.lookup_typeinfo_or_incomplete_defn_error(fullnames.REVERSE_MANY_TO_ONE_DESCRIPTOR)
459+
458460
@cached_property
459461
def many_to_many_descriptor(self) -> TypeInfo:
460462
return self.lookup_typeinfo_or_incomplete_defn_error(fullnames.MANY_TO_MANY_DESCRIPTOR)
@@ -466,23 +468,21 @@ def process_relation(self, relation: ForeignObjectRel) -> None:
466468
# explicitly declared(i.e. non-inferred) reverse accessors alone
467469
return
468470

469-
related_model_cls = self.django_context.get_field_related_model_cls(relation)
470-
related_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error(related_model_cls)
471+
to_model_cls = self.django_context.get_field_related_model_cls(relation)
472+
to_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error(to_model_cls)
471473

472474
if isinstance(relation, OneToOneRel):
473475
self.add_new_node_to_model_class(
474476
attname,
475477
Instance(
476478
self.reverse_one_to_one_descriptor,
477-
[Instance(self.model_classdef.info, []), Instance(related_model_info, [])],
479+
[Instance(self.model_classdef.info, []), Instance(to_model_info, [])],
478480
),
479481
)
480482
return
481483

482484
elif isinstance(relation, ManyToManyRel):
483485
# TODO: 'relation' should be based on `TypeInfo` instead of Django runtime.
484-
to_fullname = helpers.get_class_fullname(relation.remote_field.model)
485-
to_model_info = self.lookup_typeinfo_or_incomplete_defn_error(to_fullname)
486486
assert relation.through is not None
487487
through_fullname = helpers.get_class_fullname(relation.through)
488488
through_model_info = self.lookup_typeinfo_or_incomplete_defn_error(through_fullname)
@@ -492,79 +492,65 @@ def process_relation(self, relation: ForeignObjectRel) -> None:
492492
)
493493
return
494494

495-
related_manager_info = None
496-
try:
497-
related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS)
498-
default_manager = related_model_info.names.get("_default_manager")
499-
if not default_manager:
500-
raise helpers.IncompleteDefnException()
501-
except helpers.IncompleteDefnException as exc:
502-
if not self.api.final_iteration:
503-
raise exc
504-
505-
# If a django model has a Manager class that cannot be
506-
# resolved statically (if it is generated in a way where we
507-
# cannot import it, like `objects = my_manager_factory()`),
508-
# we fallback to the default related manager, so you at
509-
# least get a base level of working type checking.
510-
#
511-
# See https://github.com/typeddjango/django-stubs/pull/993
512-
# for more information on when this error can occur.
513-
fallback_manager = self.get_or_create_manager_with_any_fallback(related_manager=True)
514-
if fallback_manager is not None:
515-
self.add_new_node_to_model_class(
516-
attname, Instance(fallback_manager, [Instance(related_model_info, [])])
517-
)
518-
related_model_fullname = related_model_cls.__module__ + "." + related_model_cls.__name__
519-
self.ctx.api.fail(
520-
(
521-
"Couldn't resolve related manager for relation "
522-
f"{relation.name!r} (from {related_model_fullname}."
523-
f"{relation.field})."
524-
),
525-
self.ctx.cls,
526-
code=MANAGER_MISSING,
527-
)
528-
return
529-
530-
# Check if the related model has a related manager subclassed from the default manager
495+
related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS)
531496
# TODO: Support other reverse managers than `_default_manager`
532-
default_reverse_manager_info = helpers.get_reverse_manager_info(
533-
self.api, model_info=related_model_info, derived_from="_default_manager"
534-
)
535-
if default_reverse_manager_info:
536-
self.add_new_node_to_model_class(attname, Instance(default_reverse_manager_info, []))
537-
return
497+
default_manager = to_model_info.names.get("_default_manager")
498+
if default_manager is None and not self.api.final_iteration:
499+
raise helpers.IncompleteDefnException()
500+
elif (
501+
# When we get no default manager we can't customize the reverse manager any
502+
# further and will just fall back to the manager declared on the descriptor
503+
default_manager is None
504+
# '_default_manager' attribute is a node type we can't process
505+
or not isinstance(default_manager.type, Instance)
506+
# Already has a related manager subclassed from the default manager
507+
or helpers.get_reverse_manager_info(self.api, model_info=to_model_info, derived_from="_default_manager")
508+
is not None
509+
# When the default manager isn't custom there's no need to create a new type
510+
# as `RelatedManager` has `models.Manager` as base
511+
or default_manager.type.type.fullname == fullnames.MANAGER_CLASS_FULLNAME
512+
):
513+
if default_manager is None and self.api.final_iteration:
514+
# If a django model has a Manager class that cannot be
515+
# resolved statically (if it is generated in a way where we
516+
# cannot import it, like `objects = my_manager_factory()`),
517+
#
518+
# See https://github.com/typeddjango/django-stubs/pull/993
519+
# for more information on when this error can occur.
520+
self.ctx.api.fail(
521+
(
522+
f"Couldn't resolve related manager {attname!r} for relation "
523+
f"'{to_model_info.fullname}.{relation.field.name}'."
524+
),
525+
self.ctx.cls,
526+
code=MANAGER_MISSING,
527+
)
538528

539-
# The reverse manager we're looking for doesn't exist. So we
540-
# create it. The (default) reverse manager type is built from a
541-
# RelatedManager and the default manager on the related model
542-
parametrized_related_manager_type = Instance(related_manager_info, [Instance(related_model_info, [])])
543-
default_manager_type = default_manager.type
544-
assert default_manager_type is not None
545-
assert isinstance(default_manager_type, Instance)
546-
# When the default manager isn't custom there's no need to create a new type
547-
# as `RelatedManager` has `models.Manager` as base
548-
if default_manager_type.type.fullname == fullnames.MANAGER_CLASS_FULLNAME:
549-
self.add_new_node_to_model_class(attname, parametrized_related_manager_type)
529+
self.add_new_node_to_model_class(
530+
attname, Instance(self.reverse_many_to_one_descriptor, [Instance(to_model_info, [])])
531+
)
550532
return
551533

534+
# Create a reverse manager subclassed from the default manager of the related
535+
# model and 'RelatedManager'
536+
related_manager = Instance(related_manager_info, [Instance(to_model_info, [])])
552537
# The reverse manager is based on the related model's manager, so it makes most sense to add the new
553538
# related manager in that module
554539
new_related_manager_info = helpers.add_new_class_for_module(
555-
module=self.api.modules[related_model_info.module_name],
556-
name=f"{related_model_cls.__name__}_RelatedManager",
557-
bases=[parametrized_related_manager_type, default_manager_type],
540+
module=self.api.modules[to_model_info.module_name],
541+
name=f"{to_model_info.name}_RelatedManager",
542+
bases=[related_manager, default_manager.type],
558543
)
559-
new_related_manager_info.metadata["django"] = {"related_manager_to_model": related_model_info.fullname}
560544
# Stash the new reverse manager type fullname on the related model, so we don't duplicate
561545
# or have to create it again for other reverse relations
562546
helpers.set_reverse_manager_info(
563-
related_model_info,
547+
to_model_info,
564548
derived_from="_default_manager",
565549
fullname=new_related_manager_info.fullname,
566550
)
567-
self.add_new_node_to_model_class(attname, Instance(new_related_manager_info, []))
551+
self.add_new_node_to_model_class(
552+
attname, Instance(self.reverse_many_to_one_descriptor, [Instance(to_model_info, [])])
553+
)
568554

569555
def run_with_model_cls(self, model_cls: Type[Model]) -> None:
570556
# add related managers etc.

scripts/stubtest/allowlist_todo.txt

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,6 @@ django.contrib.auth.models.User.is_staff
140140
django.contrib.auth.models.User.is_superuser
141141
django.contrib.auth.models.User.last_login
142142
django.contrib.auth.models.User.last_name
143-
django.contrib.auth.models.User.logentry_set
144143
django.contrib.auth.models.User.password
145144
django.contrib.auth.models.User.user_permissions
146145
django.contrib.auth.models.User.username
@@ -169,9 +168,7 @@ django.contrib.contenttypes.management.inject_rename_contenttypes_operations
169168
django.contrib.contenttypes.models.ContentType.app_label
170169
django.contrib.contenttypes.models.ContentType.app_labeled_name
171170
django.contrib.contenttypes.models.ContentType.id
172-
django.contrib.contenttypes.models.ContentType.logentry_set
173171
django.contrib.contenttypes.models.ContentType.model
174-
django.contrib.contenttypes.models.ContentType.permission_set
175172
django.contrib.contenttypes.models.ContentTypeManager.__init__
176173
django.contrib.contenttypes.models.ContentTypeManager.__slotnames__
177174
django.contrib.flatpages.admin.FlatPageAdmin
@@ -515,7 +512,6 @@ django.contrib.sites.models.Site.domain
515512
django.contrib.sites.models.Site.flatpage_set
516513
django.contrib.sites.models.Site.id
517514
django.contrib.sites.models.Site.name
518-
django.contrib.sites.models.Site.redirect_set
519515
django.contrib.sites.models.SiteManager.__slotnames__
520516
django.contrib.staticfiles.finders.BaseStorageFinder.storage
521517
django.contrib.staticfiles.finders.DefaultStorageFinder.storage

0 commit comments

Comments
 (0)