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

Do many related manager creation during semantic analysis #2231

Merged
merged 2 commits into from
Jun 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 2 additions & 37 deletions mypy_django_plugin/transformers/manytomany.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from mypy.nodes import AssignmentStmt, Expression, MemberExpr, NameExpr, Node, RefExpr, StrExpr, TypeInfo
from mypy.plugin import FunctionContext, MethodContext
from mypy.semanal import SemanticAnalyzer
from mypy.types import Instance, ProperType, TypeVarType, UninhabitedType
from mypy.types import Instance, ProperType, UninhabitedType
from mypy.types import Type as MypyType

from mypy_django_plugin.django.context import DjangoContext
Expand Down Expand Up @@ -198,40 +198,5 @@ def refine_many_to_many_related_manager(ctx: MethodContext) -> MypyType:
checker, to=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

# Create a reusable generic subclass that is generic over a 'through' model,
# explicitly declared it'd could have looked something like below
#
# class X(models.Model): ...
# _Through = TypeVar("_Through", bound=models.Model)
# class X_ManyRelatedManager(ManyRelatedManager[X, _Through], type(X._default_manager), Generic[_Through]): ...
_through_type_var = many_related_manager.type.defn.type_vars[1]
assert isinstance(_through_type_var, TypeVarType)
generic_to_many_related_manager = many_related_manager.copy_modified(
args=[
# Keep the same '_To' as the (parent) `ManyRelatedManager` instance
many_related_manager.args[0],
# But reset the '_Through' `TypeVar` declared for `ManyRelatedManager`
_through_type_var.copy_modified(),
]
)
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=[generic_to_many_related_manager, default_manager_node.type],
)
# Reuse the '_Through' `TypeVar` from `ManyRelatedManager` in our subclass
related_manager_info.defn.type_vars = [_through_type_var.copy_modified()]
related_manager_info.add_type_vars()
related_manager_info.metadata["django"] = {"related_manager_to_model": related_model_instance.type.fullname}
# Track the existence of our manager subclass, by tying it to model it operates on
helpers.set_many_to_many_manager_info(
to=related_model_instance.type,
derived_from="_default_manager",
manager_info=related_manager_info,
)

return ctx.default_return_type
return Instance(related_manager_info, [through_model_instance])
64 changes: 60 additions & 4 deletions mypy_django_plugin/transformers/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
from mypy.plugins import common
from mypy.semanal import SemanticAnalyzer
from mypy.typeanal import TypeAnalyser
from mypy.types import AnyType, Instance, ProperType, TypedDictType, TypeOfAny, TypeType, get_proper_type
from mypy.types import AnyType, Instance, ProperType, TypedDictType, TypeOfAny, TypeType, TypeVarType, get_proper_type
from mypy.types import Type as MypyType
from mypy.typevars import fill_typevars

Expand Down Expand Up @@ -608,9 +608,15 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:

class ProcessManyToManyFields(ModelClassInitializer):
"""
Processes 'ManyToManyField()' fields and generates any implicit through tables that
Django also generates. It won't do anything if the model is abstract or for fields
where an explicit 'through' argument has been passed.
Processes 'ManyToManyField()' fields and;

- Generates any implicit through tables that Django also generates. It won't do
anything if the model is abstract or for fields where an explicit 'through'
argument has been passed.
- Creates related managers for both ends of the many to many relationship

TODO: Move the 'related_name' contribution from 'AddReverseLookups' to here. As it
makes sense to add it when processing ManyToManyField
"""

def statements(self) -> Iterable[Statement]:
Expand Down Expand Up @@ -675,6 +681,11 @@ def run(self) -> None:
model_fullname=f"{self.model_classdef.info.module_name}.{through_model_name}",
m2m_args=args,
)
# Create a 'ManyRelatedManager' class for the processed model
self.create_many_related_manager(Instance(self.model_classdef.info, []))
if isinstance(args.to.model, Instance):
# Create a 'ManyRelatedManager' class for the related model
self.create_many_related_manager(args.to.model)

@cached_property
def default_pk_instance(self) -> Instance:
Expand Down Expand Up @@ -715,6 +726,10 @@ def manager_info(self) -> TypeInfo:
def fk_field_types(self) -> FieldDescriptorTypes:
return get_field_descriptor_types(self.fk_field, is_set_nullable=False, is_get_nullable=False)

@cached_property
def many_related_manager(self) -> TypeInfo:
return self.lookup_typeinfo_or_incomplete_defn_error(fullnames.MANY_RELATED_MANAGER)

def get_pk_instance(self, model: TypeInfo, /) -> Instance:
"""
Get a primary key instance of provided model's type info. If primary key can't be resolved,
Expand Down Expand Up @@ -852,6 +867,47 @@ def resolve_many_to_many_arguments(self, call: CallExpr, /, context: Context) ->

return M2MArguments(to=to, through=through)

def create_many_related_manager(self, model: Instance) -> None:
"""
Creates a generic manager that subclasses both 'ManyRelatedManager' and the
default manager of the given model. These are normally used on both models
involved in a ManyToManyField.

The manager classes are generic over a '_Through' model, meaning that they can
be reused for multiple many to many relations.
"""
if helpers.get_many_to_many_manager_info(self.api, to=model.type, derived_from="_default_manager") is not None:
return

default_manager_node = model.type.names.get("_default_manager")
if default_manager_node is None:
raise helpers.IncompleteDefnException()
elif not isinstance(default_manager_node.type, Instance):
return

# Create a reusable generic subclass that is generic over a 'through' model,
# explicitly declared it'd could have looked something like below
#
# class X(models.Model): ...
# _Through = TypeVar("_Through", bound=models.Model)
# class X_ManyRelatedManager(ManyRelatedManager[X, _Through], type(X._default_manager), Generic[_Through]): ...
through_type_var = self.many_related_manager.defn.type_vars[1]
assert isinstance(through_type_var, TypeVarType)
generic_to_many_related_manager = Instance(self.many_related_manager, [model, through_type_var.copy_modified()])
related_manager_info = helpers.add_new_class_for_module(
module=self.api.modules[model.type.module_name],
name=f"{model.type.name}_ManyRelatedManager",
bases=[generic_to_many_related_manager, default_manager_node.type],
)
# Reuse the '_Through' `TypeVar` from `ManyRelatedManager` in our subclass
related_manager_info.defn.type_vars = [through_type_var.copy_modified()]
related_manager_info.add_type_vars()
# Track the existence of our manager subclass, by tying it to the model it
# operates on
helpers.set_many_to_many_manager_info(
to=model.type, derived_from="_default_manager", manager_info=related_manager_info
)


class MetaclassAdjustments(ModelClassInitializer):
@classmethod
Expand Down
7 changes: 7 additions & 0 deletions scripts/stubtest/allowlist.txt
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,13 @@ django.core.files.storage.default_storage
django.contrib.admin.models.LogEntry_RelatedManager
django.contrib.auth.models.Permission_RelatedManager

# '<Model>_ManyRelatedManager' entries are plugin generated and these subclasses only exist
# _locally/dynamically_ runtime -- Created via
# 'django.db.models.fields.related_descriptors.create_forward_many_to_many_manager'
django.contrib.auth.models.Group_ManyRelatedManager
django.contrib.auth.models.Permission_ManyRelatedManager
django.contrib.auth.models.User_ManyRelatedManager

Comment on lines +30 to +36
Copy link
Member Author

Choose a reason for hiding this comment

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

These are only here since the stubs includes = ManyToManyField(...) in the following places:

class Group(models.Model):
objects: ClassVar[GroupManager]
name = models.CharField(max_length=150)
permissions = models.ManyToManyField(Permission)

class PermissionsMixin(models.Model):
is_superuser = models.BooleanField()
groups = models.ManyToManyField(Group)
user_permissions = models.ManyToManyField(Permission)

Once they would move to a type hint form like: field: ManyToManyField[..., ...] they will disappear from the allowlist.

# BaseArchive abstract methods that take no argument, but typed with arguments to match the Archive and TarArchive Implementations
django.utils.archive.BaseArchive.list
django.utils.archive.BaseArchive.extract
Expand Down
Loading