From 830d74b4937781a5732f8384b46d9b38d832436b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigurd=20Lj=C3=B8dal?= <544451+ljodal@users.noreply.github.com> Date: Wed, 13 Jul 2022 09:04:44 +0200 Subject: [PATCH] Add support for inline `from_queryset` in model classes (#1045) * Add support for inline from_queryset in model classes This adds support for calling .from_queryset()() inline in models, for example like this: class MyModel(models.Model): objects = MyManager.from_queryset(MyQuerySet)() This is done by inspecting the class body in the transform_class_hook * Fix missing methods on copied manager * Add test and other minor tweaks * Always create manager at module level When the manager is added at the class level, which happened when it was created inline in the model body, it's not possible to retrieve the manager again based on fullname. That lead to problems with inheritance and the default manager. --- mypy_django_plugin/lib/helpers.py | 44 +++-- mypy_django_plugin/main.py | 6 - mypy_django_plugin/transformers/managers.py | 181 +++++++++--------- mypy_django_plugin/transformers/models.py | 118 ++++++++---- tests/typecheck/fields/test_related.yml | 6 +- .../managers/querysets/test_from_queryset.yml | 64 ++++--- tests/typecheck/managers/test_managers.yml | 4 +- 7 files changed, 244 insertions(+), 179 deletions(-) diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index 1ce73fd8b..abbdccb2d 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/mypy_django_plugin/lib/helpers.py @@ -208,6 +208,23 @@ def is_annotated_model_fullname(model_cls_fullname: str) -> bool: return model_cls_fullname.startswith(WITH_ANNOTATIONS_FULLNAME + "[") +def create_type_info(name: str, module: str, bases: List[Instance]) -> TypeInfo: + + # make new class expression + classdef = ClassDef(name, Block([])) + classdef.fullname = module + "." + name + + # make new TypeInfo + new_typeinfo = TypeInfo(SymbolTable(), classdef, module) + new_typeinfo.bases = bases + calculate_mro(new_typeinfo) + new_typeinfo.calculate_metaclass_type() + + classdef.info = new_typeinfo + + return new_typeinfo + + def add_new_class_for_module( module: MypyFile, name: str, @@ -217,15 +234,7 @@ def add_new_class_for_module( ) -> TypeInfo: new_class_unique_name = checker.gen_unique_name(name, module.names) - # make new class expression - classdef = ClassDef(new_class_unique_name, Block([])) - classdef.fullname = module.fullname + "." + new_class_unique_name - - # make new TypeInfo - new_typeinfo = TypeInfo(SymbolTable(), classdef, module.fullname) - new_typeinfo.bases = bases - calculate_mro(new_typeinfo) - new_typeinfo.calculate_metaclass_type() + new_typeinfo = create_type_info(new_class_unique_name, module.fullname, bases) # add fields if fields: @@ -237,7 +246,6 @@ def add_new_class_for_module( MDEF, var, plugin_generated=True, no_serialize=no_serialize ) - classdef.info = new_typeinfo module.names[new_class_unique_name] = SymbolTableNode( GDEF, new_typeinfo, plugin_generated=True, no_serialize=no_serialize ) @@ -382,29 +390,25 @@ def copy_method_to_another_class( method_node: FuncDef, return_type: Optional[MypyType] = None, original_module_name: Optional[str] = None, -) -> None: +) -> bool: semanal_api = get_semanal_api(ctx) if method_node.type is None: - if not semanal_api.final_iteration: - semanal_api.defer() - return - arguments, return_type = build_unannotated_method_args(method_node) add_method_to_class( semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type ) - return + return True method_type = method_node.type if not isinstance(method_type, CallableType): if not semanal_api.final_iteration: semanal_api.defer() - return + return False if return_type is None: return_type = bind_or_analyze_type(method_type.ret_type, semanal_api, original_module_name) if return_type is None: - return + return False # We build the arguments from the method signature (`CallableType`), because if we were to # use the arguments from the method node (`FuncDef.arguments`) we're not compatible with @@ -417,7 +421,7 @@ def copy_method_to_another_class( ): bound_arg_type = bind_or_analyze_type(arg_type, semanal_api, original_module_name) if bound_arg_type is None: - return + return False if arg_name is None and hasattr(method_node, "arguments"): arg_name = method_node.arguments[pos].variable.name arguments.append( @@ -435,6 +439,8 @@ def copy_method_to_another_class( semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type ) + return True + def add_new_manager_base(api: SemanticAnalyzerPluginInterface, fullname: str) -> None: sym = api.lookup_fully_qualified_or_none(fullnames.MANAGER_CLASS_FULLNAME) diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index 6d3550d62..ad449324d 100644 --- a/mypy_django_plugin/main.py +++ b/mypy_django_plugin/main.py @@ -24,7 +24,6 @@ from mypy_django_plugin.transformers import fields, forms, init_create, meta, querysets, request, settings from mypy_django_plugin.transformers.managers import ( create_new_manager_class_from_from_queryset_method, - fail_if_manager_type_created_in_model_body, resolve_manager_method, ) from mypy_django_plugin.transformers.models import ( @@ -237,11 +236,6 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M django_context=self.django_context, ) - elif method_name == "from_queryset": - info = self._get_typeinfo_or_none(class_fullname) - if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME): - return fail_if_manager_type_created_in_model_body - return None def get_base_class_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]: diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 28ac77418..ce41b6ee9 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -8,7 +8,6 @@ FuncBase, FuncDef, MemberExpr, - NameExpr, OverloadedFuncDef, RefExpr, StrExpr, @@ -16,13 +15,12 @@ TypeInfo, Var, ) -from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext +from mypy.plugin import AttributeContext, DynamicClassDefContext, SemanticAnalyzerPluginInterface from mypy.types import AnyType, CallableType, Instance, ProperType from mypy.types import Type as MypyType from mypy.types import TypeOfAny from typing_extensions import Final -from mypy_django_plugin import errorcodes from mypy_django_plugin.lib import fullnames, helpers MANAGER_METHODS_RETURNING_QUERYSET: Final = frozenset( @@ -182,81 +180,110 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte """ semanal_api = helpers.get_semanal_api(ctx) + # TODO: Emit an error when called in a class scope + if semanal_api.is_class_scope(): + return + # 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 isinstance(manager_node.node, TypeInfo): # This is just a deferral run where our work is already finished return - callee = ctx.call.callee - assert isinstance(callee, MemberExpr) - assert isinstance(callee.expr, RefExpr) - - base_manager_info = callee.expr.node - if base_manager_info is None: - if not semanal_api.final_iteration: - semanal_api.defer() + new_manager_info = create_manager_info_from_from_queryset_call(ctx.api, ctx.call, ctx.name) + if new_manager_info is None: + if not ctx.api.final_iteration: + ctx.api.defer() return - assert isinstance(base_manager_info, TypeInfo) + # So that the plugin will reparameterize the manager when it is constructed inside of a Model definition + helpers.add_new_manager_base(semanal_api, new_manager_info.fullname) - passed_queryset = ctx.call.args[0] - assert isinstance(passed_queryset, NameExpr) - derived_queryset_fullname = passed_queryset.fullname - if derived_queryset_fullname is None: - # In some cases, due to the way the semantic analyzer works, only passed_queryset.name is available. - # But it should be analyzed again, so this isn't a problem. - return +def create_manager_info_from_from_queryset_call( + api: SemanticAnalyzerPluginInterface, call_expr: CallExpr, name: Optional[str] = None +) -> Optional[TypeInfo]: + """ + Extract manager and queryset TypeInfo from a from_queryset call. + """ - base_manager_instance = fill_typevars(base_manager_info) - assert isinstance(base_manager_instance, Instance) - new_manager_info = semanal_api.basic_new_typeinfo( - ctx.name, basetype_or_fallback=base_manager_instance, line=ctx.call.line - ) + if ( + # Check that this is a from_queryset call on a manager subclass + not isinstance(call_expr.callee, MemberExpr) + or not isinstance(call_expr.callee.expr, RefExpr) + or not isinstance(call_expr.callee.expr.node, TypeInfo) + or not call_expr.callee.expr.node.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME) + or not call_expr.callee.name == "from_queryset" + # Check that the call has one or two arguments and that the first is a + # QuerySet subclass + or not 1 <= len(call_expr.args) <= 2 + or not isinstance(call_expr.args[0], RefExpr) + or not isinstance(call_expr.args[0].node, TypeInfo) + or not call_expr.args[0].node.has_base(fullnames.QUERYSET_CLASS_FULLNAME) + ): + return None - sym = semanal_api.lookup_fully_qualified_or_none(derived_queryset_fullname) - assert sym is not None - if sym.node is None: - if not semanal_api.final_iteration: - semanal_api.defer() - else: - # inherit from Any to prevent false-positives, if queryset class cannot be resolved - new_manager_info.fallback_to_any = True - return + base_manager_info, queryset_info = call_expr.callee.expr.node, call_expr.args[0].node + if queryset_info.fullname is None: + # In some cases, due to the way the semantic analyzer works, only + # passed_queryset.name is available. But it should be analyzed again, + # so this isn't a problem. + return None - derived_queryset_info = sym.node - assert isinstance(derived_queryset_info, TypeInfo) - - new_manager_info.line = ctx.call.line - new_manager_info.type_vars = base_manager_info.type_vars - new_manager_info.defn.type_vars = base_manager_info.defn.type_vars - new_manager_info.defn.line = ctx.call.line - new_manager_info.metaclass_type = new_manager_info.calculate_metaclass_type() - # Stash the queryset fullname which was passed to .from_queryset - # So that our 'resolve_manager_method' attribute hook can fetch the method from that QuerySet class - new_manager_info.metadata["django"] = {"from_queryset_manager": derived_queryset_fullname} - - if len(ctx.call.args) > 1: - expr = ctx.call.args[1] - assert isinstance(expr, StrExpr) - custom_manager_generated_name = expr.value + if len(call_expr.args) == 2 and isinstance(call_expr.args[1], StrExpr): + manager_name = call_expr.args[1].value else: - custom_manager_generated_name = base_manager_info.name + "From" + derived_queryset_info.name + manager_name = f"{base_manager_info.name}From{queryset_info.name}" - custom_manager_generated_fullname = ".".join(["django.db.models.manager", custom_manager_generated_name]) + new_manager_info = create_manager_class(api, base_manager_info, name or manager_name, call_expr.line) + + popuplate_manager_from_queryset(new_manager_info, queryset_info) + + manager_fullname = ".".join(["django.db.models.manager", manager_name]) + + base_manager_info = new_manager_info.mro[1] base_manager_info.metadata.setdefault("from_queryset_managers", {}) - base_manager_info.metadata["from_queryset_managers"][custom_manager_generated_fullname] = new_manager_info.fullname + base_manager_info.metadata["from_queryset_managers"][manager_fullname] = new_manager_info.fullname + + # Add the new manager to the current module + module = api.modules[api.cur_mod_id] + module.names[name or manager_name] = SymbolTableNode( + GDEF, new_manager_info, plugin_generated=True, no_serialize=False + ) + + return new_manager_info - # So that the plugin will reparameterize the manager when it is constructed inside of a Model definition - helpers.add_new_manager_base(semanal_api, new_manager_info.fullname) - class_def_context = ClassDefContext(cls=new_manager_info.defn, reason=ctx.call, api=semanal_api) - self_type = fill_typevars(new_manager_info) - assert isinstance(self_type, Instance) +def create_manager_class( + api: SemanticAnalyzerPluginInterface, base_manager_info: TypeInfo, name: str, line: int +) -> TypeInfo: + + base_manager_instance = fill_typevars(base_manager_info) + assert isinstance(base_manager_instance, Instance) + + manager_info = helpers.create_type_info(name, api.cur_mod_id, bases=[base_manager_instance]) + manager_info.line = line + manager_info.type_vars = base_manager_info.type_vars + manager_info.defn.type_vars = base_manager_info.defn.type_vars + manager_info.defn.line = line + manager_info.metaclass_type = manager_info.calculate_metaclass_type() + + return manager_info + + +def popuplate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeInfo) -> None: + """ + Add methods from the QuerySet class to the manager. + """ + + # Stash the queryset fullname which was passed to .from_queryset So that + # our 'resolve_manager_method' attribute hook can fetch the method from + # that QuerySet class + django_metadata = helpers.get_django_metadata(manager_info) + django_metadata["from_queryset_manager"] = queryset_info.fullname # We collect and mark up all methods before django.db.models.query.QuerySet as class members - for class_mro_info in derived_queryset_info.mro: + for class_mro_info in queryset_info.mro: if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME: break for name, sym in class_mro_info.names.items(): @@ -270,39 +297,19 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte # queryset_method: Any = ... # helpers.add_new_sym_for_info( - new_manager_info, + manager_info, name=name, sym_type=AnyType(TypeOfAny.special_form), ) - # For methods on BaseManager that return a queryset we need to update the - # return type to be the actual queryset subclass used. This is done by - # adding the methods as attributes with type Any to the manager class, - # similar to how custom queryset methods are handled above. The actual type - # of these methods are resolved in resolve_manager_method. - for name in MANAGER_METHODS_RETURNING_QUERYSET: + # For methods on BaseManager that return a queryset we need to update + # the return type to be the actual queryset subclass used. This is done + # by adding the methods as attributes with type Any to the manager + # class. The actual type of these methods are resolved in + # resolve_manager_method. + for method_name in MANAGER_METHODS_RETURNING_QUERYSET: helpers.add_new_sym_for_info( - new_manager_info, - name=name, + manager_info, + name=method_name, sym_type=AnyType(TypeOfAny.special_form), ) - - # Insert the new manager (dynamic) class - assert semanal_api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, new_manager_info, plugin_generated=True)) - - -def fail_if_manager_type_created_in_model_body(ctx: MethodContext) -> MypyType: - """ - Method hook that checks if method `.from_queryset` is called inside a model class body. - - Doing so won't, for instance, trigger the dynamic class hook(`create_new_manager_class_from_from_queryset_method`) - for managers. - """ - api = helpers.get_typechecker_api(ctx) - outer_model_info = api.scope.active_class() - if not outer_model_info or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME): - # Not inside a model class definition - return ctx.default_return_type - - api.fail("`.from_queryset` called from inside model class body", ctx.context, code=errorcodes.MANAGER_UNTYPED) - return ctx.default_return_type diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 6b8e50276..a4d2f7a1d 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -1,11 +1,11 @@ from typing import Dict, List, Optional, Type, Union, cast -from django.db.models.base import Model +from django.db.models import Manager, Model from django.db.models.fields import DateField, DateTimeField, Field from django.db.models.fields.related import ForeignKey from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel, OneToOneRel from mypy.checker import TypeChecker -from mypy.nodes import ARG_STAR2, Argument, AssignmentStmt, Context, FuncDef, NameExpr, TypeInfo, Var +from mypy.nodes import ARG_STAR2, Argument, AssignmentStmt, CallExpr, Context, FuncDef, NameExpr, TypeInfo, Var from mypy.plugin import AnalyzeTypeContext, AttributeContext, CheckerPluginInterface, ClassDefContext from mypy.plugins import common from mypy.semanal import SemanticAnalyzer @@ -17,9 +17,9 @@ from mypy_django_plugin.errorcodes import MANAGER_MISSING from mypy_django_plugin.lib import fullnames, helpers from mypy_django_plugin.lib.fullnames import ANNOTATIONS_FULLNAME, ANY_ATTR_ALLOWED_CLASS_FULLNAME, MODEL_CLASS_FULLNAME -from mypy_django_plugin.lib.helpers import add_new_class_for_module from mypy_django_plugin.transformers import fields from mypy_django_plugin.transformers.fields import get_field_descriptor_types +from mypy_django_plugin.transformers.managers import create_manager_info_from_from_queryset_call class ModelClassInitializer: @@ -198,12 +198,12 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty bases = [] for original_base in base_manager_info.bases: if self.is_any_parametrized_manager(original_base): - if original_base.type is None: - raise helpers.IncompleteDefnException() - original_base = helpers.reparametrize_instance(original_base, [Instance(self.model_classdef.info, [])]) bases.append(original_base) + # TODO: This adds the manager to the module, even if we end up + # deferring. That can be avoided by not adding it to the module first, + # but rather waiting until we know we won't defer new_manager_info = self.add_new_class_for_current_module(name, bases) # copy fields to a new manager new_cls_def_context = ClassDefContext(cls=new_manager_info.defn, reason=self.ctx.reason, api=self.api) @@ -212,13 +212,15 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty for name, sym in base_manager_info.names.items(): # replace self type with new class, if copying method if isinstance(sym.node, FuncDef): - helpers.copy_method_to_another_class( + copied_method = helpers.copy_method_to_another_class( new_cls_def_context, self_type=custom_manager_type, new_method_name=name, method_node=sym.node, original_module_name=base_manager_info.module_name, ) + if not copied_method and not self.api.final_iteration: + raise helpers.IncompleteDefnException() continue new_sym = sym.copy() @@ -236,23 +238,23 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: incomplete_manager_defs = set() for manager_name, manager in model_cls._meta.managers_map.items(): + # If the manager is already typed do nothing + manager_node = self.model_classdef.info.names.get(manager_name, None) + if manager_node and manager_node.type is not None: + continue + manager_class_name = manager.__class__.__name__ manager_fullname = helpers.get_class_fullname(manager.__class__) - try: - manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname) - except helpers.IncompleteDefnException as exc: - # Check if manager is a generated (dynamic class) manager - base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0]) - generated_managers = self.get_generated_manager_mappings(base_manager_fullname) - if manager_fullname not in generated_managers: - # Manager doesn't appear to be generated. Track that we encountered an - # incomplete definition and skip - incomplete_manager_defs.add(manager_name) - continue - manager_info = self.lookup_typeinfo(generated_managers[manager_fullname]) - if manager_info is None: - continue + manager_info = self.lookup_typeinfo(manager_fullname) + if manager_info is None: + manager_info = self.create_manager_from_from_queryset(manager_name) + if manager_info is None: + manager_info = self.get_dynamic_manager(manager_fullname, manager) + + if manager_info is None: + incomplete_manager_defs.add(manager_name) + continue is_dynamically_generated = manager_info.metadata.get("django", {}).get("from_queryset_manager") is not None if manager_name not in self.model_classdef.info.names or is_dynamically_generated: @@ -264,18 +266,20 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: # related manager. custom_model_manager_name = manager.model.__name__ + "_" + manager_class_name try: - custom_manager_type = self.create_new_model_parametrized_manager( + manager_type = self.create_new_model_parametrized_manager( custom_model_manager_name, base_manager_info=manager_info ) except helpers.IncompleteDefnException: continue - self.add_new_node_to_model_class(manager_name, custom_manager_type) + self.add_new_node_to_model_class(manager_name, manager_type) + + if incomplete_manager_defs: + if not self.api.final_iteration: + # Unless we're on the final round, see if another round could + # figure out all manager types + raise helpers.IncompleteDefnException() - if incomplete_manager_defs and not self.api.final_iteration: - # Unless we're on the final round, see if another round could figure out all manager types - raise helpers.IncompleteDefnException() - elif self.api.final_iteration: for manager_name in incomplete_manager_defs: # We act graceful and set the type as the bare minimum we know of # (Django's default) before finishing. And emit an error, to allow for @@ -289,22 +293,59 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: manager_name, Instance(django_manager_info, [Instance(self.model_classdef.info, [])]) ) # Find expression for e.g. `objects = SomeManager()` - manager_expr = [ - expr - for expr in self.ctx.cls.defs.body - if ( - isinstance(expr, AssignmentStmt) - and isinstance(expr.lvalues[0], NameExpr) - and expr.lvalues[0].name == manager_name - ) - ] + manager_expr = self.get_manager_expression(manager_name) manager_fullname = f"{self.model_classdef.fullname}.{manager_name}" self.api.fail( f'Could not resolve manager type for "{manager_fullname}"', - manager_expr[0] if manager_expr else self.ctx.cls, + manager_expr if manager_expr else self.ctx.cls, code=MANAGER_MISSING, ) + def get_manager_expression(self, name: str) -> Optional[AssignmentStmt]: + # TODO: What happens if the manager is defined multiple times? + for expr in self.ctx.cls.defs.body: + if ( + isinstance(expr, AssignmentStmt) + and isinstance(expr.lvalues[0], NameExpr) + and expr.lvalues[0].name == name + ): + return expr + + return None + + def get_dynamic_manager(self, fullname: str, manager: Manager) -> Optional[TypeInfo]: + """ + Try to get a dynamically defined manager + """ + + # Check if manager is a generated (dynamic class) manager + base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0]) + generated_managers = self.get_generated_manager_mappings(base_manager_fullname) + + generated_manager_name: Optional[str] = generated_managers.get(fullname, None) + if generated_manager_name is None: + return None + + return self.lookup_typeinfo(generated_manager_name) + + def create_manager_from_from_queryset(self, name: str) -> Optional[TypeInfo]: + """ + Try to create a manager from a .from_queryset call: + + class MyModel(models.Model): + objects = MyManager.from_queryset(MyQuerySet)() + """ + + assign_statement = self.get_manager_expression(name) + if assign_statement is None: + return None + + expr = assign_statement.rvalue + if not isinstance(expr, CallExpr) or not isinstance(expr.callee, CallExpr): + return None + + return create_manager_info_from_from_queryset_call(self.api, expr.callee) + class AddDefaultManagerAttribute(ModelClassInitializer): def run_with_model_cls(self, model_cls: Type[Model]) -> None: @@ -313,6 +354,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: default_manager_cls = model_cls._meta.default_manager.__class__ default_manager_fullname = helpers.get_class_fullname(default_manager_cls) + try: default_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(default_manager_fullname) except helpers.IncompleteDefnException as exc: @@ -580,7 +622,7 @@ def get_or_create_annotated_type( else: annotated_model_type = api.named_generic_type(ANY_ATTR_ALLOWED_CLASS_FULLNAME, []) - annotated_typeinfo = add_new_class_for_module( + annotated_typeinfo = helpers.add_new_class_for_module( model_module_file, type_name, bases=[model_type] if fields_dict is not None else [model_type, annotated_model_type], diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index af416c79f..c7f726732 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -678,7 +678,6 @@ def custom(self) -> None: pass - # Note, that we cannot resolve dynamic calls for custom managers: class Transaction(models.Model): objects = BaseManager.from_queryset(TransactionQuerySet) def test(self) -> None: @@ -689,9 +688,8 @@ class TransactionLog(models.Model): transaction = models.ForeignKey(Transaction, on_delete=models.CASCADE) out: | - myapp/models:9: error: `.from_queryset` called from inside model class body - myapp/models:11: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.TransactionLog]" - myapp/models:13: note: Revealed type is "Any" + myapp/models:10: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.TransactionLog]" + myapp/models:12: note: Revealed type is "Any" - case: resolve_primary_keys_for_foreign_keys_with_abstract_self_model diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index aadd114e7..ae69499f8 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -359,12 +359,17 @@ class MyModel(models.Model): objects = NewManager() -- case: from_queryset_in_model_class_body_yields_message +- case: test_queryset_in_model_class_body main: | from myapp.models import MyModel - reveal_type(MyModel.base_manager) # N: Revealed type is "myapp.models.BaseManagerFromMyQuerySet[myapp.models.MyModel]" - reveal_type(MyModel.manager) # N: Revealed type is "myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]" - reveal_type(MyModel.custom_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.custom) # N: Revealed type is "def () -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.all().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.custom().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects2) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" installed_apps: - myapp files: @@ -372,29 +377,42 @@ - path: myapp/models.py content: | from django.db import models - from django.db.models.manager import BaseManager - class MyQuerySet(models.QuerySet["MyModel"]): - def queryset_method(self) -> int: - return 1 + class MyManager(models.Manager["MyModel"]): + pass - class MyManager(BaseManager): - ... + class MyQuerySet(models.QuerySet["MyModel"]): + def custom(self) -> "MyQuerySet": + pass - BaseManagerFromMyQuerySet = BaseManager.from_queryset(MyQuerySet) - ManagerFromMyQuerySet = models.Manager.from_queryset(MyQuerySet) - MyManagerFromMyQuerySet = MyManager.from_queryset(MyQuerySet) class MyModel(models.Model): - objects1 = BaseManager.from_queryset(MyQuerySet)() # E: `.from_queryset` called from inside model class body - objects2 = BaseManager.from_queryset(MyQuerySet) # E: `.from_queryset` called from inside model class body - objects3 = models.Manager.from_queryset(MyQuerySet)() # E: `.from_queryset` called from inside model class body - objects4 = models.Manager.from_queryset(MyQuerySet) # E: `.from_queryset` called from inside model class body - objects5 = MyManager.from_queryset(MyQuerySet) # E: `.from_queryset` called from inside model class body - objects6 = MyManager.from_queryset(MyQuerySet)() # E: `.from_queryset` called from inside model class body - # Initiating the manager type is fine - base_manager = BaseManagerFromMyQuerySet() - manager = ManagerFromMyQuerySet() - custom_manager = MyManagerFromMyQuerySet() + objects = MyManager.from_queryset(MyQuerySet)() + objects2 = MyManager.from_queryset(MyQuerySet)() + +- case: test_queryset_in_model_class_body_subclass + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.BaseManagerFromBaseQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class BaseManager(models.Manager["BaseModel"]): + pass + + class BaseQuerySet(models.QuerySet["BaseModel"]): + def custom(self) -> "BaseQuerySet": + pass + + class BaseModel(models.Model): + objects = BaseManager.from_queryset(BaseQuerySet)() + + class MyModel(BaseModel): + pass - case: from_queryset_includes_methods_returning_queryset main: | diff --git a/tests/typecheck/managers/test_managers.yml b/tests/typecheck/managers/test_managers.yml index 99778f82d..51ee4467e 100644 --- a/tests/typecheck/managers/test_managers.yml +++ b/tests/typecheck/managers/test_managers.yml @@ -332,14 +332,14 @@ - case: custom_manager_returns_proper_model_types main: | from myapp.models import User - reveal_type(User.objects) # N: Revealed type is "myapp.models.User_MyManager2[myapp.models.User]" + reveal_type(User.objects) # N: Revealed type is "myapp.models.User_MyManager[myapp.models.User]" reveal_type(User.objects.select_related()) # N: Revealed type is "django.db.models.query._QuerySet[myapp.models.User, myapp.models.User]" reveal_type(User.objects.get()) # N: Revealed type is "myapp.models.User" reveal_type(User.objects.get_instance()) # N: Revealed type is "builtins.int" reveal_type(User.objects.get_instance_untyped('hello')) # N: Revealed type is "Any" from myapp.models import ChildUser - reveal_type(ChildUser.objects) # N: Revealed type is "myapp.models.ChildUser_MyManager2[myapp.models.ChildUser]" + reveal_type(ChildUser.objects) # N: Revealed type is "myapp.models.ChildUser_MyManager[myapp.models.ChildUser]" reveal_type(ChildUser.objects.select_related()) # N: Revealed type is "django.db.models.query._QuerySet[myapp.models.ChildUser, myapp.models.ChildUser]" reveal_type(ChildUser.objects.get()) # N: Revealed type is "myapp.models.ChildUser" reveal_type(ChildUser.objects.get_instance()) # N: Revealed type is "builtins.int"