Skip to content

Commit

Permalink
Implement __len filter for ArrayField
Browse files Browse the repository at this point in the history
  • Loading branch information
henadzit committed Feb 5, 2025
1 parent 2c32925 commit 74b55f5
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 0 deletions.
14 changes: 14 additions & 0 deletions tests/fields/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -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), [])
3 changes: 3 additions & 0 deletions tortoise/backends/base_postgres/executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -32,6 +33,7 @@
json_filter,
posix_regex,
search,
array_length,
)


Expand All @@ -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(
Expand Down
6 changes: 6 additions & 0 deletions tortoise/contrib/postgres/array_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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)
14 changes: 14 additions & 0 deletions tortoise/filters.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
##############################################################################
Expand Down Expand Up @@ -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,
},
}


Expand Down

0 comments on commit 74b55f5

Please sign in to comment.