Skip to content

Commit f041914

Browse files
UnknownPlatypusUnknownPlatypus
authored
Fix union / intersection / difference to support QuerySet using values_list and values (#2829)
1 parent de8e0b4 commit f041914

File tree

4 files changed

+87
-12
lines changed

4 files changed

+87
-12
lines changed

‎django-stubs/db/models/manager.pyi‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -127,9 +127,9 @@ class Manager(BaseManager[_T]):
127127
def complex_filter(self, filter_obj: Any) -> QuerySet[_T]: ...
128128
def count(self) -> int: ...
129129
async def acount(self) -> int: ...
130-
def union(self, *other_qs: QuerySet[Model], all: bool = ...) -> QuerySet[_T]: ...
131-
def intersection(self, *other_qs: QuerySet[Model]) -> QuerySet[_T]: ...
132-
def difference(self, *other_qs: QuerySet[Model]) -> QuerySet[_T]: ...
130+
def union(self, *other_qs: QuerySet[Model, Any], all: bool = ...) -> QuerySet[_T]: ...
131+
def intersection(self, *other_qs: QuerySet[Model, Any]) -> QuerySet[_T]: ...
132+
def difference(self, *other_qs: QuerySet[Model, Any]) -> QuerySet[_T]: ...
133133
def select_for_update(
134134
self, nowait: bool = ..., skip_locked: bool = ..., of: Sequence[str] = ..., no_key: bool = ...
135135
) -> QuerySet[_T]: ...

‎django-stubs/db/models/query.pyi‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -175,9 +175,9 @@ class QuerySet(Iterable[_Row], Sized, Generic[_Model, _Row]):
175175
def complex_filter(self, filter_obj: Any) -> Self: ...
176176
def count(self) -> int: ...
177177
async def acount(self) -> int: ...
178-
def union(self, *other_qs: QuerySet[Model], all: bool = False) -> Self: ...
179-
def intersection(self, *other_qs: QuerySet[Model]) -> Self: ...
180-
def difference(self, *other_qs: QuerySet[Model]) -> Self: ...
178+
def union(self, *other_qs: QuerySet[Model, Any], all: bool = False) -> Self: ...
179+
def intersection(self, *other_qs: QuerySet[Model, Any]) -> Self: ...
180+
def difference(self, *other_qs: QuerySet[Model, Any]) -> Self: ...
181181
def select_for_update(
182182
self, nowait: bool = False, skip_locked: bool = False, of: Sequence[str] = (), no_key: bool = False
183183
) -> Self: ...

‎tests/typecheck/managers/querysets/test_basic_methods.yml‎

Lines changed: 78 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,23 +70,98 @@
7070
7171
- case: queryset_method_of_union
7272
main: |
73-
from typing_extensions import reveal_type
73+
from typing_extensions import TypedDict, reveal_type
74+
from django.contrib.auth.models import User
75+
from django.db import models
7476
from myapp.models import MyModel1, MyModel2
7577
kls: type[MyModel1 | MyModel2] = MyModel1
78+
7679
reveal_type(kls.objects) # N: Revealed type is "django.db.models.manager.Manager[myapp.models.MyModel1] | django.db.models.manager.Manager[myapp.models.MyModel2]"
7780
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]"
7881
reveal_type(kls.objects.get()) # N: Revealed type is "myapp.models.MyModel1 | myapp.models.MyModel2"
82+
83+
# Regular QuerySet
84+
foos = MyModel1.objects.all()
85+
bars = MyModel2.objects.all()
86+
87+
reveal_type(foos.union(bars)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, myapp.models.MyModel1]"
88+
reveal_type(foos.intersection(bars)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, myapp.models.MyModel1]"
89+
reveal_type(foos.difference(bars)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, myapp.models.MyModel1]"
90+
91+
# `values_list()` QuerySet
92+
foos_list = MyModel1.objects.all().values_list("name")
93+
bars_list = MyModel2.objects.all().values_list("name")
94+
95+
reveal_type(foos_list.union(bars_list)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, tuple[builtins.str]]"
96+
reveal_type(foos_list.intersection(bars_list)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, tuple[builtins.str]]"
97+
reveal_type(foos_list.difference(bars_list)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, tuple[builtins.str]]"
98+
99+
def union_values_list() -> models.QuerySet[models.Model, tuple[str]]:
100+
union = foos_list.union(bars_list)
101+
return union
102+
103+
# `values()` QuerySet -- One field, same name
104+
foos_dict = MyModel1.objects.all().values("name")
105+
bars_dict = MyModel2.objects.all().values("name")
106+
107+
reveal_type(foos_dict.union(bars_dict)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]"
108+
reveal_type(foos_dict.intersection(bars_dict)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]"
109+
reveal_type(foos_dict.difference(bars_dict)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]"
110+
111+
# `values()` QuerySet -- One field, different name, the first one takes precedence.
112+
# cf https://docs.djangoproject.com/en/5.2/ref/models/querysets/#django.db.models.query.QuerySet.union
113+
foos_dict2 = MyModel1.objects.all().values("name")
114+
bars_dict2 = MyModel2.objects.all().values("id")
115+
116+
reveal_type(foos_dict2.union(bars_dict2)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]"
117+
reveal_type(foos_dict2.intersection(bars_dict2)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]"
118+
reveal_type(foos_dict2.difference(bars_dict2)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel1, TypedDict({'name': builtins.str})]"
119+
120+
# `values()` QuerySet -- Multiple field, same names
121+
foos_dict3 = MyModel1.objects.all().values("name", "id")
122+
bars_dict3 = MyModel2.objects.all().values("name", "id")
123+
124+
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})]"
125+
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})]"
126+
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})]"
127+
128+
# Mixing `values()` and `values_list()` is apparently ok and takes the type of the first queryset
129+
foos_dict4 = MyModel1.objects.all().values("name", "id")
130+
bars_list4 = MyModel2.objects.all().values_list("name", "id")
131+
132+
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})]"
133+
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})]"
134+
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})]"
135+
136+
# TODO: Should be an error (mixed number of selected columns) -- runtime raises 'The used SELECT statements have a different number of columns'
137+
foos_dict5 = MyModel1.objects.all().values("name", "id")
138+
bars_dict5 = MyModel2.objects.all().values("name")
139+
140+
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})]"
141+
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})]"
142+
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})]"
143+
144+
# TODO: Should be an error (mixed number of columns) -- runtime raises 'The used SELECT statements have a different number of columns'
145+
model2_qs = MyModel2.objects.all()
146+
user_qs = User.objects.all()
147+
148+
reveal_type(model2_qs.union(bars_dict5)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel2, myapp.models.MyModel2]"
149+
reveal_type(model2_qs.intersection(bars_dict5)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel2, myapp.models.MyModel2]"
150+
reveal_type(model2_qs.difference(bars_dict5)) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel2, myapp.models.MyModel2]"
151+
152+
79153
installed_apps:
154+
- django.contrib.auth
80155
- myapp
81156
files:
82157
- path: myapp/__init__.py
83158
- path: myapp/models.py
84159
content: |
85160
from django.db import models
86161
class MyModel1(models.Model):
87-
pass
162+
name = models.TextField(blank=False, null=False)
88163
class MyModel2(models.Model):
89-
pass
164+
name = models.TextField(blank=False, null=False)
90165
91166
- case: select_related_returns_queryset
92167
main: |

‎tests/typecheck/managers/querysets/test_from_queryset.yml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -665,20 +665,20 @@
665665
reveal_type(MyModel.objects.annotate) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
666666
reveal_type(MyModel.objects.complex_filter) # N: Revealed type is "def (filter_obj: Any) -> myapp.models.MyQuerySet"
667667
reveal_type(MyModel.objects.defer) # N: Revealed type is "Overload(def (None) -> myapp.models.MyQuerySet, def (*fields: builtins.str) -> myapp.models.MyQuerySet)"
668-
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"
668+
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"
669669
reveal_type(MyModel.objects.distinct) # N: Revealed type is "def (*field_names: builtins.str) -> myapp.models.MyQuerySet"
670670
reveal_type(MyModel.objects.exclude) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
671671
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"
672672
reveal_type(MyModel.objects.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet"
673-
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"
673+
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"
674674
reveal_type(MyModel.objects.none) # N: Revealed type is "def () -> myapp.models.MyQuerySet"
675675
reveal_type(MyModel.objects.only) # N: Revealed type is "def (*fields: builtins.str) -> myapp.models.MyQuerySet"
676676
reveal_type(MyModel.objects.order_by) # N: Revealed type is "def (*field_names: builtins.str | django.db.models.expressions.Combinable) -> myapp.models.MyQuerySet"
677677
reveal_type(MyModel.objects.prefetch_related) # N: Revealed type is "Overload(def (None) -> myapp.models.MyQuerySet, def (*lookups: builtins.str | django.db.models.query.Prefetch[_LookupT`-1, _PrefetchedQuerySetT`-2 = django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model], _ToAttrT`-3 = builtins.str]) -> myapp.models.MyQuerySet)"
678678
reveal_type(MyModel.objects.reverse) # N: Revealed type is "def () -> myapp.models.MyQuerySet"
679679
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"
680680
reveal_type(MyModel.objects.select_related) # N: Revealed type is "Overload(def (None) -> myapp.models.MyQuerySet, def (*fields: builtins.str) -> myapp.models.MyQuerySet)"
681-
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"
681+
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"
682682
reveal_type(MyModel.objects.using) # N: Revealed type is "def (alias: builtins.str | None) -> myapp.models.MyQuerySet"
683683
installed_apps:
684684
- myapp

0 commit comments

Comments
 (0)