diff --git a/README.md b/README.md index 51582f054..596ef3074 100644 --- a/README.md +++ b/README.md @@ -145,9 +145,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: @@ -163,12 +161,12 @@ class MyModel(models.Model): 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: @@ -188,9 +186,9 @@ class MyModel(models.Model): objects = MyModelManager() -def use_my_model(): - foo = MyModel.objects.get(id=1) - return foo.xyz # Gives an error +def use_my_model() -> int: + foo = MyModel.objects.get(id=1) # Should now be `MyModel` + return foo.xyz # Gives an error ``` ### How do I annotate cases where I called QuerySet.annotate? diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index b78fad401..b8e171994 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/mypy_django_plugin/lib/helpers.py @@ -387,29 +387,27 @@ 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, ) -> bool: - semanal_api = get_semanal_api(ctx) if method_node.type is None: 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 True 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 False 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 False @@ -422,7 +420,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 False if arg_name is None and hasattr(method_node, "arguments"): @@ -438,9 +436,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) return True diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index dde0b20d5..772c889b3 100644 --- a/mypy_django_plugin/main.py +++ b/mypy_django_plugin/main.py @@ -24,6 +24,7 @@ from mypy_django_plugin.transformers import fields, forms, init_create, meta, querysets, request, settings from mypy_django_plugin.transformers.functional import resolve_str_promise_attribute from mypy_django_plugin.transformers.managers import ( + create_new_manager_class_from_as_manager_method, create_new_manager_class_from_from_queryset_method, resolve_manager_method, ) @@ -301,11 +302,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 39ee933cc..28eae8639 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -15,7 +15,8 @@ TypeInfo, Var, ) -from mypy.plugin import AttributeContext, DynamicClassDefContext, SemanticAnalyzerPluginInterface +from mypy.plugin import AttributeContext, DynamicClassDefContext +from mypy.semanal import SemanticAnalyzer from mypy.semanal_shared import has_placeholder from mypy.types import AnyType, CallableType, Instance, ProperType from mypy.types import Type as MypyType @@ -150,7 +151,6 @@ def get_method_type_from_reverse_manager( def resolve_manager_method_from_instance(instance: Instance, method_name: str, ctx: AttributeContext) -> MypyType: - api = helpers.get_typechecker_api(ctx) method_type = get_method_type_from_dynamic_manager( api, method_name, instance @@ -164,9 +164,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. @@ -197,12 +199,12 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte return # Don't redeclare the manager class if we've already defined it. - manager_node = semanal_api.lookup_current_scope(ctx.name) - if manager_node and isinstance(manager_node.node, TypeInfo): + 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 - new_manager_info = create_manager_info_from_from_queryset_call(ctx.api, ctx.call, ctx.name) + new_manager_info = create_manager_info_from_from_queryset_call(semanal_api, ctx.call, ctx.name) if new_manager_info is None: if not ctx.api.final_iteration: ctx.api.defer() @@ -212,8 +214,17 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte helpers.add_new_manager_base(semanal_api, new_manager_info.fullname) +def register_dynamically_created_manager(fullname: str, manager_name: str, manager_base: TypeInfo) -> None: + 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 we put new type(s) in the + # module currently being processed, so we'll map those together through metadata. + runtime_fullname = ".".join(["django.db.models.manager", manager_name]) + manager_base.metadata["from_queryset_managers"][runtime_fullname] = fullname + + def create_manager_info_from_from_queryset_call( - api: SemanticAnalyzerPluginInterface, call_expr: CallExpr, name: Optional[str] = None + api: SemanticAnalyzer, call_expr: CallExpr, name: Optional[str] = None ) -> Optional[TypeInfo]: """ Extract manager and queryset TypeInfo from a from_queryset call. @@ -247,30 +258,48 @@ def create_manager_info_from_from_queryset_call( else: manager_name = f"{base_manager_info.name}From{queryset_info.name}" - try: - new_manager_info = create_manager_class(api, base_manager_info, name or manager_name, call_expr.line) - except helpers.IncompleteDefnException: - return None - - popuplate_manager_from_queryset(new_manager_info, queryset_info) - - manager_fullname = ".".join(["django.db.models.manager", manager_name]) - - base_manager_info = new_manager_info.mro[1] - base_manager_info.metadata.setdefault("from_queryset_managers", {}) - base_manager_info.metadata["from_queryset_managers"][manager_fullname] = new_manager_info.fullname + # Always look in global scope, as that's where we'll declare dynamic manager classes + manager_sym = api.globals.get(manager_name) + if ( + manager_sym is not None + and isinstance(manager_sym.node, TypeInfo) + and manager_sym.node.has_base(base_manager_info.fullname) + and manager_sym.node.metadata.get("django", {}).get("from_queryset_manager") == queryset_info.fullname + ): + # Reuse an identical, already generated, manager + new_manager_info = manager_sym.node + else: + # Create a new `TypeInfo` instance for the manager type + try: + new_manager_info = create_manager_class( + api=api, + base_manager_info=base_manager_info, + name=manager_name, + line=call_expr.line, + with_unique_name=name is not None and name != manager_name, + ) + except helpers.IncompleteDefnException: + return None + + populate_manager_from_queryset(new_manager_info, queryset_info) + register_dynamically_created_manager( + fullname=new_manager_info.fullname, + manager_name=manager_name, + manager_base=base_manager_info, + ) # Add the new manager to the current module module = api.modules[api.cur_mod_id] - module.names[name or manager_name] = SymbolTableNode( - GDEF, new_manager_info, plugin_generated=True, no_serialize=False - ) + if name is not None and name != new_manager_info.name: + # Unless names are equal, there's 2 symbol names that needs the manager info + module.names[name] = SymbolTableNode(GDEF, new_manager_info, plugin_generated=True) + module.names[new_manager_info.name] = SymbolTableNode(GDEF, new_manager_info, plugin_generated=True) return new_manager_info def create_manager_class( - api: SemanticAnalyzerPluginInterface, base_manager_info: TypeInfo, name: str, line: int + api: SemanticAnalyzer, base_manager_info: TypeInfo, name: str, line: int, with_unique_name: bool ) -> TypeInfo: base_manager_instance = fill_typevars(base_manager_info) @@ -280,17 +309,24 @@ def create_manager_class( if any(has_placeholder(type_var) for type_var in base_manager_info.defn.type_vars): raise helpers.IncompleteDefnException - manager_info = helpers.create_type_info(name, api.cur_mod_id, bases=[base_manager_instance]) + if with_unique_name: + manager_info = helpers.add_new_class_for_module( + module=api.modules[api.cur_mod_id], + name=name, + bases=[base_manager_instance], + ) + else: + manager_info = helpers.create_type_info(name, api.cur_mod_id, bases=[base_manager_instance]) + manager_info.line = line manager_info.type_vars = base_manager_info.type_vars manager_info.defn.type_vars = base_manager_info.defn.type_vars manager_info.defn.line = line - manager_info.metaclass_type = manager_info.calculate_metaclass_type() return manager_info -def popuplate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeInfo) -> None: +def populate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeInfo) -> None: """ Add methods from the QuerySet class to the manager. """ @@ -318,7 +354,7 @@ def popuplate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeI helpers.add_new_sym_for_info( manager_info, name=name, - sym_type=AnyType(TypeOfAny.special_form), + sym_type=AnyType(TypeOfAny.implementation_artifact), ) # For methods on BaseManager that return a queryset we need to update @@ -330,5 +366,103 @@ def popuplate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeI helpers.add_new_sym_for_info( manager_info, name=method_name, - sym_type=AnyType(TypeOfAny.special_form), + sym_type=AnyType(TypeOfAny.implementation_artifact), ) + + +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 + + 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.has_base(fullnames.MANAGER_CLASS_FULLNAME) + 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: + # Create a new `TypeInfo` instance for the manager type + try: + new_manager_info = create_manager_class( + api=semanal_api, + base_manager_info=manager_base, + name=manager_class_name, + line=ctx.call.line, + with_unique_name=True, + ) + except helpers.IncompleteDefnException: + if not semanal_api.final_iteration: + semanal_api.defer() + return + + populate_manager_from_queryset(new_manager_info, queryset_info) + register_dynamically_created_manager( + fullname=new_manager_info.fullname, + manager_name=manager_class_name, + manager_base=manager_base, + ) + + # 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) + + # 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) + ) + # Add the new manager to the current module + 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), + ) diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index e6d325010..a4732a325 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -126,7 +126,9 @@ def get_or_create_manager_with_any_fallback(self, related_manager: bool = False) # class. The actual type of these methods are resolved in # resolve_manager_method. for method_name in MANAGER_METHODS_RETURNING_QUERYSET: - helpers.add_new_sym_for_info(manager_info, name=method_name, sym_type=AnyType(TypeOfAny.special_form)) + helpers.add_new_sym_for_info( + manager_info, name=method_name, sym_type=AnyType(TypeOfAny.implementation_artifact) + ) manager_info.metadata["django"] = { "any_fallback_manager": True, @@ -289,14 +291,14 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty # but rather waiting until we know we won't defer new_manager_info = self.add_new_class_for_current_module(name, bases) # copy fields to a new manager - new_cls_def_context = ClassDefContext(cls=new_manager_info.defn, reason=self.ctx.reason, api=self.api) 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): copied_method = 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, @@ -316,37 +318,57 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty return custom_manager_type + def lookup_manager(self, fullname: str, manager: "Manager[Any]") -> Optional[TypeInfo]: + manager_info = self.lookup_typeinfo(fullname) + if manager_info is None: + manager_info = self.get_dynamic_manager(fullname, manager) + return manager_info + + def is_manager_dynamically_generated(self, manager_info: Optional[TypeInfo]) -> bool: + if manager_info is None: + return False + return manager_info.metadata.get("django", {}).get("from_queryset_manager") is not None + + def reparametrize_dynamically_created_manager(self, manager_name: str, manager_info: Optional[TypeInfo]) -> None: + if not self.is_manager_dynamically_generated(manager_info): + return + + assert manager_info is not None + # Reparameterize dynamically created manager with model type + manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) + self.add_new_node_to_model_class(manager_name, manager_type) + def run_with_model_cls(self, model_cls: Type[Model]) -> None: manager_info: Optional[TypeInfo] incomplete_manager_defs = set() for manager_name, manager in model_cls._meta.managers_map.items(): - # If the manager is already typed do nothing manager_node = self.model_classdef.info.names.get(manager_name, None) - if manager_node and manager_node.type is not None: - continue - - manager_class_name = manager.__class__.__name__ manager_fullname = helpers.get_class_fullname(manager.__class__) + manager_info = self.lookup_manager(manager_fullname, manager) - manager_info = self.lookup_typeinfo(manager_fullname) - if manager_info is None: + if manager_node and manager_node.type is not None: + # Manager is already typed -> do nothing unless it's a dynamically generated manager + self.reparametrize_dynamically_created_manager(manager_name, manager_info) + continue + elif manager_info is None: + # We couldn't find a manager type, see if we should create one manager_info = self.create_manager_from_from_queryset(manager_name) - if manager_info is None: - manager_info = self.get_dynamic_manager(manager_fullname, manager) if manager_info is None: incomplete_manager_defs.add(manager_name) continue - is_dynamically_generated = manager_info.metadata.get("django", {}).get("from_queryset_manager") is not None - if manager_name not in self.model_classdef.info.names or is_dynamically_generated: + if manager_name not in self.model_classdef.info.names or self.is_manager_dynamically_generated( + manager_info + ): manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class(manager_name, manager_type) 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 # related manager. + manager_class_name = manager.__class__.__name__ custom_model_manager_name = manager.model.__name__ + "_" + manager_class_name try: manager_type = self.create_new_model_parametrized_manager( 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..bc614acf4 --- /dev/null +++ b/tests/typecheck/managers/querysets/test_as_manager.yml @@ -0,0 +1,214 @@ +- 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() + +- case: handles_type_collision_with_from_queryset + main: | + from myapp.models import MyModel, FromQuerySet + reveal_type(FromQuerySet) # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.ManagerFromModelQuerySet[_T`1]" + reveal_type(MyModel.from_queryset) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.as_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["MyModel"]): + ... + + FromQuerySet = models.Manager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + from_queryset = FromQuerySet() + as_manager = ModelQuerySet.as_manager() diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index ae69499f8..68b8304ea 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 "def [_T <: django.db.models.base.Model] () -> myapp.models.ManagerFromModelQuerySet[_T`1]" + 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" @@ -453,7 +478,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(). @@ -493,3 +517,110 @@ 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 "def [_T <: django.db.models.base.Model] () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]" + reveal_type(SecondManager) # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]" + 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: resolves_name_collision_with_other_module_level_object + main: | + from myapp.models import MyModel, Generated, BaseManagerFromModelQuerySet + reveal_type(BaseManagerFromModelQuerySet) # N: Revealed type is "builtins.int" + reveal_type(Generated()) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet1[]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet1[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 = 1 + Generated = BaseManager.from_queryset(ModelQuerySet) + class MyModel(models.Model): + objects = Generated() + +- case: accepts_explicit_none_as_class_name + main: | + from myapp.models import PositionalNone, NoneAsKwarg + reveal_type(PositionalNone) # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]" + reveal_type(NoneAsKwarg) # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]" + 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 "def [_T <: django.db.models.base.Model] () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]" + 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)) 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