Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve ManyToManyDescriptor and fix Model.<manytomany>.through typing #1805

Merged
merged 9 commits into from
Nov 27, 2023
2 changes: 1 addition & 1 deletion django-stubs/db/models/fields/related.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -282,7 +282,7 @@ class ManyToManyField(RelatedField[Any, Any], Generic[_To, _M]):
) -> None: ...
# class access
@overload
def __get__(self, instance: None, owner: Any) -> ManyToManyDescriptor[_M]: ...
def __get__(self, instance: None, owner: Any) -> ManyToManyDescriptor[_To, _M]: ...
intgr marked this conversation as resolved.
Show resolved Hide resolved
# Model instance access
@overload
def __get__(self, instance: Model, owner: Any) -> ManyRelatedManager[_To]: ...
Expand Down
15 changes: 10 additions & 5 deletions django-stubs/db/models/fields/related_descriptors.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ from typing_extensions import Self
_M = TypeVar("_M", bound=Model)
_F = TypeVar("_F", bound=Field)
_From = TypeVar("_From", bound=Model)
_Through = TypeVar("_Through", bound=Model)
_To = TypeVar("_To", bound=Model)

class ForeignKeyDeferredAttribute(DeferredAttribute):
Expand Down Expand Up @@ -84,7 +85,7 @@ class ReverseManyToOneDescriptor:
@overload
def __get__(self, instance: None, cls: Any = ...) -> Self: ...
@overload
def __get__(self, instance: Model, cls: Any = ...) -> type[RelatedManager[Any]]: ...
def __get__(self, instance: Model, cls: Any = ...) -> RelatedManager[Any]: ...
def __set__(self, instance: Any, value: Any) -> NoReturn: ...

# Fake class, Django defines 'RelatedManager' inside a function body
Expand All @@ -104,7 +105,7 @@ def create_reverse_many_to_one_manager(
superclass: type[BaseManager[_M]], rel: ManyToOneRel
) -> type[RelatedManager[_M]]: ...

class ManyToManyDescriptor(ReverseManyToOneDescriptor, Generic[_M]):
class ManyToManyDescriptor(ReverseManyToOneDescriptor, Generic[_To, _Through]):
"""
In the example::

Expand All @@ -117,13 +118,17 @@ class ManyToManyDescriptor(ReverseManyToOneDescriptor, Generic[_M]):

# 'field' here is 'rel.field'
rel: ManyToManyRel # type: ignore[assignment]
field: ManyToManyField[Any, _M] # type: ignore[assignment]
field: ManyToManyField[_To, _Through] # type: ignore[assignment]
reverse: bool
def __init__(self, rel: ManyToManyRel, reverse: bool = ...) -> None: ...
@property
def through(self) -> type[_M]: ...
def through(self) -> type[_Through]: ...
@cached_property
def related_manager_cls(self) -> type[ManyRelatedManager[Any]]: ... # type: ignore[override]
def related_manager_cls(self) -> type[ManyRelatedManager[_To]]: ... # type: ignore[override]
@overload # type: ignore[override]
def __get__(self, instance: None, cls: Any = ...) -> Self: ...
@overload
def __get__(self, instance: Model, cls: Any = ...) -> ManyRelatedManager[_To]: ...
Comment on lines +130 to +131
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just realised that the plugin needs to make the returned ManyRelatedManager here more accurate, so that it's possible to find custom queryset/manager methods.

Quite surprised we're missing tests for that actually.

It probably would've been easiest if ManyRelatedManager was another type argument, but I don't think it can. Since it's not possible to pass any dynamically generated managers as type arguments (ref #1699).

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've adjusted so that we imitate Django's behaviour by having the plugin create a manager class with ManyRelatedManager and the _To model's default manager as bases.

Ref: https://github.com/django/django/blob/46df3ab244e1688bd186f0bfbfea6a354097a910/django/db/models/fields/related_descriptors.py#L1032


# Fake class, Django defines 'ManyRelatedManager' inside a function body
class ManyRelatedManager(BaseManager[_M], Generic[_M]):
Expand Down
2 changes: 2 additions & 0 deletions mypy_django_plugin/lib/fullnames.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@
}

REVERSE_ONE_TO_ONE_DESCRIPTOR = "django.db.models.fields.related_descriptors.ReverseOneToOneDescriptor"
MANY_TO_MANY_DESCRIPTOR = "django.db.models.fields.related_descriptors.ManyToManyDescriptor"
MANY_RELATED_MANAGER = "django.db.models.fields.related_descriptors.ManyRelatedManager"
RELATED_FIELDS_CLASSES = frozenset(
(
FOREIGN_OBJECT_FULLNAME,
Expand Down
14 changes: 14 additions & 0 deletions mypy_django_plugin/lib/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,20 @@ def get_django_metadata_bases(
return get_django_metadata(model_info).setdefault(key, cast(Dict[str, int], {}))


def get_reverse_manager_info(
api: Union[TypeChecker, SemanticAnalyzer], model_info: TypeInfo, derived_from: str
) -> Optional[TypeInfo]:
manager_fullname = get_django_metadata(model_info).get("reverse_managers", {}).get(derived_from)
if not manager_fullname:
return None

return lookup_fully_qualified_typeinfo(api, manager_fullname)


def set_reverse_manager_info(model_info: TypeInfo, derived_from: str, fullname: str) -> None:
get_django_metadata(model_info).setdefault("reverse_managers", {})[derived_from] = fullname


class IncompleteDefnException(Exception):
pass

Expand Down
7 changes: 6 additions & 1 deletion mypy_django_plugin/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 fields, forms, init_create, manytomany, 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,
Expand Down Expand Up @@ -187,6 +187,11 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
info = self._get_typeinfo_or_none(class_fullname)
if info and info.has_base(fullnames.FORM_MIXIN_CLASS_FULLNAME):
return forms.extract_proper_type_for_get_form
elif method_name == "__get__" and class_fullname in {
intgr marked this conversation as resolved.
Show resolved Hide resolved
fullnames.MANYTOMANY_FIELD_FULLNAME,
fullnames.MANY_TO_MANY_DESCRIPTOR,
}:
return manytomany.refine_many_to_many_related_manager

manager_classes = self._get_current_manager_bases()

Expand Down
51 changes: 50 additions & 1 deletion mypy_django_plugin/transformers/manytomany.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from mypy.checker import TypeChecker
from mypy.nodes import AssignmentStmt, Expression, MemberExpr, NameExpr, StrExpr, TypeInfo
from mypy.plugin import FunctionContext
from mypy.plugin import FunctionContext, MethodContext
from mypy.semanal import SemanticAnalyzer
from mypy.types import Instance, ProperType, UninhabitedType
from mypy.types import Type as MypyType
Expand Down Expand Up @@ -151,3 +151,52 @@ def get_model_from_expression(
if model_info is not None:
return Instance(model_info, [])
return None


def refine_many_to_many_related_manager(ctx: MethodContext) -> MypyType:
"""
Updates the 'ManyRelatedManager' returned by e.g. 'ManyToManyDescriptor' to be a subclass
of 'ManyRelatedManager' and the related model's default manager.
"""
if (
not isinstance(ctx.default_return_type, Instance)
# Only change a return type being 'ManyRelatedManager'
or ctx.default_return_type.type.fullname != fullnames.MANY_RELATED_MANAGER
):
return ctx.default_return_type

# This is a call to '__get__' overload with a model instance of 'ManyToManyDescriptor'.
# Returning a 'ManyRelatedManager'. Which we want to, just like Django, build from the
# default manager of the related model.
many_related_manager = ctx.default_return_type
# Require first type argument of 'ManyRelatedManager' to be a model
if (
not many_related_manager.args
or not isinstance(many_related_manager.args[0], Instance)
or many_related_manager.args[0].type.metaclass_type is None
or many_related_manager.args[0].type.metaclass_type.type.fullname != fullnames.MODEL_METACLASS_FULLNAME
intgr marked this conversation as resolved.
Show resolved Hide resolved
):
return ctx.default_return_type

checker = helpers.get_typechecker_api(ctx)
related_model_instance = many_related_manager.args[0].copy_modified()
related_manager_info = helpers.get_reverse_manager_info(
checker, related_model_instance.type, derived_from="_default_manager"
)
if related_manager_info is None:
default_manager_node = related_model_instance.type.names.get("_default_manager")
if default_manager_node is None or not isinstance(default_manager_node.type, Instance):
return ctx.default_return_type

related_manager_info = helpers.add_new_class_for_module(
module=checker.modules[related_model_instance.type.module_name],
name=f"{related_model_instance.type.name}_ManyRelatedManager",
bases=[many_related_manager, default_manager_node.type],
)
related_manager_info.metadata["django"] = {"related_manager_to_model": related_model_instance.type.fullname}
helpers.set_reverse_manager_info(
related_model_instance.type,
derived_from="_default_manager",
fullname=related_manager_info.fullname,
)
return Instance(related_manager_info, [])
46 changes: 25 additions & 21 deletions mypy_django_plugin/transformers/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,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 ForeignObjectRel, OneToOneRel
from django.db.models.fields.reverse_related import ManyToManyRel, OneToOneRel
from mypy.checker import TypeChecker
from mypy.nodes import (
ARG_STAR2,
Expand Down Expand Up @@ -448,23 +448,15 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:


class AddReverseLookups(ModelClassInitializer):
def get_reverse_manager_info(self, model_info: TypeInfo, derived_from: str) -> Optional[TypeInfo]:
manager_fullname = helpers.get_django_metadata(model_info).get("reverse_managers", {}).get(derived_from)
if not manager_fullname:
return None

symbol = self.api.lookup_fully_qualified_or_none(manager_fullname)
if symbol is None or not isinstance(symbol.node, TypeInfo):
return None
return symbol.node
@cached_property
def reverse_one_to_one_descriptor(self) -> TypeInfo:
return self.lookup_typeinfo_or_incomplete_defn_error(fullnames.REVERSE_ONE_TO_ONE_DESCRIPTOR)

def set_reverse_manager_info(self, model_info: TypeInfo, derived_from: str, fullname: str) -> None:
helpers.get_django_metadata(model_info).setdefault("reverse_managers", {})[derived_from] = fullname
@cached_property
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:
reverse_one_to_one_descriptor = self.lookup_typeinfo_or_incomplete_defn_error(
fullnames.REVERSE_ONE_TO_ONE_DESCRIPTOR
)
# add related managers
for relation in self.django_context.get_model_relations(model_cls):
attname = relation.get_accessor_name()
Expand All @@ -487,13 +479,25 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
self.add_new_node_to_model_class(
attname,
Instance(
reverse_one_to_one_descriptor,
self.reverse_one_to_one_descriptor,
[Instance(self.model_classdef.info, []), Instance(related_model_info, [])],
),
)
continue

if isinstance(relation, ForeignObjectRel):
elif isinstance(relation, ManyToManyRel):
intgr marked this conversation as resolved.
Show resolved Hide resolved
# TODO: 'relation' should be based on `TypeInfo` instead of Django runtime..
intgr marked this conversation as resolved.
Show resolved Hide resolved
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, [])]
),
)
else:
intgr marked this conversation as resolved.
Show resolved Hide resolved
related_manager_info = None
try:
related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(
Expand Down Expand Up @@ -534,8 +538,8 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:

# 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 = self.get_reverse_manager_info(
model_info=related_model_info, derived_from="_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, []))
Expand Down Expand Up @@ -564,7 +568,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
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
self.set_reverse_manager_info(
helpers.set_reverse_manager_info(
related_model_info,
derived_from="_default_manager",
fullname=new_related_manager_info.fullname,
Expand Down
4 changes: 1 addition & 3 deletions mypy_django_plugin/transformers/orm_lookups.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,7 @@ def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext)
lookup_kwargs = ctx.arg_names[1] if len(ctx.arg_names) >= 2 else []
provided_lookup_types = ctx.arg_types[1] if len(ctx.arg_types) >= 2 else []

assert isinstance(ctx.type, Instance)

if not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
if not isinstance(ctx.type, Instance) or not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
return ctx.default_return_type

model_cls_fullname = ctx.type.args[0].type.fullname
Expand Down
2 changes: 0 additions & 2 deletions scripts/stubtest/allowlist.txt
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,7 @@ django.db.models.fields.related_descriptors.RelatedManager
# _locally/dynamically_ runtime -- Created via
# 'django.db.models.fields.related_descriptors.create_reverse_many_to_one_manager'
django.contrib.admin.models.LogEntry_RelatedManager
django.contrib.auth.models.Group_RelatedManager
django.contrib.auth.models.Permission_RelatedManager
django.contrib.auth.models.User_RelatedManager

# BaseArchive abstract methods that take no argument, but typed with arguments to match the Archive and TarArchive Implementations
django.utils.archive.BaseArchive.list
Expand Down
1 change: 0 additions & 1 deletion scripts/stubtest/allowlist_todo.txt
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,6 @@ django.contrib.auth.models.GroupManager.__slotnames__
django.contrib.auth.models.Permission.codename
django.contrib.auth.models.Permission.content_type
django.contrib.auth.models.Permission.content_type_id
django.contrib.auth.models.Permission.group_set
django.contrib.auth.models.Permission.id
django.contrib.auth.models.Permission.name
django.contrib.auth.models.Permission.user_set
Expand Down
Loading