From db6d239900e199468632a7f3e98e352333406bca Mon Sep 17 00:00:00 2001 From: Stephen Moore Date: Sun, 24 Mar 2024 10:07:32 +1100 Subject: [PATCH] Determine the type of queryset methods on unions Prior to this if a type was a union of multiple django model managers then it wouldn't be able to determine the type of the queryset methods on them. This change will mean that in the case of a union, the type of the queryset method will be a union of the types of all the queryset methods on each item in the union --- mypy_django_plugin/transformers/managers.py | 33 ++++++++++++++- .../managers/querysets/test_as_manager.yml | 42 +++++++++++++++++++ .../managers/querysets/test_basic_methods.yml | 24 +++++++++++ .../managers/querysets/test_from_queryset.yml | 37 ++++++++++++++++ .../managers/querysets/test_union_type.yml | 2 +- 5 files changed, 136 insertions(+), 2 deletions(-) diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index c1d0e91779..9bf29db1eb 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/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,26 @@ 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): + if any(isinstance(instance, AnyType) for instance in ctx.type.items): + ctx.api.fail( + f'Unable to resolve return type of queryset/manager method "{method_name}" on a union involving Any', + ctx.context, + ) + return AnyType(TypeOfAny.from_error) + elif any(not isinstance(instance, Instance) for instance in ctx.type.items): + ctx.api.fail( + f'Unable to resolve return type of queryset/manager method "{method_name}"', + ctx.context, + ) + return AnyType(TypeOfAny.from_error) + else: + resolved = [ + 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(tuple(resolved))) else: ctx.api.fail(f'Unable to resolve return type of queryset/manager method "{method_name}"', ctx.context) return AnyType(TypeOfAny.from_error) diff --git a/tests/typecheck/managers/querysets/test_as_manager.yml b/tests/typecheck/managers/querysets/test_as_manager.yml index 13bf7894af..ceb199df9f 100644 --- a/tests/typecheck/managers/querysets/test_as_manager.yml +++ b/tests/typecheck/managers/querysets/test_as_manager.yml @@ -140,6 +140,48 @@ class MyModel(models.Model): objects = ModelQuerySet.as_manager() +- case: includes_custom_queryset_methods_on_unions + main: | + from myapp.models import MyModel1, MyModel2 + import sys + if sys.version_info <= (3, 8): + import typing + kls: typing.Type[MyModel1 | MyModel2] = MyModel1 + else: + kls: type[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 diff --git a/tests/typecheck/managers/querysets/test_basic_methods.yml b/tests/typecheck/managers/querysets/test_basic_methods.yml index 8c37cba063..314e98f6a8 100644 --- a/tests/typecheck/managers/querysets/test_basic_methods.yml +++ b/tests/typecheck/managers/querysets/test_basic_methods.yml @@ -59,6 +59,30 @@ class User(models.Model): pass +- case: queryset_method_of_union + main: | + from myapp.models import MyModel1, MyModel2 + import sys + if sys.version_info <= (3, 8): + import typing + kls: typing.Type[MyModel1 | MyModel2] = MyModel1 + else: + kls: type[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 diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index 7a595a6337..15c6f392ad 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -276,6 +276,43 @@ class MyModel(models.Model): objects = NewManager() +- case: from_queryset_with_manager_of_union + main: | + from myapp.models import MyModel1, MyModel2 + import sys + if sys.version_info <= (3, 8): + import typing + kls: typing.Type[MyModel1 | MyModel2] = MyModel1 + else: + kls: type[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 diff --git a/tests/typecheck/managers/querysets/test_union_type.yml b/tests/typecheck/managers/querysets/test_union_type.yml index 27b9886ec8..6863825103 100644 --- a/tests/typecheck/managers/querysets/test_union_type.yml +++ b/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: