Add .group_by(...) support (#363)

CHANGELOG.rst

@@ -9,8 +9,10 @@ Changelog
 
 0.16
 ====
+
 0.16.8
 ------
+- Add ``group by`` support
 - Fixed regression where ``GROUP BY`` class is missing for an aggregate with a specified order.
 
 0.16.7

docs/examples/basic.rst

@@ -53,6 +53,15 @@ Functions
 
 .. _example_schema_create:
 
+Group By
+========
+.. literalinclude::  ../../examples/group_by.py
+
+
+.. rst-class:: html-toggle
+
+.. _example_aggregation:
+
 Schema creation
 ===============
 .. literalinclude::  ../../examples/schema_create.py

examples/group_by.py

@@ -0,0 +1,76 @@
+from tortoise import Model, Tortoise, fields, run_async
+from tortoise.functions import Avg, Count, Sum
+
+
+class Author(Model):
+    name = fields.CharField(max_length=255)
+
+
+class Book(Model):
+    name = fields.CharField(max_length=255)
+    author = fields.ForeignKeyField("models.Author", related_name="books")
+    rating = fields.FloatField()
+
+
+async def run():
+    await Tortoise.init(db_url="sqlite://:memory:", modules={"models": ["__main__"]})
+    await Tortoise.generate_schemas()
+
+    a1 = await Author.create(name="author1")
+    a2 = await Author.create(name="author2")
+    for i in range(10):
+        await Book.create(name=f"book{i}", author=a1, rating=i)
+    for i in range(5):
+        await Book.create(name=f"book{i}", author=a2, rating=i)
+
+    ret = await Book.annotate(count=Count("id")).group_by("author_id").values("author_id", "count")
+    print(ret)
+    # >>> [{'author_id': 1, 'count': 10}, {'author_id': 2, 'count': 5}]
+
+    ret = (
+        await Book.annotate(count=Count("id"))
+        .filter(count__gt=6)
+        .group_by("author_id")
+        .values("author_id", "count")
+    )
+    print(ret)
+    # >>> [{'author_id': 1, 'count': 10}]
+
+    ret = await Book.annotate(sum=Sum("rating")).group_by("author_id").values("author_id", "sum")
+    print(ret)
+    # >>> [{'author_id': 1, 'sum': 45.0}, {'author_id': 2, 'sum': 10.0}]
+
+    ret = (
+        await Book.annotate(sum=Sum("rating"))
+        .filter(sum__gt=11)
+        .group_by("author_id")
+        .values("author_id", "sum")
+    )
+    print(ret)
+    # >>> [{'author_id': 1, 'sum': 45.0}]
+
+    ret = await Book.annotate(avg=Avg("rating")).group_by("author_id").values("author_id", "avg")
+    print(ret)
+    # >>> [{'author_id': 1, 'avg': 4.5}, {'author_id': 2, 'avg': 2.0}]
+
+    ret = (
+        await Book.annotate(avg=Avg("rating"))
+        .filter(avg__gt=3)
+        .group_by("author_id")
+        .values("author_id", "avg")
+    )
+    print(ret)
+    # >>> [{'author_id': 1, 'avg': 4.5}]
+
+    # and use .values_list()
+    ret = (
+        await Book.annotate(count=Count("id"))
+        .group_by("author_id")
+        .values_list("author_id", "count")
+    )
+    print(ret)
+    # >>> [(1, 10), (2, 5)]
+
+
+if __name__ == "__main__":
+    run_async(run())

tests/test_group_by.py

@@ -0,0 +1,161 @@
+from tests.testmodels import Author, Book
+from tortoise.contrib import test
+from tortoise.functions import Avg, Count, Sum
+
+
+class TestGroupBy(test.TestCase):
+    async def setUp(self) -> None:
+        self.a1 = await Author.create(name="author1")
+        self.a2 = await Author.create(name="author2")
+        for i in range(10):
+            await Book.create(name=f"book{i}", author=self.a1, rating=i)
+        for i in range(5):
+            await Book.create(name=f"book{i}", author=self.a2, rating=i)
+
+    async def test_count_group_by(self):
+        ret = (
+            await Book.annotate(count=Count("id"))
+            .group_by("author_id")
+            .values("author_id", "count")
+        )
+
+        for item in ret:
+            author_id = item.get("author_id")
+            count = item.get("count")
+            if author_id == self.a1.pk:
+                self.assertEqual(count, 10)
+            elif author_id == self.a2.pk:
+                self.assertEqual(count, 5)
+
+    async def test_count_filter_group_by(self):
+        ret = (
+            await Book.annotate(count=Count("id"))
+            .filter(count__gt=6)
+            .group_by("author_id")
+            .values("author_id", "count")
+        )
+        self.assertEqual(len(ret), 1)
+        self.assertEqual(ret[0].get("count"), 10)
+
+    async def test_sum_group_by(self):
+        ret = (
+            await Book.annotate(sum=Sum("rating")).group_by("author_id").values("author_id", "sum")
+        )
+        for item in ret:
+            author_id = item.get("author_id")
+            sum_ = item.get("sum")
+            if author_id == self.a1.pk:
+                self.assertEqual(sum_, 45.0)
+            elif author_id == self.a2.pk:
+                self.assertEqual(sum_, 10.0)
+
+    async def test_sum_filter_group_by(self):
+        ret = (
+            await Book.annotate(sum=Sum("rating"))
+            .filter(sum__gt=11)
+            .group_by("author_id")
+            .values("author_id", "sum")
+        )
+        self.assertEqual(len(ret), 1)
+        self.assertEqual(ret[0].get("sum"), 45.0)
+
+    async def test_avg_group_by(self):
+        ret = (
+            await Book.annotate(avg=Avg("rating")).group_by("author_id").values("author_id", "avg")
+        )
+
+        for item in ret:
+            author_id = item.get("author_id")
+            avg = item.get("avg")
+            if author_id == self.a1.pk:
+                self.assertEqual(avg, 4.5)
+            elif author_id == self.a2.pk:
+                self.assertEqual(avg, 2.0)
+
+    async def test_avg_filter_group_by(self):
+        ret = (
+            await Book.annotate(avg=Avg("rating"))
+            .filter(avg__gt=3)
+            .group_by("author_id")
+            .values_list("author_id", "avg")
+        )
+        self.assertEqual(len(ret), 1)
+        self.assertEqual(ret[0][1], 4.5)
+
+    async def test_count_values_list_group_by(self):
+        ret = (
+            await Book.annotate(count=Count("id"))
+            .group_by("author_id")
+            .values_list("author_id", "count")
+        )
+
+        for item in ret:
+            author_id = item[0]
+            count = item[1]
+            if author_id == self.a1.pk:
+                self.assertEqual(count, 10)
+            elif author_id == self.a2.pk:
+                self.assertEqual(count, 5)
+
+    async def test_count_values_list_filter_group_by(self):
+        ret = (
+            await Book.annotate(count=Count("id"))
+            .filter(count__gt=6)
+            .group_by("author_id")
+            .values_list("author_id", "count")
+        )
+        self.assertEqual(len(ret), 1)
+        self.assertEqual(ret[0][1], 10)
+
+    async def test_sum_values_list_group_by(self):
+        ret = (
+            await Book.annotate(sum=Sum("rating"))
+            .group_by("author_id")
+            .values_list("author_id", "sum")
+        )
+        for item in ret:
+            author_id = item[0]
+            sum_ = item[1]
+            if author_id == self.a1.pk:
+                self.assertEqual(sum_, 45.0)
+            elif author_id == self.a2.pk:
+                self.assertEqual(sum_, 10.0)
+
+    async def test_sum_values_list_filter_group_by(self):
+        ret = (
+            await Book.annotate(sum=Sum("rating"))
+            .filter(sum__gt=11)
+            .group_by("author_id")
+            .values_list("author_id", "sum")
+        )
+        self.assertEqual(len(ret), 1)
+        self.assertEqual(ret[0][1], 45.0)
+
+    async def test_avg_values_list_group_by(self):
+        ret = (
+            await Book.annotate(avg=Avg("rating"))
+            .group_by("author_id")
+            .values_list("author_id", "avg")
+        )
+
+        for item in ret:
+            author_id = item[0]
+            avg = item[1]
+            if author_id == self.a1.pk:
+                self.assertEqual(avg, 4.5)
+            elif author_id == self.a2.pk:
+                self.assertEqual(avg, 2.0)
+
+    async def test_avg_values_list_filter_group_by(self):
+        ret = (
+            await Book.annotate(avg=Avg("rating"))
+            .filter(avg__gt=3)
+            .group_by("author_id")
+            .values_list("author_id", "avg")
+        )
+        self.assertEqual(len(ret), 1)
+        self.assertEqual(ret[0][1], 4.5)
+
+    async def test_implicit_group_by(self):
+        ret = await Author.annotate(count=Count("books")).filter(count__gt=6)
+        self.assertEqual(ret[0].count, 10)

tortoise/expressions.py

@@ -1,7 +1,7 @@
-from typing import TYPE_CHECKING, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Optional, Tuple, Type
 
 from pypika import Field
-from pypika.terms import ArithmeticExpression
+from pypika.terms import ArithmeticExpression, Term
 
 from tortoise.exceptions import FieldError
 
@@ -12,10 +12,8 @@
 class F(Field):  # type: ignore
     @classmethod
     def resolver_arithmetic_expression(
-        cls,
-        model: "Type[Model]",
-        arithmetic_expression_or_field: Union[ArithmeticExpression, Field],
-    ) -> Tuple[Union[ArithmeticExpression, Field], Optional[Field]]:
+        cls, model: "Type[Model]", arithmetic_expression_or_field: Term,
+    ) -> Tuple[Term, Optional[Field]]:
         field_object = None
 
         if isinstance(arithmetic_expression_or_field, Field):