diff --git a/tests/fields/test_array.py b/tests/fields/test_array.py index a4fafa541..58ec1a68f 100644 --- a/tests/fields/test_array.py +++ b/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), []) diff --git a/tortoise/backends/base_postgres/executor.py b/tortoise/backends/base_postgres/executor.py index 9720b8a4c..ccf02dbc2 100644 --- a/tortoise/backends/base_postgres/executor.py +++ b/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( diff --git a/tortoise/contrib/postgres/array_functions.py b/tortoise/contrib/postgres/array_functions.py index c6247978e..0d84b6d16 100644 --- a/tortoise/contrib/postgres/array_functions.py +++ b/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) diff --git a/tortoise/filters.py b/tortoise/filters.py index e4f49550b..de9afdbd7 100644 --- a/tortoise/filters.py +++ b/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, + }, }