Implement __len filter for ArrayField

Author: henadzit

tests/fields/test_array.py

@@ -135,3 +135,17 @@ async def test_overlap_ints(self):
             "array", flat=True
         )
         self.assertEqual(sorted(list(found)), [[1, 2, 3], [2, 3, 4], [3, 4, 5]])
+
+    async def test_array_length(self):
+        await testmodels.ArrayFields.create(array=[1, 2, 3])
+        await testmodels.ArrayFields.create(array=[1])
+        await testmodels.ArrayFields.create(array=[1, 2])
+
+        found = await testmodels.ArrayFields.filter(array__len=3).values_list("array", flat=True)
+        self.assertEqual(list(found), [[1, 2, 3]])
+
+        found = await testmodels.ArrayFields.filter(array__len=1).values_list("array", flat=True)
+        self.assertEqual(list(found), [[1]])
+
+        found = await testmodels.ArrayFields.filter(array__len=0).values_list("array", flat=True)
+        self.assertEqual(list(found), [])

tortoise/backends/base_postgres/executor.py

@@ -11,6 +11,7 @@
     postgres_array_contains,
     postgres_array_contained_by,
     postgres_array_overlap,
+    postgres_array_length,
 )
 from tortoise.contrib.postgres.json_functions import (
     postgres_json_contained_by,
@@ -32,6 +33,7 @@
     json_filter,
     posix_regex,
     search,
+    array_length,
 )
 
 
@@ -52,6 +54,7 @@ class BasePostgresExecutor(BaseExecutor):
         json_filter: postgres_json_filter,
         posix_regex: postgres_posix_regex,
         insensitive_posix_regex: postgres_insensitive_posix_regex,
+        array_length: postgres_array_length,
     }
 
     def _prepare_insert_statement(

tortoise/contrib/postgres/array_functions.py

@@ -2,6 +2,7 @@
 from typing import Any, Sequence, Union
 
 from pypika_tortoise.terms import Array, BasicCriterion, Criterion, Term
+from pypika_tortoise.functions import Function
 
 
 class PostgresArrayOperators(str, Enum):
@@ -23,3 +24,8 @@ def postgres_array_contained_by(field: Term, value: Sequence[Any]) -> Criterion:
 
 def postgres_array_overlap(field: Term, value: Sequence[Any]) -> Criterion:
     return BasicCriterion(PostgresArrayOperators.OVERLAP, field, Array(*value))
+
+
+def postgres_array_length(field: Term, value: int) -> Criterion:
+    """Returns a criterion that checks if array length equals the given value"""
+    return Function("array_length", field, 1).eq(value)

tortoise/filters.py

@@ -73,6 +73,10 @@ def string_encoder(value: Any, instance: "Model", field: Field) -> str:
     return str(value)
 
 
+def int_encoder(value: Any, instance: "Model", field: Field) -> int:
+    return int(value)
+
+
 def json_encoder(value: Any, instance: "Model", field: Field) -> dict:
     return value
 
@@ -238,6 +242,10 @@ def array_overlap(field: Term, value: Union[Any, Sequence[Any]]) -> Criterion:
     raise NotImplementedError("must be overridden in each executor")
 
 
+def array_length(field: Term, value: int) -> Criterion:
+    raise NotImplementedError("must be overridden in each executor")
+
+
 ##############################################################################
 # Filter resolvers
 ##############################################################################
@@ -430,6 +438,12 @@ def get_array_filter(field_name: str, source_field: str) -> dict[str, FilterInfo
             "source_field": source_field,
             "operator": array_overlap,
         },
+        f"{field_name}__len": {
+            "field": field_name,
+            "source_field": source_field,
+            "operator": array_length,
+            "value_encoder": int_encoder,
+        },
     }