diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 028cb43d9..acf45147a 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -3,23 +3,36 @@ from mypy.checker import TypeChecker from mypy.nodes import ( GDEF, + AssignmentStmt, CallExpr, Decorator, FuncBase, FuncDef, MemberExpr, + NameExpr, Node, OverloadedFuncDef, RefExpr, StrExpr, SymbolTableNode, + TypeAlias, TypeInfo, Var, ) from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext from mypy.semanal import SemanticAnalyzer from mypy.semanal_shared import has_placeholder -from mypy.types import AnyType, CallableType, FunctionLike, Instance, Overloaded, ProperType, TypeOfAny, TypeVarType +from mypy.types import ( + AnyType, + CallableType, + FunctionLike, + Instance, + Overloaded, + PlaceholderType, + ProperType, + TypeOfAny, + TypeVarType, +) from mypy.types import Type as MypyType from mypy.typevars import fill_typevars @@ -58,10 +71,11 @@ def get_method_type_from_dynamic_manager( Attempt to resolve a method on a manager that was built from '.from_queryset' """ - manager_type_info = manager_instance.type + manager_type_info = manager_instance.type.get_containing_type_info(method_name) if ( - "django" not in manager_type_info.metadata + manager_type_info is None + or "django" not in manager_type_info.metadata or "from_queryset_manager" not in manager_type_info.metadata["django"] ): # Manager isn't dynamically added @@ -347,24 +361,73 @@ def create_manager_info_from_from_queryset_call( """ Extract manager and queryset TypeInfo from a from_queryset call. """ - if ( # Check that this is a from_queryset call on a manager subclass not isinstance(call_expr.callee, MemberExpr) + or not call_expr.callee.name == "from_queryset" or not isinstance(call_expr.callee.expr, RefExpr) or not isinstance(call_expr.callee.expr.node, TypeInfo) or not call_expr.callee.expr.node.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME) - or not call_expr.callee.name == "from_queryset" + ): + return None + + module = api.modules[api.cur_mod_id] + base_manager_info = call_expr.callee.expr.node + # Mypy parses `NewManager = .from_queryset()` as a `Var` + # assignment, but what is happening runtime is rather that `NewManager` becomes an + # alias to another type. But mypy can't see that type. Before anything else happens, + # we want to adjust this. But we don't have any type to alias yet and we can't be + # sure we _can_ have it yet either. Due to the state of eventual dependencies we + # might need to defer, so to begin with we insert a placeholder type. + if name is not None: # New manager is aliased + assignment_symbol = module.names.get(name) + if ( + assignment_symbol is not None + and assignment_symbol.fullname is not None + and not assignment_symbol.plugin_generated + and not isinstance(assignment_symbol.node, TypeAlias) + ): + assignment_stmt = next( + ( + stmt + for stmt in module.defs + if ( + isinstance(stmt, AssignmentStmt) + and len(stmt.lvalues) == 1 + and isinstance(stmt.lvalues[0], NameExpr) + and stmt.lvalues[0].name == name + ) + ), + None, + ) + if assignment_stmt is not None: + # Mypy don't allow call expressions to be aliases, but to get any kind + # of handle to the invisible type we spoof mypy that this statement is a + # type alias.. + assignment_stmt.is_alias_def = True + alias = TypeAlias( + target=PlaceholderType( + fullname=assignment_symbol.fullname, + args=list(base_manager_info.defn.type_vars), + line=call_expr.line, + ), + fullname=assignment_symbol.fullname, + line=call_expr.line, + column=0, + ) + module.names[name] = SymbolTableNode(GDEF, alias, plugin_generated=True) + + if ( # Check that the call has one or two arguments and that the first is a # QuerySet subclass - or not 1 <= len(call_expr.args) <= 2 + not 1 <= len(call_expr.args) <= 2 or not isinstance(call_expr.args[0], RefExpr) or not isinstance(call_expr.args[0].node, TypeInfo) or not call_expr.args[0].node.has_base(fullnames.QUERYSET_CLASS_FULLNAME) ): return None - base_manager_info, queryset_info = call_expr.callee.expr.node, call_expr.args[0].node + queryset_info = call_expr.args[0].node if queryset_info.fullname is None: # In some cases, due to the way the semantic analyzer works, only # passed_queryset.name is available. But it should be analyzed again, @@ -408,10 +471,19 @@ def create_manager_info_from_from_queryset_call( # Add the new manager to the current module # TODO: use proper SemanticAnalyzer API for that. - module = api.modules[api.cur_mod_id] - 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) + if name is not None: + assignment_symbol = module.names.get(name) + if ( + assignment_symbol is not None + and name != new_manager_info.name + and assignment_symbol.plugin_generated + and isinstance(assignment_symbol.node, TypeAlias) + and isinstance(assignment_symbol.node.target, PlaceholderType) + ): + # Unless names are equal, there's 2 symbol names that needs to be updated + manager_instance = fill_typevars(new_manager_info) + assert isinstance(manager_instance, Instance) + assignment_symbol.node.target = manager_instance module.names[new_manager_info.name] = SymbolTableNode(GDEF, new_manager_info, plugin_generated=True) return new_manager_info diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index c8f9d5d38..32c72a7bd 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -642,8 +642,8 @@ reveal_type(user.article_set) # N: Revealed type is "myapp.models.Article_RelatedManager" reveal_type(user.book_set.add) # N: Revealed type is "def (*objs: Union[myapp.models.Book, builtins.int], *, bulk: builtins.bool =)" reveal_type(user.article_set.add) # N: Revealed type is "def (*objs: Union[myapp.models.Article, builtins.int], *, bulk: builtins.bool =)" - reveal_type(user.book_set.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet[myapp.models.Book]" - reveal_type(user.article_set.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet[myapp.models.Article]" + reveal_type(user.book_set.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet" + reveal_type(user.article_set.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet" reveal_type(user.book_set.queryset_method()) # N: Revealed type is "builtins.int" reveal_type(user.article_set.queryset_method()) # N: Revealed type is "builtins.int" installed_apps: @@ -798,11 +798,11 @@ from myapp.models.user import User reveal_type(Store().purchases) # N: Revealed type is "myapp.models.purchase.Purchase_RelatedManager" reveal_type(Store().purchases.queryset_method()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet" - reveal_type(Store().purchases.filter()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase]" + reveal_type(Store().purchases.filter()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet" reveal_type(Store().purchases.filter().queryset_method()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet" reveal_type(User().purchases) # N: Revealed type is "myapp.models.purchase.Purchase_RelatedManager" reveal_type(User().purchases.queryset_method()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet" - reveal_type(User().purchases.filter()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase]" + reveal_type(User().purchases.filter()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet" reveal_type(User().purchases.filter().queryset_method()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet" installed_apps: - myapp diff --git a/tests/typecheck/managers/querysets/test_as_manager.yml b/tests/typecheck/managers/querysets/test_as_manager.yml index 830448ca6..d80a4c259 100644 --- a/tests/typecheck/managers/querysets/test_as_manager.yml +++ b/tests/typecheck/managers/querysets/test_as_manager.yml @@ -310,7 +310,7 @@ - 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(FromQuerySet) # N: Revealed type is "def () -> 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: diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index 59f0cd63c..a1e9a21f8 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -158,7 +158,7 @@ - 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]" + reveal_type(NewManager) # N: Revealed type is "def () -> myapp.models.ManagerFromModelQuerySet[_T`1]" installed_apps: - myapp files: @@ -197,7 +197,7 @@ - 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.ModelBaseManagerFromModelQuerySet[]" + reveal_type(NewManager()) # N: Revealed type is "myapp.models.ModelBaseManagerFromModelQuerySet[_T`1]" reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ModelBaseManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel.objects.get()) # N: Revealed type is "myapp.models.MyModel" reveal_type(MyModel.objects.manager_only_method()) # N: Revealed type is "builtins.int" @@ -228,7 +228,7 @@ reveal_type(MyModel.objects.get()) # N: Revealed type is "myapp.models.MyModel" 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(OtherManager()) # N: Revealed type is "myapp.models.X[_T`1]" 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" @@ -572,8 +572,8 @@ - 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(FirstManager) # N: Revealed type is "def () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]" + reveal_type(SecondManager) # N: Revealed type is "def () -> 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: @@ -619,7 +619,7 @@ 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(Generated()) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet1[_T`1]" reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet1[myapp.models.MyModel]" installed_apps: - myapp @@ -641,8 +641,8 @@ - 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]" + reveal_type(PositionalNone) # N: Revealed type is "def () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]" + reveal_type(NoneAsKwarg) # N: Revealed type is "def () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]" installed_apps: - myapp files: @@ -661,7 +661,7 @@ - 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]" + reveal_type(StrCallable) # N: Revealed type is "def () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]" installed_apps: - myapp files: @@ -675,3 +675,34 @@ ... StrCallable = BaseManager.from_queryset(ModelQuerySet, class_name=str(1)) + +- case: test_type_generated_from_queryset_can_be_inherited + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects) + reveal_type(MyModel.objects.my_method()) + reveal_type(MyModel.objects.queryset_method()) + out: | + main:2: note: Revealed type is "myapp.models.MyModelManager[myapp.models.MyModel]" + main:3: note: Revealed type is "builtins.str" + main:4: note: Revealed type is "builtins.int" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from __future__ import annotations + from django.db import models + + class CustomQuerySet(models.QuerySet["MyModel"]): + def queryset_method(self) -> int: + ... + + GeneratedManager = models.Manager.from_queryset(CustomQuerySet) + + class MyModelManager(GeneratedManager): + def my_method(self) -> str: + ... + class MyModel(models.Model): + objects = MyModelManager()