Determine the type of queryset methods on unions

mypy_django_plugin/transformers/managers.py

@@ -19,7 +19,18 @@
 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,
+    ProperType,
+    TypeOfAny,
+    TypeVarType,
+    UnionType,
+    get_proper_type,
+)
 from mypy.types import Type as MypyType
 from mypy.typevars import fill_typevars
 
@@ -274,6 +285,13 @@ def resolve_manager_method(ctx: AttributeContext) -> MypyType:
 
     if isinstance(ctx.type, Instance):
         return resolve_manager_method_from_instance(instance=ctx.type, method_name=method_name, ctx=ctx)
+    elif isinstance(ctx.type, UnionType) and all(isinstance(instance, Instance) for instance in ctx.type.items):
+        resolved = tuple(
+            resolve_manager_method_from_instance(instance=instance, method_name=method_name, ctx=ctx)
+            for instance in ctx.type.items
+            if isinstance(instance, Instance)
+        )
+        return get_proper_type(UnionType(resolved))
     else:
         ctx.api.fail(f'Unable to resolve return type of queryset/manager method "{method_name}"', ctx.context)
         return AnyType(TypeOfAny.from_error)

tests/typecheck/managers/querysets/test_as_manager.yml

@@ -140,6 +140,44 @@
                 class MyModel(models.Model):
                     objects = ModelQuerySet.as_manager()
 
+-   case: includes_custom_queryset_methods_on_unions
+    main: |
+        from myapp.models import MyModel1, MyModel2
+        import typing
+        kls: typing.Type[typing.Union[MyModel1, MyModel2]] = MyModel1
+        reveal_type(kls.objects.custom_queryset_method())  # N: Revealed type is "Union[myapp.models.ModelQuerySet1, myapp.models.ModelQuerySet2]"
+        reveal_type(kls.objects.all().custom_queryset_method())  # N: Revealed type is "Union[myapp.models.ModelQuerySet1, myapp.models.ModelQuerySet2]"
+        reveal_type(kls.objects.returns_thing())  # N: Revealed type is "Union[builtins.int, builtins.str]"
+        reveal_type(kls.objects.get())  # N: Revealed type is "Union[myapp.models.MyModel1, myapp.models.MyModel2]"
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from typing import Sequence
+
+                class ModelQuerySet1(models.QuerySet["MyModel1"]):
+                    def custom_queryset_method(self) -> "ModelQuerySet1":
+                        return self.all()
+
+                    def returns_thing(self) -> int:
+                      return 1
+
+                class ModelQuerySet2(models.QuerySet["MyModel2"]):
+                    def custom_queryset_method(self) -> "ModelQuerySet2":
+                        return self.all()
+
+                    def returns_thing(self) -> str:
+                      return "asdf"
+
+                class MyModel1(models.Model):
+                    objects = ModelQuerySet1.as_manager()
+
+                class MyModel2(models.Model):
+                    objects = ModelQuerySet2.as_manager()
+
 -   case: handles_call_outside_of_model_class_definition
     main: |
         from myapp.models import MyModel, MyModelManager

tests/typecheck/managers/querysets/test_basic_methods.yml

@@ -59,6 +59,26 @@
                 class User(models.Model):
                     pass
 
+-   case: queryset_method_of_union
+    main: |
+        from myapp.models import MyModel1, MyModel2
+        import typing
+        kls: typing.Type[typing.Union[MyModel1, MyModel2]] = MyModel1
+        reveal_type(kls.objects)  # N: Revealed type is "Union[django.db.models.manager.Manager[myapp.models.MyModel1], django.db.models.manager.Manager[myapp.models.MyModel2]]"
+        reveal_type(kls.objects.all())  # N: Revealed type is "Union[django.db.models.query._QuerySet[myapp.models.MyModel1, myapp.models.MyModel1], django.db.models.query._QuerySet[myapp.models.MyModel2, myapp.models.MyModel2]]"
+        reveal_type(kls.objects.get())  # N: Revealed type is "Union[myapp.models.MyModel1, myapp.models.MyModel2]"
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class MyModel1(models.Model):
+                    pass
+                class MyModel2(models.Model):
+                    pass
+
 -   case: select_related_returns_queryset
     main: |
         from myapp.models import Book

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -276,6 +276,39 @@
                 class MyModel(models.Model):
                     objects = NewManager()
 
+-   case: from_queryset_with_manager_of_union
+    main: |
+        from myapp.models import MyModel1, MyModel2
+        import typing
+        kls: typing.Type[typing.Union[MyModel1, MyModel2]] = MyModel1
+        reveal_type(kls.objects)  # N: Revealed type is "Union[myapp.models.ManagerFromModelQuerySet1[myapp.models.MyModel1], myapp.models.ManagerFromModelQuerySet2[myapp.models.MyModel2]]"
+        reveal_type(kls.objects.all())  # N: Revealed type is "Union[myapp.models.ModelQuerySet1[myapp.models.MyModel1], myapp.models.ModelQuerySet2[myapp.models.MyModel2]]"
+        reveal_type(kls.objects.get())  # N: Revealed type is "Union[myapp.models.MyModel1, myapp.models.MyModel2]"
+        reveal_type(kls.objects.queryset_method())  # N: Revealed type is "Union[builtins.int, builtins.str]"
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+
+                class ModelQuerySet1(models.QuerySet["MyModel1"]):
+                    def queryset_method(self) -> int:
+                        return 1
+
+                class ModelQuerySet2(models.QuerySet["MyModel2"]):
+                    def queryset_method(self) -> str:
+                        return 'hello'
+
+                NewManager1 = models.Manager.from_queryset(ModelQuerySet1)
+                class MyModel1(models.Model):
+                    objects = NewManager1()
+
+                NewManager2 = models.Manager.from_queryset(ModelQuerySet2)
+                class MyModel2(models.Model):
+                    objects = NewManager2()
+
 -   case: from_queryset_returns_intersection_of_manager_and_queryset
     main: |
         from myapp.models import MyModel, NewManager

tests/typecheck/managers/querysets/test_union_type.yml

@@ -10,7 +10,7 @@
 
         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.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"  [misc]
+        reveal_type(model_cls.objects.my_method())  # N: Revealed type is "myapp.models.MyQuerySet"
     installed_apps:
         - myapp
     files: