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"