Skip to content

Commit fc21bed

Browse files
long2icelong2ice
authored
Add .group_by(...) support (tortoise#363)
1 parent eb0fd88 commit fc21bed

File tree

6 files changed

+279
-6
lines changed

6 files changed

+279
-6
lines changed

‎CHANGELOG.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,10 @@ Changelog
99

1010
0.16
1111
====
12+
1213
0.16.8
1314
------
15+
- Add ``group by`` support
1416
- Fixed regression where ``GROUP BY`` class is missing for an aggregate with a specified order.
1517

1618
0.16.7

‎docs/examples/basic.rst‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,15 @@ Functions
5353

5454
.. _example_schema_create:
5555

56+
Group By
57+
========
58+
.. literalinclude:: ../../examples/group_by.py
59+
60+
61+
.. rst-class:: html-toggle
62+
63+
.. _example_aggregation:
64+
5665
Schema creation
5766
===============
5867
.. literalinclude:: ../../examples/schema_create.py

‎examples/group_by.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
from tortoise import Model, Tortoise, fields, run_async
2+
from tortoise.functions import Avg, Count, Sum
3+
4+
5+
class Author(Model):
6+
name = fields.CharField(max_length=255)
7+
8+
9+
class Book(Model):
10+
name = fields.CharField(max_length=255)
11+
author = fields.ForeignKeyField("models.Author", related_name="books")
12+
rating = fields.FloatField()
13+
14+
15+
async def run():
16+
await Tortoise.init(db_url="sqlite://:memory:", modules={"models": ["__main__"]})
17+
await Tortoise.generate_schemas()
18+
19+
a1 = await Author.create(name="author1")
20+
a2 = await Author.create(name="author2")
21+
for i in range(10):
22+
await Book.create(name=f"book{i}", author=a1, rating=i)
23+
for i in range(5):
24+
await Book.create(name=f"book{i}", author=a2, rating=i)
25+
26+
ret = await Book.annotate(count=Count("id")).group_by("author_id").values("author_id", "count")
27+
print(ret)
28+
# >>> [{'author_id': 1, 'count': 10}, {'author_id': 2, 'count': 5}]
29+
30+
ret = (
31+
await Book.annotate(count=Count("id"))
32+
.filter(count__gt=6)
33+
.group_by("author_id")
34+
.values("author_id", "count")
35+
)
36+
print(ret)
37+
# >>> [{'author_id': 1, 'count': 10}]
38+
39+
ret = await Book.annotate(sum=Sum("rating")).group_by("author_id").values("author_id", "sum")
40+
print(ret)
41+
# >>> [{'author_id': 1, 'sum': 45.0}, {'author_id': 2, 'sum': 10.0}]
42+
43+
ret = (
44+
await Book.annotate(sum=Sum("rating"))
45+
.filter(sum__gt=11)
46+
.group_by("author_id")
47+
.values("author_id", "sum")
48+
)
49+
print(ret)
50+
# >>> [{'author_id': 1, 'sum': 45.0}]
51+
52+
ret = await Book.annotate(avg=Avg("rating")).group_by("author_id").values("author_id", "avg")
53+
print(ret)
54+
# >>> [{'author_id': 1, 'avg': 4.5}, {'author_id': 2, 'avg': 2.0}]
55+
56+
ret = (
57+
await Book.annotate(avg=Avg("rating"))
58+
.filter(avg__gt=3)
59+
.group_by("author_id")
60+
.values("author_id", "avg")
61+
)
62+
print(ret)
63+
# >>> [{'author_id': 1, 'avg': 4.5}]
64+
65+
# and use .values_list()
66+
ret = (
67+
await Book.annotate(count=Count("id"))
68+
.group_by("author_id")
69+
.values_list("author_id", "count")
70+
)
71+
print(ret)
72+
# >>> [(1, 10), (2, 5)]
73+
74+
75+
if __name__ == "__main__":
76+
run_async(run())

‎tests/test_group_by.py‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
from tests.testmodels import Author, Book
2+
from tortoise.contrib import test
3+
from tortoise.functions import Avg, Count, Sum
4+
5+
6+
class TestGroupBy(test.TestCase):
7+
async def setUp(self) -> None:
8+
self.a1 = await Author.create(name="author1")
9+
self.a2 = await Author.create(name="author2")
10+
for i in range(10):
11+
await Book.create(name=f"book{i}", author=self.a1, rating=i)
12+
for i in range(5):
13+
await Book.create(name=f"book{i}", author=self.a2, rating=i)
14+
15+
async def test_count_group_by(self):
16+
ret = (
17+
await Book.annotate(count=Count("id"))
18+
.group_by("author_id")
19+
.values("author_id", "count")
20+
)
21+
22+
for item in ret:
23+
author_id = item.get("author_id")
24+
count = item.get("count")
25+
if author_id == self.a1.pk:
26+
self.assertEqual(count, 10)
27+
elif author_id == self.a2.pk:
28+
self.assertEqual(count, 5)
29+
30+
async def test_count_filter_group_by(self):
31+
ret = (
32+
await Book.annotate(count=Count("id"))
33+
.filter(count__gt=6)
34+
.group_by("author_id")
35+
.values("author_id", "count")
36+
)
37+
self.assertEqual(len(ret), 1)
38+
self.assertEqual(ret[0].get("count"), 10)
39+
40+
async def test_sum_group_by(self):
41+
ret = (
42+
await Book.annotate(sum=Sum("rating")).group_by("author_id").values("author_id", "sum")
43+
)
44+
for item in ret:
45+
author_id = item.get("author_id")
46+
sum_ = item.get("sum")
47+
if author_id == self.a1.pk:
48+
self.assertEqual(sum_, 45.0)
49+
elif author_id == self.a2.pk:
50+
self.assertEqual(sum_, 10.0)
51+
52+
async def test_sum_filter_group_by(self):
53+
ret = (
54+
await Book.annotate(sum=Sum("rating"))
55+
.filter(sum__gt=11)
56+
.group_by("author_id")
57+
.values("author_id", "sum")
58+
)
59+
self.assertEqual(len(ret), 1)
60+
self.assertEqual(ret[0].get("sum"), 45.0)
61+
62+
async def test_avg_group_by(self):
63+
ret = (
64+
await Book.annotate(avg=Avg("rating")).group_by("author_id").values("author_id", "avg")
65+
)
66+
67+
for item in ret:
68+
author_id = item.get("author_id")
69+
avg = item.get("avg")
70+
if author_id == self.a1.pk:
71+
self.assertEqual(avg, 4.5)
72+
elif author_id == self.a2.pk:
73+
self.assertEqual(avg, 2.0)
74+
75+
async def test_avg_filter_group_by(self):
76+
ret = (
77+
await Book.annotate(avg=Avg("rating"))
78+
.filter(avg__gt=3)
79+
.group_by("author_id")
80+
.values_list("author_id", "avg")
81+
)
82+
self.assertEqual(len(ret), 1)
83+
self.assertEqual(ret[0][1], 4.5)
84+
85+
async def test_count_values_list_group_by(self):
86+
ret = (
87+
await Book.annotate(count=Count("id"))
88+
.group_by("author_id")
89+
.values_list("author_id", "count")
90+
)
91+
92+
for item in ret:
93+
author_id = item[0]
94+
count = item[1]
95+
if author_id == self.a1.pk:
96+
self.assertEqual(count, 10)
97+
elif author_id == self.a2.pk:
98+
self.assertEqual(count, 5)
99+
100+
async def test_count_values_list_filter_group_by(self):
101+
ret = (
102+
await Book.annotate(count=Count("id"))
103+
.filter(count__gt=6)
104+
.group_by("author_id")
105+
.values_list("author_id", "count")
106+
)
107+
self.assertEqual(len(ret), 1)
108+
self.assertEqual(ret[0][1], 10)
109+
110+
async def test_sum_values_list_group_by(self):
111+
ret = (
112+
await Book.annotate(sum=Sum("rating"))
113+
.group_by("author_id")
114+
.values_list("author_id", "sum")
115+
)
116+
for item in ret:
117+
author_id = item[0]
118+
sum_ = item[1]
119+
if author_id == self.a1.pk:
120+
self.assertEqual(sum_, 45.0)
121+
elif author_id == self.a2.pk:
122+
self.assertEqual(sum_, 10.0)
123+
124+
async def test_sum_values_list_filter_group_by(self):
125+
ret = (
126+
await Book.annotate(sum=Sum("rating"))
127+
.filter(sum__gt=11)
128+
.group_by("author_id")
129+
.values_list("author_id", "sum")
130+
)
131+
self.assertEqual(len(ret), 1)
132+
self.assertEqual(ret[0][1], 45.0)
133+
134+
async def test_avg_values_list_group_by(self):
135+
ret = (
136+
await Book.annotate(avg=Avg("rating"))
137+
.group_by("author_id")
138+
.values_list("author_id", "avg")
139+
)
140+
141+
for item in ret:
142+
author_id = item[0]
143+
avg = item[1]
144+
if author_id == self.a1.pk:
145+
self.assertEqual(avg, 4.5)
146+
elif author_id == self.a2.pk:
147+
self.assertEqual(avg, 2.0)
148+
149+
async def test_avg_values_list_filter_group_by(self):
150+
ret = (
151+
await Book.annotate(avg=Avg("rating"))
152+
.filter(avg__gt=3)
153+
.group_by("author_id")
154+
.values_list("author_id", "avg")
155+
)
156+
self.assertEqual(len(ret), 1)
157+
self.assertEqual(ret[0][1], 4.5)
158+
159+
async def test_implicit_group_by(self):
160+
ret = await Author.annotate(count=Count("books")).filter(count__gt=6)
161+
self.assertEqual(ret[0].count, 10)

‎tortoise/expressions.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
from typing import TYPE_CHECKING, Optional, Tuple, Type, Union
1+
from typing import TYPE_CHECKING, Optional, Tuple, Type
22

33
from pypika import Field
4-
from pypika.terms import ArithmeticExpression
4+
from pypika.terms import ArithmeticExpression, Term
55

66
from tortoise.exceptions import FieldError
77

@@ -12,10 +12,8 @@
1212
class F(Field): # type: ignore
1313
@classmethod
1414
def resolver_arithmetic_expression(
15-
cls,
16-
model: "Type[Model]",
17-
arithmetic_expression_or_field: Union[ArithmeticExpression, Field],
18-
) -> Tuple[Union[ArithmeticExpression, Field], Optional[Field]]:
15+
cls, model: "Type[Model]", arithmetic_expression_or_field: Term,
16+
) -> Tuple[Term, Optional[Field]]:
1917
field_object = None
2018

2119
if isinstance(arithmetic_expression_or_field, Field):

0 commit comments

Comments
 (0)