diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 9fada10..b8bc729 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -3,6 +3,7 @@ from typing import Optional, Union from mypy.checker import TypeChecker, fill_typevars from mypy.nodes import ( GDEF, + CallExpr, Decorator, FuncBase, FuncDef, @@ -96,27 +97,39 @@ 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.type + ) or get_method_type_from_reverse_manager(api, method_name, instance.type) + + return method_type if method_type is not None else ctx.default_attr_type + + 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. """ - api = helpers.get_typechecker_api(ctx) # Skip (method) type that is currently something other than Any if not isinstance(ctx.default_attr_type, AnyType): 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. - assert isinstance(ctx.context, MemberExpr) - method_name = ctx.context.name - manager_instance = ctx.type - assert isinstance(manager_instance, Instance) - method_type = get_method_type_from_dynamic_manager( - api, method_name, manager_instance.type - ) or get_method_type_from_reverse_manager(api, method_name, manager_instance.type) + if isinstance(ctx.context, MemberExpr): + method_name = ctx.context.name + elif isinstance(ctx.context, CallExpr) and isinstance(ctx.context.callee, MemberExpr): + method_name = ctx.context.callee.name + else: + ctx.api.fail("Unable to resolve return type of queryset/manager method", ctx.context) + return AnyType(TypeOfAny.from_error) - return method_type if method_type is not None else ctx.default_attr_type + if isinstance(ctx.type, Instance): + return resolve_manager_method_from_instance(instance=ctx.type, method_name=method_name, ctx=ctx) + else: + ctx.api.fail(f'Unable to resolve return type of queryset/manager method "{method_name}"', ctx.context) + return AnyType(TypeOfAny.from_error) def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext) -> None: diff --git a/tests/typecheck/managers/querysets/test_union_type.yml b/tests/typecheck/managers/querysets/test_union_type.yml new file mode 100644 index 0000000..7289603 --- /dev/null +++ b/tests/typecheck/managers/querysets/test_union_type.yml @@ -0,0 +1,33 @@ +- case: union_queryset_custom_method + main: | + from typing import Union + from django.db.models import QuerySet + from myapp.models import User, Order + + instance: Union[Order, User] = User() + + 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]]" + model_cls.objects.my_method() # E: Unable to resolve return type of queryset/manager method "my_method" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from __future__ import annotations + from django.db import models + + class MyQuerySet(models.QuerySet): + def my_method(self) -> MyQuerySet: + pass + + UserManager = models.Manager.from_queryset(MyQuerySet) + class User(models.Model): + objects = UserManager() + + OrderManager = models.Manager.from_queryset(MyQuerySet) + class Order(models.Model): + objects = OrderManager()