Merge 6f7bac9 into 9673ccd

CHANGELOG.rst

@@ -11,7 +11,8 @@ Changelog
 ====
 0.16.17
 -------
-- Add `on_delete` in `ManyToManyField`.
+- Add `on_delete` in `ManyToManyField`. (#508)
+- Support `F` expression in `annotate`. (#475)
 
 0.16.16
 -------

tests/test_queryset.py

@@ -317,6 +317,15 @@ async def test_order_by_bad_value(self):
         with self.assertRaisesRegex(FieldError, "Unknown field badid for model IntFields"):
             await IntFields.all().order_by("badid").values_list()
 
+    async def test_annotate_order_expression(self):
+        data = (
+            await IntFields.annotate(idp=F("id") + 1)
+            .order_by("-idp")
+            .first()
+            .values_list("id", "idp")
+        )[0]
+        self.assertEqual(data[0] + 1, data[1])
+
     async def test_get_raw_sql(self):
         sql = IntFields.all().sql()
         self.assertRegex(sql, r"^SELECT.+FROM.+")

tortoise/models.py

@@ -16,9 +16,11 @@
     Tuple,
     Type,
     TypeVar,
+    Union,
 )
 
 from pypika import Order, Query, Table
+from pypika.terms import Term
 
 from tortoise.backends.base.client import BaseDBAsyncClient
 from tortoise.exceptions import (
@@ -1041,9 +1043,9 @@ def exclude(cls: Type[MODEL], *args: Q, **kwargs: Any) -> QuerySet[MODEL]:
         return QuerySet(cls).exclude(*args, **kwargs)
 
     @classmethod
-    def annotate(cls: Type[MODEL], **kwargs: Function) -> QuerySet[MODEL]:
+    def annotate(cls: Type[MODEL], **kwargs: Union[Function, Term]) -> QuerySet[MODEL]:
         """
-        Annotates the result set with extra Functions/Aggregations.
+        Annotates the result set with extra Functions/Aggregations/Expressions.
 
         :param kwargs: Parameter name and the Function/Aggregation to annotate with.
         """

tortoise/queryset.py

@@ -184,8 +184,11 @@ def resolve_ordering(
                 )
             elif field_name in annotations:
                 annotation = annotations[field_name]
-                annotation_info = annotation.resolve(self.model, table)
-                self.query = self.query.orderby(annotation_info["field"], order=ordering[1])
+                if isinstance(annotation, Term):
+                    self.query = self.query.orderby(annotation, order=ordering[1])
+                else:
+                    annotation_info = annotation.resolve(self.model, table)
+                    self.query = self.query.orderby(annotation_info["field"], order=ordering[1])
             else:
                 field_object = model._meta.fields_map.get(field_name)
 
@@ -209,7 +212,10 @@ def _resolve_annotate(self) -> bool:
         table = self.model._meta.basetable
         annotation_info = {}
         for key, annotation in self._annotations.items():
-            annotation_info[key] = annotation.resolve(self.model, table)
+            if isinstance(annotation, Term):
+                annotation_info[key] = {"joins": [], "field": annotation}
+            else:
+                annotation_info[key] = annotation.resolve(self.model, table)
 
         for key, info in annotation_info.items():
             for join in info["joins"]:
@@ -429,8 +435,8 @@ def annotate(self, **kwargs: Function) -> "QuerySet[MODEL]":
 
         queryset = self._clone()
         for key, annotation in kwargs.items():
-            if not isinstance(annotation, Function):
-                raise TypeError("value is expected to be Function instance")
+            if not isinstance(annotation, (Function, Term)):
+                raise TypeError("value is expected to be Function/Term instance")
             queryset._annotations[key] = annotation
             queryset._custom_filters.update(get_filters_for_field(key, None, key))
         return queryset
@@ -1076,7 +1082,8 @@ def resolve_to_python_value(self, model: Type[MODEL], field: str) -> Callable:
             return lambda x: x
 
         if field in self.annotations:
-            field_object = self.annotations[field].field_object
+            annotation = self.annotations[field]
+            field_object = getattr(annotation, "field_object", None)
             if field_object:
                 return field_object.to_python_value
             return lambda x: x