From ad718345b6cc23636eac32aa40c2821c79ec1a9c Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Sun, 26 Jun 2022 22:16:42 +0200 Subject: [PATCH] Implement support for `.as_manager()` --- README.md | 19 +- mypy_django_plugin/lib/helpers.py | 24 +- mypy_django_plugin/main.py | 8 +- mypy_django_plugin/transformers/managers.py | 291 +++++++++++++----- mypy_django_plugin/transformers/models.py | 4 +- .../managers/querysets/test_as_manager.yml | 192 ++++++++++++ .../managers/querysets/test_from_queryset.yml | 102 +++++- .../managers/querysets/test_union_type.yml | 2 +- 8 files changed, 524 insertions(+), 118 deletions(-) create mode 100644 tests/typecheck/managers/querysets/test_as_manager.yml diff --git a/README.md b/README.md index a392f80ab8..32a469ba1b 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/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index 1ce73fd8ba..c8a9f48dc2 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 6d3550d627..058b95a04f 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,16 @@ 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"): + if fullname.endswith(".from_queryset"): class_name, _, _ = fullname.rpartition(".") 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 fullname.endswith(".as_manager"): + class_name, _, _ = fullname.rpartition(".") + 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 28ac774188..6b1906604b 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -16,10 +16,10 @@ TypeInfo, Var, ) -from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext -from mypy.types import AnyType, CallableType, Instance, ProperType +from mypy.plugin import AttributeContext, DynamicClassDefContext, MethodContext +from mypy.semanal import SemanticAnalyzer +from mypy.types import AnyType, CallableType, Instance, ProperType, TypeOfAny, TypeType from mypy.types import Type as MypyType -from mypy.types import TypeOfAny from typing_extensions import Final from mypy_django_plugin import errorcodes @@ -175,6 +175,67 @@ 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.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: + helpers.add_new_sym_for_info( + manager_info, + name=name, + sym_type=AnyType(TypeOfAny.special_form), + ) + + def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext) -> None: """ Insert a new manager class node for a: ' = .from_queryset()'. @@ -183,8 +244,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 +253,183 @@ 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) - 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_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 - 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} + queryset_info = queryset_sym.node + assert isinstance(queryset_info, TypeInfo) if len(ctx.call.args) > 1: expr = ctx.call.args[1] assert isinstance(expr, StrExpr) - custom_manager_generated_name = expr.value + manager_class_name = expr.value else: - custom_manager_generated_name = base_manager_info.name + "From" + derived_queryset_info.name + manager_class_name = manager_base.name + "From" + queryset_info.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 + 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.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() - # 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) + try: + 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, + ) + except helpers.IncompleteDefnException: + if not semanal_api.final_iteration: + semanal_api.defer() + return - 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) + symbol_kind = semanal_api.current_symbol_kind() + # Annotate the module variable as `: Type[]` as the model + # type won't be defined on variable level. + var = Var( + name=ctx.name, + type=TypeType(Instance(new_manager_info, [AnyType(TypeOfAny.from_omitted_generics)])), + ) + 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 _only_ like the type and not the + # `Var` to exist.. + 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) + ) - # 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), - ) - # 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), - ) +def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext) -> None: + """ + Insert a new manager class node for a + + ``` + = .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 - # Insert the new manager (dynamic) class - assert semanal_api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, new_manager_info, plugin_generated=True)) + 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 + + try: + 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, + ) + except helpers.IncompleteDefnException: + if not semanal_api.final_iteration: + semanal_api.defer() + return + + # 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 6b8e502762..abaed7d743 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -206,14 +206,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, 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 0000000000..2cebf26e84 --- /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 2a2ae5db42..59df446469 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,49 @@ 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[Any]]" + reveal_type(SecondManager) # N: Revealed type is "Type[myapp.models.BaseManagerFromModelQuerySet[Any]]" + 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() diff --git a/tests/typecheck/managers/querysets/test_union_type.yml b/tests/typecheck/managers/querysets/test_union_type.yml index 7289603ec5..421343a8aa 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