diff --git a/django-stubs/db/models/manager.pyi b/django-stubs/db/models/manager.pyi index fe4312dba..46e4012c0 100644 --- a/django-stubs/db/models/manager.pyi +++ b/django-stubs/db/models/manager.pyi @@ -127,9 +127,9 @@ class Manager(BaseManager[_T]): def complex_filter(self, filter_obj: Any) -> QuerySet[_T]: ... def count(self) -> int: ... async def acount(self) -> int: ... - def union(self, *other_qs: QuerySet[Model], all: bool = ...) -> QuerySet[_T]: ... - def intersection(self, *other_qs: QuerySet[Model]) -> QuerySet[_T]: ... - def difference(self, *other_qs: QuerySet[Model]) -> QuerySet[_T]: ... + def union(self, *other_qs: QuerySet[Model, Any], all: bool = ...) -> QuerySet[_T]: ... + def intersection(self, *other_qs: QuerySet[Model, Any]) -> QuerySet[_T]: ... + def difference(self, *other_qs: QuerySet[Model, Any]) -> QuerySet[_T]: ... def select_for_update( self, nowait: bool = ..., skip_locked: bool = ..., of: Sequence[str] = ..., no_key: bool = ... ) -> QuerySet[_T]: ... diff --git a/django-stubs/db/models/query.pyi b/django-stubs/db/models/query.pyi index a6d840200..4c58066bf 100644 --- a/django-stubs/db/models/query.pyi +++ b/django-stubs/db/models/query.pyi @@ -175,9 +175,9 @@ class QuerySet(Iterable[_Row], Sized, Generic[_Model, _Row]): def complex_filter(self, filter_obj: Any) -> Self: ... def count(self) -> int: ... async def acount(self) -> int: ... - def union(self, *other_qs: QuerySet[Model], all: bool = False) -> Self: ... - def intersection(self, *other_qs: QuerySet[Model]) -> Self: ... - def difference(self, *other_qs: QuerySet[Model]) -> Self: ... + def union(self, *other_qs: QuerySet[Model, Any], all: bool = False) -> Self: ... + def intersection(self, *other_qs: QuerySet[Model, Any]) -> Self: ... + def difference(self, *other_qs: QuerySet[Model, Any]) -> Self: ... def select_for_update( self, nowait: bool = False, skip_locked: bool = False, of: Sequence[str] = (), no_key: bool = False ) -> Self: ... diff --git a/tests/typecheck/managers/querysets/test_basic_methods.yml b/tests/typecheck/managers/querysets/test_basic_methods.yml index b85fcce97..394a66ea0 100644 --- a/tests/typecheck/managers/querysets/test_basic_methods.yml +++ b/tests/typecheck/managers/querysets/test_basic_methods.yml @@ -70,13 +70,88 @@ - case: queryset_method_of_union main: | - from typing_extensions import reveal_type + from typing_extensions import TypedDict, reveal_type + from django.contrib.auth.models import User + from django.db import models from myapp.models import MyModel1, MyModel2 kls: type[MyModel1 | MyModel2] = MyModel1 + reveal_type(kls.objects) # N: Revealed type is "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 "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 "myapp.models.MyModel1 | myapp.models.MyModel2" + + # Regular QuerySet + foos = MyModel1.objects.all() + bars = MyModel2.objects.all() + + reveal_type(foos.union(bars)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, myapp.models.MyModel1]" + reveal_type(foos.intersection(bars)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, myapp.models.MyModel1]" + reveal_type(foos.difference(bars)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, myapp.models.MyModel1]" + + # `values_list()` QuerySet + foos_list = MyModel1.objects.all().values_list("name") + bars_list = MyModel2.objects.all().values_list("name") + + reveal_type(foos_list.union(bars_list)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, tuple[builtins.str]]" + reveal_type(foos_list.intersection(bars_list)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, tuple[builtins.str]]" + reveal_type(foos_list.difference(bars_list)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, tuple[builtins.str]]" + + def union_values_list() -> models.QuerySet[models.Model, tuple[str]]: + union = foos_list.union(bars_list) + return union + + # `values()` QuerySet -- One field, same name + foos_dict = MyModel1.objects.all().values("name") + bars_dict = MyModel2.objects.all().values("name") + + reveal_type(foos_dict.union(bars_dict)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]" + reveal_type(foos_dict.intersection(bars_dict)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]" + reveal_type(foos_dict.difference(bars_dict)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]" + + # `values()` QuerySet -- One field, different name, the first one takes precedence. + # cf https://docs.djangoproject.com/en/5.2/ref/models/querysets/#django.db.models.query.QuerySet.union + foos_dict2 = MyModel1.objects.all().values("name") + bars_dict2 = MyModel2.objects.all().values("id") + + reveal_type(foos_dict2.union(bars_dict2)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]" + reveal_type(foos_dict2.intersection(bars_dict2)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]" + reveal_type(foos_dict2.difference(bars_dict2)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]" + + # `values()` QuerySet -- Multiple field, same names + foos_dict3 = MyModel1.objects.all().values("name", "id") + bars_dict3 = MyModel2.objects.all().values("name", "id") + + reveal_type(foos_dict3.union(bars_dict3)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str, 'id': builtins.int})]" + reveal_type(foos_dict3.intersection(bars_dict3)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str, 'id': builtins.int})]" + reveal_type(foos_dict3.difference(bars_dict3)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str, 'id': builtins.int})]" + + # Mixing `values()` and `values_list()` is apparently ok and takes the type of the first queryset + foos_dict4 = MyModel1.objects.all().values("name", "id") + bars_list4 = MyModel2.objects.all().values_list("name", "id") + + reveal_type(foos_dict4.union(bars_list4)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str, 'id': builtins.int})]" + reveal_type(foos_dict4.intersection(bars_list4)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str, 'id': builtins.int})]" + reveal_type(foos_dict4.difference(bars_list4)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str, 'id': builtins.int})]" + + # TODO: Should be an error (mixed number of selected columns) -- runtime raises 'The used SELECT statements have a different number of columns' + foos_dict5 = MyModel1.objects.all().values("name", "id") + bars_dict5 = MyModel2.objects.all().values("name") + + reveal_type(foos_dict5.union(bars_dict5)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str, 'id': builtins.int})]" + reveal_type(foos_dict5.intersection(bars_dict5)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str, 'id': builtins.int})]" + reveal_type(foos_dict5.difference(bars_dict5)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str, 'id': builtins.int})]" + + # TODO: Should be an error (mixed number of columns) -- runtime raises 'The used SELECT statements have a different number of columns' + model2_qs = MyModel2.objects.all() + user_qs = User.objects.all() + + reveal_type(model2_qs.union(bars_dict5)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel2, myapp.models.MyModel2]" + reveal_type(model2_qs.intersection(bars_dict5)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel2, myapp.models.MyModel2]" + reveal_type(model2_qs.difference(bars_dict5)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel2, myapp.models.MyModel2]" + + installed_apps: + - django.contrib.auth - myapp files: - path: myapp/__init__.py @@ -84,9 +159,9 @@ content: | from django.db import models class MyModel1(models.Model): - pass + name = models.TextField(blank=False, null=False) class MyModel2(models.Model): - pass + name = models.TextField(blank=False, null=False) - case: select_related_returns_queryset main: | diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index f1d8a7d6c..c2b291ee7 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -665,12 +665,12 @@ reveal_type(MyModel.objects.annotate) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.complex_filter) # N: Revealed type is "def (filter_obj: Any) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.defer) # N: Revealed type is "Overload(def (None) -> myapp.models.MyQuerySet, def (*fields: builtins.str) -> myapp.models.MyQuerySet)" - reveal_type(MyModel.objects.difference) # N: Revealed type is "def (*other_qs: django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model]) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.difference) # N: Revealed type is "def (*other_qs: django.db.models.query.QuerySet[django.db.models.base.Model, Any]) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.distinct) # N: Revealed type is "def (*field_names: builtins.str) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.exclude) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.extra) # N: Revealed type is "def (select: builtins.dict[builtins.str, Any] | None =, where: typing.Sequence[builtins.str] | None =, params: typing.Sequence[Any] | None =, tables: typing.Sequence[builtins.str] | None =, order_by: typing.Sequence[builtins.str | django.db.models.expressions.Combinable] | None =, select_params: typing.Sequence[Any] | None =) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" - reveal_type(MyModel.objects.intersection) # N: Revealed type is "def (*other_qs: django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model]) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.intersection) # N: Revealed type is "def (*other_qs: django.db.models.query.QuerySet[django.db.models.base.Model, Any]) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.none) # N: Revealed type is "def () -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.only) # N: Revealed type is "def (*fields: builtins.str) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.order_by) # N: Revealed type is "def (*field_names: builtins.str | django.db.models.expressions.Combinable) -> myapp.models.MyQuerySet" @@ -678,7 +678,7 @@ reveal_type(MyModel.objects.reverse) # N: Revealed type is "def () -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.select_for_update) # N: Revealed type is "def (nowait: builtins.bool =, skip_locked: builtins.bool =, of: typing.Sequence[builtins.str] =, no_key: builtins.bool =) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.select_related) # N: Revealed type is "Overload(def (None) -> myapp.models.MyQuerySet, def (*fields: builtins.str) -> myapp.models.MyQuerySet)" - reveal_type(MyModel.objects.union) # N: Revealed type is "def (*other_qs: django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model], all: builtins.bool =) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.union) # N: Revealed type is "def (*other_qs: django.db.models.query.QuerySet[django.db.models.base.Model, Any], all: builtins.bool =) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.using) # N: Revealed type is "def (alias: builtins.str | None) -> myapp.models.MyQuerySet" installed_apps: - myapp