diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index 554550378..e5f3ac837 100644 --- a/mypy_django_plugin/main.py +++ b/mypy_django_plugin/main.py @@ -37,8 +37,7 @@ ) from mypy_django_plugin.transformers.functional import resolve_str_promise_attribute from mypy_django_plugin.transformers.managers import ( - construct_as_manager_instance, - create_new_manager_class_from_as_manager_method, + add_as_manager_to_queryset_class, create_new_manager_class_from_from_queryset_method, reparametrize_any_manager_hook, resolve_manager_method, @@ -209,10 +208,6 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M fullnames.REVERSE_MANY_TO_ONE_DESCRIPTOR: manytoone.refine_many_to_one_related_manager, } return hooks.get(class_fullname) - elif method_name == "as_manager": - info = self._get_typeinfo_or_none(class_fullname) - if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME): - return partial(construct_as_manager_instance, info=info) if method_name in self.manager_and_queryset_method_hooks: info = self._get_typeinfo_or_none(class_fullname) @@ -250,6 +245,10 @@ def get_base_class_hook(self, fullname: str) -> Optional[Callable[[ClassDefConte # Base class is a Form class definition if fullname in self._get_current_form_bases(): return transform_form_class + + # Base class is a QuerySet class definition + if sym is not None and isinstance(sym.node, TypeInfo) and sym.node.has_base(fullnames.QUERYSET_CLASS_FULLNAME): + return add_as_manager_to_queryset_class return None def get_attribute_hook(self, fullname: str) -> Optional[Callable[[AttributeContext], MypyType]]: @@ -308,10 +307,6 @@ def get_dynamic_class_hook(self, fullname: str) -> Optional[Callable[[DynamicCla info = self._get_typeinfo_or_none(class_name) if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME): return create_new_manager_class_from_from_queryset_method - elif method_name == "as_manager": - info = self._get_typeinfo_or_none(class_name) - if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME): - return create_new_manager_class_from_as_manager_method return None def report_config_data(self, ctx: ReportConfigContext) -> Dict[str, Any]: diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index ec13bf20c..04c5cf4ea 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -16,7 +16,8 @@ SymbolTableNode, TypeInfo, ) -from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext +from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext +from mypy.plugins.common import add_method_to_class from mypy.semanal import SemanticAnalyzer from mypy.semanal_shared import has_placeholder from mypy.subtypes import find_member @@ -482,44 +483,37 @@ def populate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeIn ) -def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext) -> None: - """ - Insert a new manager class node for a - - ``` - = .as_manager() - ``` - """ +def add_as_manager_to_queryset_class(ctx: ClassDefContext) -> None: semanal_api = helpers.get_semanal_api(ctx) - # Don't redeclare the manager class if we've already defined it. - manager_node = semanal_api.lookup_current_scope(ctx.name) - if manager_node and manager_node.type is not None: - # This is just a deferral run where our work is already finished - return - manager_sym = semanal_api.lookup_fully_qualified_or_none(fullnames.MANAGER_CLASS_FULLNAME) - assert manager_sym is not None - manager_base = manager_sym.node - if manager_base is None: + def _defer() -> None: if not semanal_api.final_iteration: semanal_api.defer() - return - assert isinstance(manager_base, TypeInfo) + queryset_info = semanal_api.type + if queryset_info is None: + return _defer() - callee = ctx.call.callee - assert isinstance(callee, MemberExpr) - assert isinstance(callee.expr, RefExpr) + # either a manual `as_manager` definition or this is a deferral pass + if "as_manager" in queryset_info.names: + return - queryset_info = callee.expr.node - if queryset_info is None: - if not semanal_api.final_iteration: - semanal_api.defer() + base_as_manager = queryset_info.get("as_manager") + if ( + base_as_manager is None + or not isinstance(base_as_manager.type, CallableType) + or not isinstance(base_as_manager.type.ret_type, Instance) + ): return - assert isinstance(queryset_info, TypeInfo) + base_ret_type = base_as_manager.type.ret_type.type + + manager_sym = semanal_api.lookup_fully_qualified_or_none(fullnames.MANAGER_CLASS_FULLNAME) + if manager_sym is None or not isinstance(manager_sym.node, TypeInfo): + return _defer() - manager_class_name = manager_base.name + "From" + queryset_info.name + manager_base = manager_sym.node + manager_class_name = f"{manager_base.name}From{queryset_info.name}" current_module = semanal_api.modules[semanal_api.cur_mod_id] existing_sym = current_module.names.get(manager_class_name) if ( @@ -535,15 +529,13 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext) try: new_manager_info = create_manager_class( api=semanal_api, - base_manager_info=manager_base, + base_manager_info=base_ret_type, name=manager_class_name, - line=ctx.call.line, + line=queryset_info.line, with_unique_name=True, ) except helpers.IncompleteDefnException: - if not semanal_api.final_iteration: - semanal_api.defer() - return + return _defer() populate_manager_from_queryset(new_manager_info, queryset_info) register_dynamically_created_manager( @@ -551,38 +543,23 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext) manager_name=manager_class_name, manager_base=manager_base, ) - queryset_info.metadata.setdefault("django_as_manager_names", {}) - queryset_info.metadata["django_as_manager_names"][semanal_api.cur_mod_id] = new_manager_info.name # Add the new manager to the current module - added = semanal_api.add_symbol_table_node( - # We'll use `new_manager_info.name` instead of `manager_class_name` here - # to handle possible name collisions, as it's unique. - new_manager_info.name, + # We'll use `new_manager_info.name` instead of `manager_class_name` here + # to handle possible name collisions, as it's unique. + current_module.names[new_manager_info.name] = ( # Note that the generated manager type is always inserted at module level - SymbolTableNode(GDEF, new_manager_info, plugin_generated=True), + SymbolTableNode(GDEF, new_manager_info, plugin_generated=True) ) - assert added - -def construct_as_manager_instance(ctx: MethodContext, *, info: TypeInfo) -> MypyType: - api = helpers.get_typechecker_api(ctx) - module = helpers.get_current_module(api) - try: - manager_name = info.metadata["django_as_manager_names"][module.fullname] - except KeyError: - return ctx.default_return_type - - manager_node = api.lookup(manager_name) - if not isinstance(manager_node.node, TypeInfo): - return ctx.default_return_type - - # Whenever `.as_manager()` isn't called at class level, we want to ensure - # that the variable is an instance of our generated manager. Instead of the return - # value of `.as_manager()`. Though model argument is populated as `Any`. - # `transformers.models.AddManagers` will populate a model's manager(s), when it - # finds it on class level. - return Instance(manager_node.node, [AnyType(TypeOfAny.from_omitted_generics)]) + add_method_to_class( + semanal_api, + ctx.cls, + "as_manager", + args=[], + return_type=Instance(new_manager_info, [AnyType(TypeOfAny.from_omitted_generics)]), + is_classmethod=True, + ) def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None: