Skip to content

Commit

Permalink
Python: adds JSON.ARRTRIM command (valkey-io#2457)
Browse files Browse the repository at this point in the history
---------

Signed-off-by: Shoham Elias <shohame@amazon.com>
  • Loading branch information
shohamazon authored Oct 27, 2024
1 parent d95040c commit c197f1d
Show file tree
Hide file tree
Showing 3 changed files with 159 additions and 0 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
* Python: Add `JSON.STRAPPEND` , `JSON.STRLEN` commands ([#2372](https://github.com/valkey-io/valkey-glide/pull/2372))
* Python: Add `JSON.OBJKEYS` command ([#2395](https://github.com/valkey-io/valkey-glide/pull/2395))
* Python: Add `JSON.ARRINSERT` command ([#2464](https://github.com/valkey-io/valkey-glide/pull/2464))
* Python: Add `JSON.ARRTRIM` command ([#2457](https://github.com/valkey-io/valkey-glide/pull/2457))

#### Breaking Changes

Expand Down
55 changes: 55 additions & 0 deletions python/python/glide/async_commands/server_modules/json.py
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,61 @@ async def arrlen(
)


async def arrtrim(
client: TGlideClient,
key: TEncodable,
path: TEncodable,
start: int,
end: int,
) -> TJsonResponse[int]:
"""
Trims an array at the specified `path` within the JSON document stored at `key` so that it becomes a subarray [start, end], both inclusive.›
If `start` < 0, it is treated as 0.
If `end` >= size (size of the array), it is treated as size-1.
If `start` >= size or `start` > `end`, the array is emptied and 0 is returned.
Args:
client (TGlideClient): The client to execute the command.
key (TEncodable): The key of the JSON document.
path (TEncodable): The path within the JSON document.
start (int): The start index, inclusive.
end (int): The end index, inclusive.
Returns:
TJsonResponse[int]:
For JSONPath (`path` starts with '$'):
Returns a list of integer replies for every possible path, indicating the new length of the array, or None for JSON values matching the path that are not an array.
If a value is an empty array, its corresponding return value is 0.
If `path` doesn't exist, an empty array will be returned.
For legacy path (`path` doesn't starts with `$`):
Returns an integer representing the new length of the array.
If the array is empty, returns 0.
If multiple paths match, the length of the first trimmed array match is returned.
If `path` doesn't exist, or the value at `path` is not an array, an error is raised.
If `key` doesn't exist, an error is raised.
Examples:
>>> from glide import json
>>> await json.set(client, "doc", "$", '[[], ["a"], ["a", "b"], ["a", "b", "c"]]')
'OK'
>>> await json.arrtrim(client, "doc", "$[*]", 0, 1)
[0, 1, 2, 2]
>>> await json.get(client, "doc")
b'[[],[\"a\"],[\"a\",\"b\"],[\"a\",\"b\"]]'
>>> await json.set(client, "doc", "$", '{"children": ["John", "Jack", "Tom", "Bob", "Mike"]}')
'OK'
>>> await json.arrtrim(client, "doc", ".children", 0, 1)
2
>>> await json.get(client, "doc", ".children")
b'["John","Jack"]'
"""
return cast(
TJsonResponse[int],
await client.custom_command(["JSON.ARRTRIM", key, path, str(start), str(end)]),
)


async def clear(
client: TGlideClient,
key: TEncodable,
Expand Down
103 changes: 103 additions & 0 deletions python/python/tests/tests_server_modules/test_json.py
Original file line number Diff line number Diff line change
Expand Up @@ -1191,3 +1191,106 @@ async def test_json_debug_memory(self, glide_client: TGlideClient):
# Test for non-existent key
result = await json.debug_memory(glide_client, "non_existent_key", ".key10")
assert result == None

@pytest.mark.parametrize("cluster_mode", [True, False])
@typing.no_type_check
@pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
async def test_json_arrtrim(self, glide_client: TGlideClient):
key = get_random_string(5)

# Test with enhanced path syntax
json_value = '{"a": [0, 1, 2, 3, 4, 5, 6, 7, 8], "b": {"a": [0, 9, 10, 11, 12, 13], "c": {"a": 42}}}'
assert await json.set(glide_client, key, "$", json_value) == OK

# Basic trim
assert await json.arrtrim(glide_client, key, "$..a", 1, 7) == [7, 5, None]
assert OuterJson.loads(await json.get(glide_client, key, "$..a")) == [
[1, 2, 3, 4, 5, 6, 7],
[9, 10, 11, 12, 13],
42,
]

# Test negative start (should be treated as 0)
assert await json.arrtrim(glide_client, key, "$.a", -1, 5) == [6]
assert OuterJson.loads(await json.get(glide_client, key, "$.a")) == [
[1, 2, 3, 4, 5, 6]
]
assert await json.arrtrim(glide_client, key, ".a", -1, 5) == 6
assert OuterJson.loads(await json.get(glide_client, key, ".a")) == [
1,
2,
3,
4,
5,
6,
]

# Test end >= size (should be treated as size-1)
assert await json.arrtrim(glide_client, key, "$.a", 0, 10) == [6]
assert OuterJson.loads(await json.get(glide_client, key, "$.a")) == [
[1, 2, 3, 4, 5, 6]
]

assert await json.arrtrim(glide_client, key, ".a", 0, 10) == 6
assert OuterJson.loads(await json.get(glide_client, key, ".a")) == [
1,
2,
3,
4,
5,
6,
]

# Test start >= size (should empty the array)
assert await json.arrtrim(glide_client, key, "$.a", 7, 10) == [0]
assert OuterJson.loads(await json.get(glide_client, key, "$.a")) == [[]]

assert await json.set(glide_client, key, ".a", '["a", "b", "c"]') == OK
assert await json.arrtrim(glide_client, key, ".a", 7, 10) == 0
assert OuterJson.loads(await json.get(glide_client, key, ".a")) == []

# Test start > end (should empty the array)
assert await json.arrtrim(glide_client, key, "$..a", 2, 1) == [0, 0, None]
assert OuterJson.loads(await json.get(glide_client, key, "$..a")) == [
[],
[],
42,
]
assert await json.set(glide_client, key, "..a", '["a", "b", "c", "d"]') == OK
assert await json.arrtrim(glide_client, key, "..a", 2, 1) == 0
assert OuterJson.loads(await json.get(glide_client, key, ".a")) == []

# Multiple path match
assert await json.set(glide_client, key, "$", json_value) == OK
assert await json.arrtrim(glide_client, key, "..a", 1, 10) == 8
assert OuterJson.loads(await json.get(glide_client, key, "$..a")) == [
[1, 2, 3, 4, 5, 6, 7, 8],
[9, 10, 11, 12, 13],
42,
]

# Test with non-existent path
with pytest.raises(RequestError):
await json.arrtrim(glide_client, key, ".non_existent", 0, 1)

assert await json.arrtrim(glide_client, key, "$.non_existent", 0, 1) == []

# Test with non-array path
assert await json.arrtrim(glide_client, key, "$", 0, 1) == [None]

with pytest.raises(RequestError):
await json.arrtrim(glide_client, key, ".", 0, 1)

# Test with non-existent key
with pytest.raises(RequestError):
await json.arrtrim(glide_client, "non_existent_key", "$", 0, 1)

# Test with non-existent key
with pytest.raises(RequestError):
await json.arrtrim(glide_client, "non_existent_key", ".", 0, 1)

# Test empty array
assert await json.set(glide_client, key, "$.empty", "[]") == OK
assert await json.arrtrim(glide_client, key, "$.empty", 0, 1) == [0]
assert await json.arrtrim(glide_client, key, ".empty", 0, 1) == 0
assert OuterJson.loads(await json.get(glide_client, key, "$.empty")) == [[]]

0 comments on commit c197f1d

Please sign in to comment.