diff --git a/README.md b/README.md index a392f80ab..32a469ba1 100644 --- a/README.md +++ b/README.md @@ -140,9 +140,7 @@ And then use `AuthenticatedHttpRequest` instead of the standard `HttpRequest` fo ### My QuerySet methods are returning Any rather than my Model -`QuerySet.as_manager()` is not currently supported. - -If you are using `MyQuerySet.as_manager()`, then your `Manager`/`QuerySet` methods will all not be linked to your model. +If you are using `MyQuerySet.as_manager()`: Example: @@ -156,12 +154,12 @@ class MyModel(models.Model): bar = models.IntegerField() objects = MyModelQuerySet.as_manager() -def use_my_model(): - foo = MyModel.objects.get(id=1) # This is `Any` but it should be `MyModel` - return foo.xyz # No error, but there should be +def use_my_model() -> int: + foo = MyModel.objects.get(id=1) # Should now be `MyModel` + return foo.xyz # Gives an error ``` -There is a workaround: use `Manager.from_queryset` instead. +Or if you're using `Manager.from_queryset`: Example: @@ -177,11 +175,14 @@ class MyModel(models.Model): bar = models.IntegerField() objects = MyModelManager() -def use_my_model(): - foo = MyModel.objects.get(id=1) +def use_my_model() -> int: + foo = MyModel.objects.get(id=1) # Should now be `MyModel` return foo.xyz # Gives an error ``` +Take note that `.from_queryset` needs to be placed at _module level_ to annotate +types correctly. Calling `.from_queryset` in class definition will yield an error. + ### How do I annotate cases where I called QuerySet.annotate? Django-stubs provides a special type, `django_stubs_ext.WithAnnotations[Model]`, which indicates that the `Model` has diff --git a/mypy_django_plugin/errorcodes.py b/mypy_django_plugin/errorcodes.py index 76af6dbaa..aab6d5cb9 100644 --- a/mypy_django_plugin/errorcodes.py +++ b/mypy_django_plugin/errorcodes.py @@ -2,3 +2,4 @@ MANAGER_UNTYPED = ErrorCode("django-manager", "Untyped manager disallowed", "Django") MANAGER_MISSING = ErrorCode("django-manager-missing", "Couldn't resolve manager for model", "Django") +MODEL_ARG_MISMATCH = ErrorCode("django-model-arg", "Model argument mismatching between manager and queryset", "Django") diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index 1ce73fd8b..c8a9f48dc 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/mypy_django_plugin/lib/helpers.py @@ -376,33 +376,31 @@ def bind_or_analyze_type(t: MypyType, api: SemanticAnalyzer, module_name: Option def copy_method_to_another_class( - ctx: ClassDefContext, + api: SemanticAnalyzer, + cls: ClassDef, self_type: Instance, new_method_name: str, method_node: FuncDef, return_type: Optional[MypyType] = None, original_module_name: Optional[str] = None, ) -> None: - semanal_api = get_semanal_api(ctx) if method_node.type is None: - if not semanal_api.final_iteration: - semanal_api.defer() + if not api.final_iteration: + 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 - ) + add_method_to_class(api, cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type) return method_type = method_node.type if not isinstance(method_type, CallableType): - if not semanal_api.final_iteration: - semanal_api.defer() + if not api.final_iteration: + api.defer() return if return_type is None: - return_type = bind_or_analyze_type(method_type.ret_type, semanal_api, original_module_name) + return_type = bind_or_analyze_type(method_type.ret_type, api, original_module_name) if return_type is None: return @@ -415,7 +413,7 @@ def copy_method_to_another_class( zip(method_type.arg_types[1:], method_type.arg_kinds[1:], method_type.arg_names[1:]), start=1, ): - bound_arg_type = bind_or_analyze_type(arg_type, semanal_api, original_module_name) + bound_arg_type = bind_or_analyze_type(arg_type, api, original_module_name) if bound_arg_type is None: return if arg_name is None and hasattr(method_node, "arguments"): @@ -431,9 +429,7 @@ def copy_method_to_another_class( ) ) - add_method_to_class( - semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type - ) + add_method_to_class(api, cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type) def add_new_manager_base(api: SemanticAnalyzerPluginInterface, fullname: str) -> None: diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index 6d3550d62..51ea09f9a 100644 --- a/mypy_django_plugin/main.py +++ b/mypy_django_plugin/main.py @@ -23,6 +23,7 @@ from mypy_django_plugin.lib import fullnames, helpers 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_as_manager_method, create_new_manager_class_from_from_queryset_method, fail_if_manager_type_created_in_model_body, resolve_manager_method, @@ -303,11 +304,15 @@ def get_type_analyze_hook(self, fullname: str) -> Optional[Callable[[AnalyzeType def get_dynamic_class_hook(self, fullname: str) -> Optional[Callable[[DynamicClassDefContext], None]]: # Create a new manager class definition when a manager's '.from_queryset' classmethod is called - if fullname.endswith("from_queryset"): - class_name, _, _ = fullname.rpartition(".") + class_name, _, method_name = fullname.rpartition(".") + if method_name == "from_queryset": 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 diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 28ac77418..8a36d56e9 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -12,14 +12,17 @@ OverloadedFuncDef, RefExpr, StrExpr, + SymbolNode, SymbolTableNode, + TypeAlias, TypeInfo, Var, ) -from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext +from mypy.plugin import AttributeContext, DynamicClassDefContext, MethodContext +from mypy.semanal import SemanticAnalyzer from mypy.types import AnyType, CallableType, Instance, ProperType from mypy.types import Type as MypyType -from mypy.types import TypeOfAny +from mypy.types import TypeOfAny, TypeType, TypeVarType from typing_extensions import Final from mypy_django_plugin import errorcodes @@ -154,9 +157,11 @@ def resolve_manager_method(ctx: AttributeContext) -> MypyType: A 'get_attribute_hook' that is intended to be invoked whenever the TypeChecker encounters an attribute on a class that has 'django.db.models.BaseManager' as a base. """ - # Skip (method) type that is currently something other than Any + # Skip (method) type that is currently something other than Any of type `implementation_artifact` if not isinstance(ctx.default_attr_type, AnyType): return ctx.default_attr_type + elif ctx.default_attr_type.type_of_any != TypeOfAny.implementation_artifact: + return ctx.default_attr_type # (Current state is:) We wouldn't end up here when looking up a method from a custom _manager_. # That's why we only attempt to lookup the method for either a dynamically added or reverse manager. @@ -175,6 +180,170 @@ def resolve_manager_method(ctx: AttributeContext) -> MypyType: return AnyType(TypeOfAny.from_error) +def merge_queryset_into_manager( + api: SemanticAnalyzer, + manager_info: TypeInfo, + queryset_info: TypeInfo, + manager_base: TypeInfo, + class_name: str, +) -> Optional[TypeInfo]: + """ + Merges a queryset definition with a manager definition. + As a reference one can look at the 3-arg call definition for ``builtins.type`` to + get an idea of what kind of class we're trying to express a type for. + """ + # Stash the queryset fullname so that our 'resolve_manager_method' attribute hook + # can fetch the method from that QuerySet class + manager_info.metadata["django"] = {"from_queryset_manager": queryset_info.fullname} + + manager_base.metadata.setdefault("from_queryset_managers", {}) + # The `__module__` value of the manager type created by Django's `.from_queryset` + # is `django.db.models.manager`. But `basic_new_typeinfo` defaults to what is + # currently being processed, so we'll map that together through metadata. + manager_base.metadata["from_queryset_managers"][f"django.db.models.manager.{class_name}"] = manager_info.fullname + + # So that the plugin will reparameterize the manager when it is constructed inside of a Model definition + helpers.add_new_manager_base(api, manager_info.fullname) + + self_type = fill_typevars(manager_info) + assert isinstance(self_type, Instance) + + # We collect and mark up all methods before django.db.models.query.QuerySet as class members + 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(): + if not isinstance(sym.node, (FuncDef, Decorator)): + continue + # Insert the queryset method name as a class member. Note that the type of + # the method is set as Any. Figuring out the type is the job of the + # 'resolve_manager_method' attribute hook, which comes later. + # + # class BaseManagerFromMyQuerySet(BaseManager): + # queryset_method: Any = ... + # + helpers.add_new_sym_for_info( + manager_info, + name=name, + sym_type=AnyType(TypeOfAny.implementation_artifact), + ) + + # 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: + helpers.add_new_sym_for_info( + manager_info, + name=name, + sym_type=AnyType(TypeOfAny.implementation_artifact), + ) + + +def get_manager_model_argument_type(manager_info: TypeInfo) -> Optional[MypyType]: + """ + TODO: Fill me out + """ + manager_model_arg = next( + ( + base.args[0] + for base in helpers.iter_bases(manager_info) + if (base.type.fullname in fullnames.MANAGER_CLASSES and len(base.args) > 0) + ), + ( + AnyType(TypeOfAny.from_omitted_generics) + if manager_info.fullname == fullnames.BASE_MANAGER_CLASS_FULLNAME + else None + ), + ) + if isinstance(manager_model_arg, TypeVarType): + # Since `django.db.models.Manager` is a subclass of `BaseManager` we'll collect + # its default `TypeVar` when argument has been omitted. So now we'll refine the + # manager argument to be the type var's upper bound. + manager_model_arg = manager_model_arg.upper_bound + + return manager_model_arg + + +def intersect_manager_and_queryset_model_argument_type( + manager_info: TypeInfo, queryset_info: TypeInfo +) -> Optional[MypyType]: + """ + TODO: Fill me out + """ + # TODO: Add test with manager inheritance + # TODO: Add test with queryset inheritance + manager_model_arg = get_manager_model_argument_type(manager_info) + queryset_model_arg = next( + ( + base.args[0] + for base in helpers.iter_bases(queryset_info) + if base.type.fullname == fullnames.QUERYSET_CLASS_FULLNAME and len(base.args) > 0 + ), + ( + AnyType(TypeOfAny.from_omitted_generics) + if queryset_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME + else None + ), + ) + + if ( + isinstance(manager_model_arg, Instance) + and isinstance(queryset_model_arg, Instance) + and ( + manager_model_arg.type.fullname == queryset_model_arg.type.fullname + # Allow queryset model arg if manager model arg comes in as the default + # upper bound `django.db.models.base.Model` + or manager_model_arg.type.fullname == fullnames.MODEL_CLASS_FULLNAME + ) + ): + """ + Captures: + - CustomManager(BaseManager[CustomModel]) with CustomQuerySet(QuerySet[CustomModel]) + - CustomManager(Manager[CustomModel]) with CustomQuerySet(QuerySet[CustomModel]) + - Manager with CustomQuerySet(QuerySet[CustomModel]) + """ + return queryset_model_arg + elif ( + isinstance(manager_model_arg, AnyType) + and isinstance(queryset_model_arg, Instance) + and manager_model_arg.type_of_any == TypeOfAny.from_omitted_generics + ): + """ + Captures: + - CustomManager(BaseManager) with CustomQuerySet(QuerySet[CustomModel]) + - CustomManager(Manager) with CustomQuerySet(QuerySet[CustomModel]) + - BaseManager with CustomQuerySet(QuerySet[CustomModel]) + """ + return queryset_model_arg + elif ( + isinstance(manager_model_arg, Instance) + and isinstance(queryset_model_arg, AnyType) + and queryset_model_arg.type_of_any == TypeOfAny.from_omitted_generics + ): + """ + Captures: + - CustomManager(BaseManager[CustomModel]) with QuerySet + - CustomManager(Manager[CustomModel]) with QuerySet + """ + if manager_model_arg.type.fullname == fullnames.MODEL_CLASS_FULLNAME: + # Avoid populating default upper bound `django.db.models.base.Model` + return AnyType(TypeOfAny.from_omitted_generics) + return manager_model_arg + elif manager_model_arg is not None and queryset_model_arg is not None and manager_model_arg == queryset_model_arg: + """ + Captures: + - BaseManager with QuerySet + - Manager with QuerySet + - CustomManager(BaseManager) with CustomQuerySet(QuerySet) + - CustomManager(Manager) with CustomQuerySet(QuerySet) + """ + return queryset_model_arg + + return None + + def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext) -> None: """ Insert a new manager class node for a: ' = .from_queryset()'. @@ -183,8 +352,8 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte 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 isinstance(manager_node.node, TypeInfo): + manager_sym = semanal_api.lookup_current_scope(ctx.name) + if manager_sym and isinstance(manager_sym.node, TypeInfo): # This is just a deferral run where our work is already finished return @@ -192,103 +361,192 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte assert isinstance(callee, MemberExpr) assert isinstance(callee.expr, RefExpr) - base_manager_info = callee.expr.node - if base_manager_info is None: + manager_base = callee.expr.node + if manager_base is None: if not semanal_api.final_iteration: semanal_api.defer() return - assert isinstance(base_manager_info, TypeInfo) + assert isinstance(manager_base, TypeInfo) passed_queryset = ctx.call.args[0] - assert isinstance(passed_queryset, NameExpr) + assert isinstance(passed_queryset, NameExpr) # TODO: Add `MemberExpr`? - derived_queryset_fullname = passed_queryset.fullname - if derived_queryset_fullname is None: + if passed_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 + elif semanal_api.is_class_scope(): + # We force `.from_queryset` to be called _outside_ of a class body. So we'll + # skip doing any work if we're inside of one.. + return - 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 - ) - - sym = semanal_api.lookup_fully_qualified_or_none(derived_queryset_fullname) - assert sym is not None - if sym.node is None: + queryset_info: SymbolNode # We'll specialise this more a bit further down + class_name_param_value = None + queryset_sym = semanal_api.lookup_fully_qualified_or_none(passed_queryset.fullname) + assert queryset_sym is not None + if queryset_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 + elif isinstance(queryset_sym.node, TypeAlias): + assert isinstance(queryset_sym.node.target, Instance) + queryset_info = queryset_sym.node.target.type + # When we've encountered an alias, we take the queryset name from the alias. + # As e.g. for `django.db.models.QuerySet`, which is an alias, we'll otherwise + # get a name of `django.db.models._QuerySet` + class_name_param_value = manager_base.name + "From" + queryset_sym.node.name + else: + queryset_info = queryset_sym.node - derived_queryset_info = sym.node - assert isinstance(derived_queryset_info, TypeInfo) + assert isinstance(queryset_info, TypeInfo) + if len(ctx.call.args) > 1: + expr = ctx.call.args[1] + if isinstance(expr, StrExpr): + class_name_param_value = expr.value + # When the argument is not a `StrExpr` we're gonna have a hard time figuring + # out what the runtime value will be, so we'll just use the default instead. + + manager_class_name = class_name_param_value or manager_base.name + "From" + queryset_info.name + manager_base_instance = fill_typevars(manager_base) + assert isinstance(manager_base_instance, Instance) + # Create a new `TypeInfo` instance for the manager type + new_manager_info = semanal_api.basic_new_typeinfo( + name=manager_class_name, basetype_or_fallback=manager_base_instance, line=ctx.call.line + ) + new_manager_info.type_vars = manager_base.type_vars 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.type_vars = manager_base.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 - else: - custom_manager_generated_name = base_manager_info.name + "From" + derived_queryset_info.name + merge_queryset_into_manager( + api=semanal_api, + manager_info=new_manager_info, + queryset_info=queryset_info, + manager_base=manager_base, + class_name=manager_class_name, + ) - custom_manager_generated_fullname = ".".join(["django.db.models.manager", custom_manager_generated_name]) - base_manager_info.metadata.setdefault("from_queryset_managers", {}) - base_manager_info.metadata["from_queryset_managers"][custom_manager_generated_fullname] = new_manager_info.fullname + model_argument_type = intersect_manager_and_queryset_model_argument_type(manager_base, queryset_info) + if model_argument_type is None: + semanal_api.fail( + "Manager model argument does not match queryset model argument", + ctx.call, + code=errorcodes.MODEL_ARG_MISMATCH, + ) + model_argument_type = AnyType(TypeOfAny.from_error) - # 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) + symbol_kind = semanal_api.current_symbol_kind() + # Annotate the module variable + var = Var( + name=ctx.name, + type=TypeType(Instance(new_manager_info, [model_argument_type])), + ) + var.info = new_manager_info + var._fullname = f"{semanal_api.cur_mod_id}.{ctx.name}" + var.is_inferred = True + # Note: Order of `add_symbol_table_node` calls matter. Case being if + # `ctx.name == new_manager_info.name`, then we'd like the type to override the + # `Var`, so the `Var` won't exist in the end. + assert semanal_api.add_symbol_table_node(ctx.name, SymbolTableNode(symbol_kind, var, plugin_generated=True)) + # Insert the new manager dynamic class + assert semanal_api.add_symbol_table_node( + new_manager_info.name, SymbolTableNode(symbol_kind, new_manager_info, plugin_generated=True) + ) - 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) - # 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: - if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME: - break - for name, sym in class_mro_info.names.items(): - if not isinstance(sym.node, (FuncDef, Decorator)): - continue - # Insert the queryset method name as a class member. Note that the type of - # the method is set as Any. Figuring out the type is the job of the - # 'resolve_manager_method' attribute hook, which comes later. - # - # class BaseManagerFromMyQuerySet(BaseManager): - # queryset_method: Any = ... - # - helpers.add_new_sym_for_info( - new_manager_info, - name=name, - sym_type=AnyType(TypeOfAny.special_form), - ) +def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext) -> None: + """ + Insert a new manager class node for a - # 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: - helpers.add_new_sym_for_info( - new_manager_info, - name=name, - sym_type=AnyType(TypeOfAny.special_form), + ``` + = .as_manager() + ``` + """ + 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: + if not semanal_api.final_iteration: + semanal_api.defer() + return + + assert isinstance(manager_base, TypeInfo) + + callee = ctx.call.callee + assert isinstance(callee, MemberExpr) + assert isinstance(callee.expr, RefExpr) + + queryset_info = callee.expr.node + if queryset_info is None: + if not semanal_api.final_iteration: + semanal_api.defer() + return + + assert isinstance(queryset_info, TypeInfo) + + manager_class_name = 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 ( + existing_sym is not None + and isinstance(existing_sym.node, TypeInfo) + and existing_sym.node.metadata.get("django", {}).get("from_queryset_manager") == queryset_info.fullname + ): + # Reuse an identical, already generated, manager + new_manager_info = existing_sym.node + else: + manager_base_instance = fill_typevars(manager_base) + assert isinstance(manager_base_instance, Instance) + new_manager_info = helpers.add_new_class_for_module( + module=current_module, + name=manager_class_name, + bases=[manager_base_instance], + ) + new_manager_info.type_vars = manager_base.type_vars + new_manager_info.line = ctx.call.line + new_manager_info.defn.type_vars = manager_base.defn.type_vars + new_manager_info.defn.line = ctx.call.line + + merge_queryset_into_manager( + api=semanal_api, + manager_info=new_manager_info, + queryset_info=queryset_info, + manager_base=manager_base, + class_name=manager_class_name, ) - # Insert the new manager (dynamic) class - assert semanal_api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, new_manager_info, plugin_generated=True)) + # 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. + var = Var(name=ctx.name, type=Instance(new_manager_info, [AnyType(TypeOfAny.from_omitted_generics)])) + var.info = new_manager_info + var._fullname = f"{current_module.fullname}.{ctx.name}" + var.is_inferred = True + # Note: Order of `add_symbol_table_node` calls matters. Depending on what level + # we've found the `.as_manager()` call. Point here being that we want to replace the + # `.as_manager` return value with our newly created manager. + assert semanal_api.add_symbol_table_node( + ctx.name, SymbolTableNode(semanal_api.current_symbol_kind(), var, plugin_generated=True) + ) + assert 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, + # Note that the generated manager type is always inserted at module level + SymbolTableNode(GDEF, new_manager_info, plugin_generated=True), + ) def fail_if_manager_type_created_in_model_body(ctx: MethodContext) -> MypyType: diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 6b8e50276..7acbb1db6 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -9,17 +9,21 @@ from mypy.plugin import AnalyzeTypeContext, AttributeContext, CheckerPluginInterface, ClassDefContext from mypy.plugins import common from mypy.semanal import SemanticAnalyzer -from mypy.types import AnyType, Instance +from mypy.types import AnyType, Instance, PlaceholderType from mypy.types import Type as MypyType from mypy.types import TypedDictType, TypeOfAny from mypy_django_plugin.django.context import DjangoContext -from mypy_django_plugin.errorcodes import MANAGER_MISSING +from mypy_django_plugin.errorcodes import MANAGER_MISSING, MODEL_ARG_MISMATCH 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 ( + get_manager_model_argument_type, + intersect_manager_and_queryset_model_argument_type, +) class ModelClassInitializer: @@ -192,7 +196,11 @@ def has_any_parametrized_manager_as_base(self, info: TypeInfo) -> bool: return False def is_any_parametrized_manager(self, typ: Instance) -> bool: - return typ.type.fullname in fullnames.MANAGER_CLASSES and isinstance(typ.args[0], AnyType) + return ( + typ.type.fullname in fullnames.MANAGER_CLASSES + and isinstance(typ.args[0], AnyType) + and typ.args[0].type_of_any == TypeOfAny.from_omitted_generics + ) def create_new_model_parametrized_manager(self, name: str, base_manager_info: TypeInfo) -> Instance: bases = [] @@ -206,14 +214,14 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty 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) custom_manager_type = Instance(new_manager_info, [Instance(self.model_classdef.info, [])]) 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( - new_cls_def_context, + api=self.api, + cls=new_manager_info.defn, self_type=custom_manager_type, new_method_name=name, method_node=sym.node, @@ -231,6 +239,50 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty return custom_manager_type + def manager_is_dynamically_generated(self, manager_info: TypeInfo) -> bool: + return manager_info.metadata.get("django", {}).get("from_queryset_manager") is not None + + def get_manager_model_argument_type(self, manager_info: TypeInfo) -> Optional[MypyType]: + if self.manager_is_dynamically_generated(manager_info): + queryset_info = self.lookup_typeinfo(manager_info.metadata["django"]["from_queryset_manager"]) + # TODO: We should always be able to find the `TypeInfo` we created a + # dynamic manager class from, right? + assert queryset_info is not None + arg_type = intersect_manager_and_queryset_model_argument_type(manager_info, queryset_info) + else: + arg_type = get_manager_model_argument_type(manager_info) + + return arg_type + + def get_manager_type(self, manager_name: str, manager_info: TypeInfo) -> Optional[MypyType]: + manager_arg_type = self.get_manager_model_argument_type(manager_info) + if ( + isinstance(manager_arg_type, Instance) + and ( + manager_arg_type.type.fullname == self.model_classdef.info.fullname + # We accept overwriting declared upper bound of model argument type + or manager_arg_type.type.fullname == fullnames.MODEL_CLASS_FULLNAME + ) + ) or ( + isinstance(manager_arg_type, AnyType) and manager_arg_type.type_of_any == TypeOfAny.from_omitted_generics + ): + return Instance(manager_info, [Instance(self.model_classdef.info, [])]) + elif isinstance(manager_arg_type, PlaceholderType): + # Arg is expected to become a type, at some point, but for now we can't + # process it. + return None + else: + self.api.fail( + "Manager model argument does not match model it is assigned to", + ( + self.model_classdef.info.names[manager_name].node or self.model_classdef + if manager_name in self.model_classdef.info.names + else self.model_classdef + ), + code=MODEL_ARG_MISMATCH, + ) + return Instance(manager_info, [AnyType(TypeOfAny.from_error)]) + def run_with_model_cls(self, model_cls: Type[Model]) -> None: manager_info: Optional[TypeInfo] @@ -254,10 +306,14 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: if manager_info is None: 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: - manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) - self.add_new_node_to_model_class(manager_name, manager_type) + if manager_name not in self.model_classdef.info.names or self.manager_is_dynamically_generated( + manager_info + ): + manager_type = self.get_manager_type(manager_name, manager_info) + if manager_type is not None: + self.add_new_node_to_model_class(manager_name, manager_type) + else: + incomplete_manager_defs.add(manager_name) elif self.has_any_parametrized_manager_as_base(manager_info): # Ending up here could for instance be due to having a custom _Manager_ # that is not built from a custom QuerySet. Another example is a @@ -271,6 +327,13 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: continue self.add_new_node_to_model_class(manager_name, custom_manager_type) + else: + manager_type = self.get_manager_type(manager_name, manager_info) + if manager_type is not None: + # TODO: Can we avoid overwriting? Since `manager_name` is in type info names + self.add_new_node_to_model_class(manager_name, manager_type) + else: + incomplete_manager_defs.add(manager_name) 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 diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index af416c79f..ef114967f 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -680,18 +680,22 @@ # Note, that we cannot resolve dynamic calls for custom managers: class Transaction(models.Model): - objects = BaseManager.from_queryset(TransactionQuerySet) + objects = BaseManager.from_queryset(TransactionQuerySet)() def test(self) -> None: reveal_type(self.transactionlog_set) + reveal_type(Transaction.objects) # We use a fallback Any type: reveal_type(Transaction.objects.custom()) class TransactionLog(models.Model): transaction = models.ForeignKey(Transaction, on_delete=models.CASCADE) out: | + myapp/models:9: error: Could not resolve manager type for "myapp.models.Transaction.objects" 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:12: note: Revealed type is "django.db.models.manager.Manager[myapp.models.Transaction]" + myapp/models:14: error: "Manager[Transaction]" has no attribute "custom" + myapp/models:14: note: Revealed type is "Any" - case: resolve_primary_keys_for_foreign_keys_with_abstract_self_model diff --git a/tests/typecheck/managers/querysets/test_as_manager.yml b/tests/typecheck/managers/querysets/test_as_manager.yml new file mode 100644 index 000000000..2cebf26e8 --- /dev/null +++ b/tests/typecheck/managers/querysets/test_as_manager.yml @@ -0,0 +1,192 @@ +- case: declares_manager_type_like_django + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class MyQuerySet(models.QuerySet): + ... + + class MyModel(models.Model): + objects = MyQuerySet.as_manager() + +- case: includes_django_methods_returning_queryset + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects.none) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class MyQuerySet(models.QuerySet): + ... + + class MyModel(models.Model): + objects = MyQuerySet.as_manager() + +- case: model_gets_generated_manager_as_default_manager + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.queryset_method()) # N: Revealed type is "builtins.str" + reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class ModelQuerySet(models.QuerySet): + def queryset_method(self) -> str: + return 'hello' + + class MyModel(models.Model): + objects = ModelQuerySet.as_manager() + +- case: resolves_name_collision_with_other_module_level_object + main: | + from myapp.models import MyModel, ManagerFromModelQuerySet + reveal_type(ManagerFromModelQuerySet) # N: Revealed type is "builtins.int" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[myapp.models.MyModel]" + reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + ManagerFromModelQuerySet = 1 + + class ModelQuerySet(models.QuerySet): + ... + + class MyModel(models.Model): + objects = ModelQuerySet.as_manager() + +- case: includes_custom_queryset_methods + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects.custom_queryset_method()) # N: Revealed type is "myapp.models.ModelQuerySet" + reveal_type(MyModel.objects.all().custom_queryset_method()) # N: Revealed type is "myapp.models.ModelQuerySet" + reveal_type(MyModel.objects.returns_int_sequence()) # N: Revealed type is "typing.Sequence[builtins.int]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from typing import Sequence + + class ModelQuerySet(models.QuerySet["MyModel"]): + def custom_queryset_method(self) -> "ModelQuerySet": + return self.all() + + def returns_int_sequence(self) -> Sequence[int]: + return [1] + + class MyModel(models.Model): + objects = ModelQuerySet.as_manager() + +- case: handles_call_outside_of_model_class_definition + main: | + from myapp.models import MyModel, MyModelManager + reveal_type(MyModelManager) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[Any]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class ModelQuerySet(models.QuerySet["MyModel"]): + ... + + MyModelManager = ModelQuerySet.as_manager() + class MyModel(models.Model): + objects = MyModelManager + +- case: handles_name_collision_when_declared_outside_of_model_class_body + main: | + from myapp.models import MyModel, ManagerFromModelQuerySet + reveal_type(ManagerFromModelQuerySet) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[Any]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[myapp.models.MyModel]" + reveal_type(MyModel.objects.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class ModelQuerySet(models.QuerySet["MyModel"]): + ... + + ManagerFromModelQuerySet = ModelQuerySet.as_manager() + class MyModel(models.Model): + objects = ManagerFromModelQuerySet + +- case: reuses_generated_type_when_called_identically_for_multiple_managers + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects_1) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects_2) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects_1.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects_2.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class ModelQuerySet(models.QuerySet["MyModel"]): + ... + + class MyModel(models.Model): + objects_1 = ModelQuerySet.as_manager() + objects_2 = ModelQuerySet.as_manager() + +- case: generates_new_manager_class_when_name_colliding_with_explicit_manager + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[myapp.models.MyModel]" + reveal_type(MyModel.objects.custom_method()) # N: Revealed type is "builtins.int" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class ManagerFromModelQuerySet(models.Manager): + ... + + class ModelQuerySet(models.QuerySet["MyModel"]): + def custom_method(self) -> int: + return 1 + + class MyModel(models.Model): + objects = ModelQuerySet.as_manager() diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index 2a2ae5db4..7c532ef8e 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -1,7 +1,7 @@ - case: from_queryset_with_base_manager main: | from myapp.models import MyModel - reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" + reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel().objects.get()) # N: Revealed type is "myapp.models.MyModel" reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is "builtins.str" reveal_type(MyModel.objects.filter(id=1).queryset_method()) # N: Revealed type is "builtins.str" @@ -25,8 +25,8 @@ - case: from_queryset_queryset_imported_from_other_module main: | from myapp.models import MyModel - reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" - reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel.objects.get()) # N: Revealed type is "myapp.models.MyModel" reveal_type(MyModel.objects.queryset_method()) # N: Revealed type is "myapp.querysets.ModelQuerySet" reveal_type(MyModel.objects.queryset_method_2()) # N: Revealed type is "typing.Iterable[myapp.querysets.Custom]" @@ -74,7 +74,7 @@ - case: from_queryset_generated_manager_imported_from_other_module main: | from myapp.models import MyModel - reveal_type(MyModel.objects) # N: Revealed type is "myapp.querysets.NewManager[myapp.models.MyModel]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.querysets.BaseManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel.objects.get()) # N: Revealed type is "myapp.models.MyModel" reveal_type(MyModel.objects.queryset_method()) # N: Revealed type is "myapp.querysets.ModelQuerySet" reveal_type(MyModel.objects.queryset_method_2()) # N: Revealed type is "typing.Iterable[myapp.querysets.Custom]" @@ -120,10 +120,27 @@ class MyModel(models.Model): objects = NewManager() +- case: from_queryset_annotates_manager_variable_as_type + main: | + from myapp.models import NewManager + reveal_type(NewManager) # N: Revealed type is "Type[myapp.models.ManagerFromModelQuerySet[Any]]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class ModelQuerySet(models.QuerySet): + ... + + NewManager = models.Manager.from_queryset(ModelQuerySet) + - case: from_queryset_with_manager main: | from myapp.models import MyModel - reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" + reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel().objects.get()) # N: Revealed type is "myapp.models.MyModel" reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is "builtins.str" installed_apps: @@ -145,8 +162,8 @@ - case: from_queryset_returns_intersection_of_manager_and_queryset main: | from myapp.models import MyModel, NewManager - reveal_type(NewManager()) # N: Revealed type is "myapp.models.NewManager" - reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" + reveal_type(NewManager()) # N: Revealed type is "myapp.models.ModelBaseManagerFromModelQuerySet" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ModelBaseManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel.objects.get()) # N: Revealed type is "Any" reveal_type(MyModel.objects.manager_only_method()) # N: Revealed type is "builtins.int" reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is "builtins.str" @@ -170,12 +187,16 @@ - case: from_queryset_with_class_name_provided main: | - from myapp.models import MyModel, NewManager + from myapp.models import MyModel, NewManager, OtherModel, OtherManager reveal_type(NewManager()) # N: Revealed type is "myapp.models.NewManager" reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" reveal_type(MyModel.objects.get()) # N: Revealed type is "Any" reveal_type(MyModel.objects.manager_only_method()) # N: Revealed type is "builtins.int" reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is "builtins.str" + reveal_type(OtherManager()) # N: Revealed type is "myapp.models.X" + reveal_type(OtherModel.objects) # N: Revealed type is "myapp.models.X[myapp.models.OtherModel]" + reveal_type(OtherModel.objects.manager_only_method()) # N: Revealed type is "builtins.int" + reveal_type(OtherModel.objects.manager_and_queryset_method()) # N: Revealed type is "builtins.str" installed_apps: - myapp files: @@ -194,10 +215,14 @@ class MyModel(models.Model): objects = NewManager() + OtherManager = ModelBaseManager.from_queryset(ModelQuerySet, class_name='X') + class OtherModel(models.Model): + objects = OtherManager() + - case: from_queryset_with_class_inheritance main: | from myapp.models import MyModel - reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" + reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel().objects.get()) # N: Revealed type is "myapp.models.MyModel" reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is "builtins.str" installed_apps: @@ -221,7 +246,7 @@ - case: from_queryset_with_manager_in_another_directory_and_imports main: | from myapp.models import MyModel - reveal_type(MyModel().objects) # N: Revealed type is "myapp.managers.NewManager[myapp.models.MyModel]" + reveal_type(MyModel().objects) # N: Revealed type is "myapp.managers.ManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel().objects.get()) # N: Revealed type is "myapp.models.MyModel" reveal_type(MyModel().objects.queryset_method) # N: Revealed type is "def (param: Union[builtins.str, None] =) -> Union[builtins.str, None]" reveal_type(MyModel().objects.queryset_method('str')) # N: Revealed type is "Union[builtins.str, None]" @@ -251,7 +276,7 @@ disable_cache: true main: | from myapp.models import MyModel - reveal_type(MyModel().objects) # N: Revealed type is "myapp.managers.NewManager[myapp.models.MyModel]" + reveal_type(MyModel().objects) # N: Revealed type is "myapp.managers.ManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel().objects.get()) # N: Revealed type is "myapp.models.MyModel" reveal_type(MyModel().objects.base_queryset_method) # N: Revealed type is "def (param: Union[builtins.int, builtins.str]) -> " reveal_type(MyModel().objects.base_queryset_method(2)) # N: Revealed type is "" @@ -283,7 +308,7 @@ - case: from_queryset_with_decorated_queryset_methods main: | from myapp.models import MyModel - reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" + reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is "builtins.str" reveal_type(MyModel.objects.queryset_method_2()) # N: Revealed type is "builtins.int" installed_apps: @@ -311,9 +336,9 @@ - case: from_queryset_model_gets_generated_manager_as_default_manager main: | from myapp.models import MyModel - reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel.objects.queryset_method()) # N: Revealed type is "builtins.str" - reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" + reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" installed_apps: - myapp files: @@ -333,7 +358,7 @@ - case: from_queryset_can_resolve_explicit_any_methods main: | from myapp.models import MyModel - reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.MyManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel.objects.queryset_method(1)) # N: Revealed type is "Any" reveal_type(MyModel.objects.queryset_method) # N: Revealed type is "def (qarg: Any) -> Any" reveal_type(MyModel.objects.manager_method(2)) # N: Revealed type is "Any" @@ -435,7 +460,6 @@ class MyModel(models.Model): objects = MyManager() - # This tests a regression where mypy would generate phantom warnings about # undefined types due to unresolved types when copying methods from QuerySet to # a manager dynamically created using Manager.from_queryset(). @@ -475,3 +499,314 @@ class UserQuerySet(models.QuerySet["User"]): pass + +- case: reuses_type_when_called_twice_identically + main: | + from myapp.models import MyModel, FirstManager, SecondManager + reveal_type(FirstManager) # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(SecondManager) # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.first) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.second) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class ModelQuerySet(models.QuerySet["MyModel"]): + ... + + FirstManager = BaseManager.from_queryset(ModelQuerySet) + SecondManager = BaseManager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + first = FirstManager() + second = SecondManager() + +- case: handles_name_collision_with_generated_type + main: | + from myapp.models import MyModel, BaseManagerFromModelQuerySet + reveal_type(BaseManagerFromModelQuerySet()) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class ModelQuerySet(models.QuerySet["MyModel"]): + ... + + BaseManagerFromModelQuerySet = BaseManager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + objects = BaseManagerFromModelQuerySet() + +- case: accepts_explicit_none_as_class_name + main: | + from myapp.models import PositionalNone, NoneAsKwarg + reveal_type(PositionalNone) # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[Any]]" + reveal_type(NoneAsKwarg) # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[Any]]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class ModelQuerySet(models.QuerySet): + ... + + PositionalNone = BaseManager.from_queryset(ModelQuerySet, None) + NoneAsKwarg = BaseManager.from_queryset(ModelQuerySet, class_name=None) + +- case: uses_fallback_class_name_when_argument_is_not_string_expression + main: | + from myapp.models import StrCallable + reveal_type(StrCallable) # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[Any]]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class ModelQuerySet(models.QuerySet): + ... + + StrCallable = BaseManager.from_queryset(ModelQuerySet, class_name=str(1)) + +- case: emits_error_when_mismatching_manager_and_queryset_model_args + main: | + from myapp.models import FromBaseManager, FromManager, MyModel + reveal_type(FromBaseManager) # N: Revealed type is "Type[myapp.models.BaseModelManagerFromModelQuerySet[Any]]" + reveal_type(MyModel.base) # N: Revealed type is "myapp.models.BaseModelManagerFromModelQuerySet[Any]" + reveal_type(FromManager) # N: Revealed type is "Type[myapp.models.ManagerModelManagerFromModelQuerySet[Any]]" + reveal_type(MyModel.manager) # N: Revealed type is "myapp.models.ManagerModelManagerFromModelQuerySet[Any]" + out: | + myapp/models:13: error: Manager model argument does not match queryset model argument + myapp/models:14: error: Manager model argument does not match queryset model argument + myapp/models:16: error: Manager model argument does not match model it is assigned to + myapp/models:17: error: Manager model argument does not match model it is assigned to + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class ModelQuerySet(models.QuerySet["MyModel"]): + ... + + class BaseModelManager(BaseManager["OtherModel"]): + ... + + class ManagerModelManager(models.Manager["OtherModel"]): + ... + + FromBaseManager = BaseModelManager.from_queryset(ModelQuerySet) + FromManager = ManagerModelManager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + base = FromBaseManager() + manager = FromManager() + + class OtherModel(models.Model): + ... + +- case: populates_model_arg_from_manager_when_missing_on_queryset + main: | + from myapp.models import FromBaseManager, FromManager, MyModel + reveal_type(FromBaseManager) # N: Revealed type is "Type[myapp.models.BaseModelManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.base) # N: Revealed type is "myapp.models.BaseModelManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(FromManager) # N: Revealed type is "Type[myapp.models.ManagerModelManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.manager) # N: Revealed type is "myapp.models.ManagerModelManagerFromModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class ModelQuerySet(models.QuerySet): + ... + + class BaseModelManager(BaseManager["MyModel"]): + ... + + class ManagerModelManager(models.Manager["MyModel"]): + ... + + FromBaseManager = BaseModelManager.from_queryset(ModelQuerySet) + FromManager = ManagerModelManager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + base = FromBaseManager() + manager = FromManager() + +- case: populates_model_arg_from_queryset_when_missing_on_manager + main: | + from myapp.models import FromBaseManager, FromManager, MyModel + reveal_type(FromBaseManager) # N: Revealed type is "Type[myapp.models.BaseModelManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.base) # N: Revealed type is "myapp.models.BaseModelManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(FromManager) # N: Revealed type is "Type[myapp.models.ManagerModelManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.manager) # N: Revealed type is "myapp.models.ManagerModelManagerFromModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class ModelQuerySet(models.QuerySet["MyModel"]): + ... + + class BaseModelManager(BaseManager): + ... + + class ManagerModelManager(models.Manager): + ... + + FromBaseManager = BaseModelManager.from_queryset(ModelQuerySet) + FromManager = ManagerModelManager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + base = FromBaseManager() + manager = FromManager() + +- case: populates_model_arg_from_queryset_when_manager_base_is_django_builtin_manager + main: | + from myapp.models import FromBaseManager, FromManager, MyModel + reveal_type(FromBaseManager) # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.base) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(FromManager) # N: Revealed type is "Type[myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.manager) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class ModelQuerySet(models.QuerySet["MyModel"]): + ... + + FromBaseManager = BaseManager.from_queryset(ModelQuerySet) + FromManager = models.Manager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + base = FromBaseManager() + manager = FromManager() + +- case: populates_model_arg_from_manager_when_queryset_is_django_builtin + main: | + from myapp.models import FromBaseManager, FromManager, MyModel + reveal_type(FromBaseManager) # N: Revealed type is "Type[myapp.models.BaseModelManagerFromQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.base) # N: Revealed type is "myapp.models.BaseModelManagerFromQuerySet[myapp.models.MyModel]" + reveal_type(FromManager) # N: Revealed type is "Type[myapp.models.ManagerModelManagerFromQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.manager) # N: Revealed type is "myapp.models.ManagerModelManagerFromQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models import QuerySet + from django.db.models.manager import BaseManager + + class BaseModelManager(BaseManager["MyModel"]): + ... + + class ManagerModelManager(models.Manager["MyModel"]): + ... + + FromBaseManager = BaseModelManager.from_queryset(QuerySet) + FromManager = ManagerModelManager.from_queryset(QuerySet) + class MyModel(models.Model): + base = FromBaseManager() + manager = FromManager() + +- case: populates_model_arg_as_any_when_no_args_specified + main: | + from myapp.models import FromCustomBaseManager, FromCustomManager, FromBaseManager, FromManager, MyModel + reveal_type(FromCustomBaseManager) # N: Revealed type is "Type[myapp.models.BaseModelManagerFromModelQuerySet[Any]]" + reveal_type(MyModel.custom_base) # N: Revealed type is "myapp.models.BaseModelManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(FromCustomManager) # N: Revealed type is "Type[myapp.models.ManagerModelManagerFromModelQuerySet[Any]]" + reveal_type(MyModel.custom_manager) # N: Revealed type is "myapp.models.ManagerModelManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(FromBaseManager) # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[Any]]" + reveal_type(MyModel.base) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(FromManager) # N: Revealed type is "Type[myapp.models.ManagerFromModelQuerySet[Any]]" + reveal_type(MyModel.manager) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models import QuerySet + from django.db.models.manager import BaseManager + + class ModelQuerySet(models.QuerySet): + ... + + class BaseModelManager(BaseManager): + ... + + class ManagerModelManager(models.Manager): + ... + + FromCustomBaseManager = BaseModelManager.from_queryset(ModelQuerySet) + FromCustomManager = ManagerModelManager.from_queryset(ModelQuerySet) + FromBaseManager = BaseManager.from_queryset(ModelQuerySet) + FromManager = models.Manager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + custom_base = FromCustomBaseManager() + custom_manager = FromCustomManager() + base = FromBaseManager() + manager = FromManager() + +- case: populates_model_arg_when_same_for_manager_and_queryset + main: | + from myapp.models import FromBaseManager, FromManager, MyModel + reveal_type(FromBaseManager) # N: Revealed type is "Type[myapp.models.BaseModelManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.base) # N: Revealed type is "myapp.models.BaseModelManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(FromManager) # N: Revealed type is "Type[myapp.models.ManagerModelManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.manager) # N: Revealed type is "myapp.models.ManagerModelManagerFromModelQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models import QuerySet + from django.db.models.manager import BaseManager + + class ModelQuerySet(models.QuerySet["MyModel"]): + ... + + class BaseModelManager(BaseManager["MyModel"]): + ... + + class ManagerModelManager(models.Manager["MyModel"]): + ... + + FromBaseManager = BaseModelManager.from_queryset(ModelQuerySet) + FromManager = ManagerModelManager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + base = FromBaseManager() + manager = FromManager() diff --git a/tests/typecheck/managers/querysets/test_union_type.yml b/tests/typecheck/managers/querysets/test_union_type.yml index 7289603ec..421343a8a 100644 --- a/tests/typecheck/managers/querysets/test_union_type.yml +++ b/tests/typecheck/managers/querysets/test_union_type.yml @@ -9,7 +9,7 @@ model_cls = type(instance) reveal_type(model_cls) # N: Revealed type is "Union[Type[myapp.models.Order], Type[myapp.models.User]]" - reveal_type(model_cls.objects) # N: Revealed type is "Union[myapp.models.OrderManager[myapp.models.Order], myapp.models.UserManager[myapp.models.User]]" + reveal_type(model_cls.objects) # N: Revealed type is "Union[myapp.models.ManagerFromMyQuerySet[myapp.models.Order], myapp.models.ManagerFromMyQuerySet[myapp.models.User]]" model_cls.objects.my_method() # E: Unable to resolve return type of queryset/manager method "my_method" installed_apps: - myapp diff --git a/tests/typecheck/managers/test_managers.yml b/tests/typecheck/managers/test_managers.yml index 99778f82d..331abbb27 100644 --- a/tests/typecheck/managers/test_managers.yml +++ b/tests/typecheck/managers/test_managers.yml @@ -521,3 +521,90 @@ myapp/models:47: note: Revealed type is "django.db.models.manager.BaseManager[myapp.models.InvisibleUnresolvable]" myapp/models:49: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.Booking]" myapp/models:50: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.Booking]" + +- case: emits_error_when_manager_model_arg_does_not_match_model + main: | + from myapp.models import MyModel + reveal_type(MyModel.from_manager) # N: Revealed type is "myapp.models.FromManager[Any]" + reveal_type(MyModel.from_base_manager) # N: Revealed type is "myapp.models.FromBaseManager[Any]" + out: | + myapp/models:11: error: Manager model argument does not match model it is assigned to + myapp/models:12: error: Manager model argument does not match model it is assigned to + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class FromManager(models.Manager["OtherModel"]): + ... + + class FromBaseManager(BaseManager["OtherModel"]): + ... + + class MyModel(models.Model): + from_manager = FromManager() + from_base_manager = FromBaseManager() + + class OtherModel(models.Model): + ... + +- case: emits_error_when_imported_manager_model_arg_does_not_match_model + main: | + from myapp.models import MyModel + reveal_type(MyModel.from_manager) # N: Revealed type is "myapp.managers.FromManager[Any]" + reveal_type(MyModel.from_base_manager) # N: Revealed type is "myapp.managers.FromBaseManager[Any]" + out: | + myapp/models:5: error: Manager model argument does not match model it is assigned to + myapp/models:6: error: Manager model argument does not match model it is assigned to + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from .managers import FromManager, FromBaseManager + + class MyModel(models.Model): + from_manager = FromManager() + from_base_manager = FromBaseManager() + + class OtherModel(models.Model): + ... + - path: myapp/managers.py + content: | + from typing import TYPE_CHECKING + from django.db import models + from django.db.models.manager import BaseManager + + if TYPE_CHECKING: + from .models import OtherModel + + class FromManager(models.Manager["OtherModel"]): + ... + + class FromBaseManager(BaseManager["OtherModel"]): + ... + +- case: populates_manager_model_arg_for_django_builtin_managers + main: | + from myapp.models import MyModel + reveal_type(MyModel.from_manager) # N: Revealed type is "django.db.models.manager.Manager[myapp.models.MyModel]" + reveal_type(MyModel.from_base_manager) # N: Revealed type is "django.db.models.manager.BaseManager[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + from django.db.models.manager import BaseManager + + class MyModel(models.Model): + from_manager = models.Manager() + # TODO: Why do get complaint that we "Need type annotation ..."? + from_base_manager: BaseManager = BaseManager()