From b651b1c48d0652d547a852895074e47bda7b4107 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Sat, 22 Jun 2024 21:29:04 +0200 Subject: [PATCH 1/2] Do many related manager creation during semantic analysis When finding a `ManyToManyField` we can directly create 2 `ManyRelatedManager`s, one for each side of the relation. --- mypy_django_plugin/transformers/manytomany.py | 39 +---------- mypy_django_plugin/transformers/models.py | 64 +++++++++++++++++-- 2 files changed, 62 insertions(+), 41 deletions(-) diff --git a/mypy_django_plugin/transformers/manytomany.py b/mypy_django_plugin/transformers/manytomany.py index 4e0617dd3..c9fa3fc75 100644 --- a/mypy_django_plugin/transformers/manytomany.py +++ b/mypy_django_plugin/transformers/manytomany.py @@ -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 @@ -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]) diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 5c17b930a..91fa28fa4 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -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 @@ -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]: @@ -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: @@ -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, @@ -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 From b677bd3fafc47f75444966d64cd8cce3762cccea Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Sat, 22 Jun 2024 21:45:58 +0200 Subject: [PATCH 2/2] fixup! Do many related manager creation during semantic analysis --- scripts/stubtest/allowlist.txt | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/scripts/stubtest/allowlist.txt b/scripts/stubtest/allowlist.txt index 06bdb091d..aeccccbd7 100644 --- a/scripts/stubtest/allowlist.txt +++ b/scripts/stubtest/allowlist.txt @@ -27,6 +27,13 @@ django.core.files.storage.default_storage django.contrib.admin.models.LogEntry_RelatedManager django.contrib.auth.models.Permission_RelatedManager +# '_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 + # 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