Python: adds JSON.ARRPOP command

Signed-off-by: Shoham Elias <shohame@amazon.com>

python/python/glide/__init__.py

@@ -33,6 +33,7 @@
     UpdateOptions,
 )
 from glide.async_commands.server_modules import json
+from glide.async_commands.server_modules.json import JsonArrPopOptions, JsonGetOptions
 from glide.async_commands.sorted_set import (
     AggregationType,
     GeoSearchByBox,
@@ -184,7 +185,6 @@
     "InfBound",
     "InfoSection",
     "InsertPosition",
-    "json",
     "LexBoundary",
     "Limit",
     "ListDirection",
@@ -212,6 +212,10 @@
     "ClusterScanCursor"
     # PubSub
     "PubSubMsg",
+    # Json
+    "json",
+    "JsonGetOptions",
+    "JsonArrPopOptions",
     # Logger
     "Logger",
     "LogLevel",

python/python/glide/async_commands/server_modules/json.py

@@ -54,6 +54,36 @@ def get_options(self) -> List[str]:
         return args
 
 
+from typing import List, Optional, Union
+
+
+class JsonArrPopOptions:
+    """
+    Options for the JSON.ARRPOP command.
+
+    Args:
+        path (TEncodable): The JSON path within the document from which to pop an element.
+        index (Optional[int]): The index of the element to pop. If not specified, will pop the last element.
+            Out of boundary indexes are rounded to their respective array boundaries. Defaults to None.
+    """
+
+    def __init__(self, path: TEncodable, index: Optional[int] = None):
+        self.path = path
+        self.index = index
+
+    def get_options(self) -> List[TEncodable]:
+        """
+        Get the options as a list of arguments for the JSON.ARRPOP command.
+
+        Returns:
+            List[TEncodable]: A list containing the path and, if specified, the index.
+        """
+        args = [self.path]
+        if self.index is not None:
+            args.append(str(self.index))
+        return args
+
+
 async def set(
     client: TGlideClient,
     key: TEncodable,
@@ -139,6 +169,66 @@ async def get(
     return cast(bytes, await client.custom_command(args))
 
 
+async def arrpop(
+    client: TGlideClient,
+    key: TEncodable,
+    options: Optional[JsonArrPopOptions] = None,
+) -> Optional[TJsonResponse[bytes]]:
+    """
+    Pops the last element from the array located at the specified path within the JSON document stored at `key`.
+    If `options.index` is provided, it pops the element at that index instead of the last element.
+
+    See https://valkey.io/commands/json.arrpop/ for more details.
+
+    Args:
+        client (TGlideClient): The client to execute the command.
+        key (TEncodable): The key of the JSON document.
+        options (Optional[JsonArrPopOptions]): Options including the path and optional index. See `JsonArrPopOptions`.
+
+    Returns:
+        Optional[TJsonResponse[bytes]]:
+            For JSONPath (`path` starts with `$`):
+                Returns a list of bytes string replies for every possible path, representing the popped JSON values,
+                or None for JSON values matching the path that are not an array or are an empty array.
+                If a value is not an array, its corresponding return value is null.
+            For legacy path (`path` starts with `.`):
+                Returns a bytes string representing the popped JSON value, or None if the array at `path` is empty.
+                If multiple paths match, the value from the first matching array that is not empty is returned.
+                If the JSON value at `path` is not a array or if `path` doesn't exist, an error is raised.
+            If `key` doesn't exist, an error is raised.
+
+    Examples:
+        >>> from glide import json
+        >>> await json.set(client, "doc", "$", '{"a": [1, 2, true], "b": {"a": [3, 4, ["value", 3, false], 5], "c": {"a": 42}}}')
+            b'OK'  # JSON is successfully set
+        >>> await json.arrpop(client, "doc", JsonArrPopOptions(path="$.a", index=1))
+            [b'2']  # Pop second element from array at path $.a
+        >>> await json.arrpop(client, "doc", JsonArrPopOptions(path="$..a"))
+            [b'true', b'5', None]  # Pop last elements from all arrays matching path `..a`
+
+        #### Using a legacy path (..) to pop the first matching array
+        >>> await json.arrpop(client, "doc", JsonArrPopOptions(path="..a"))
+            b"1"  # First match popped (from array at path $.a)
+
+        #### Even though only one value is returned from `..a`, subsequent arrays are also affected
+        >>> await json.get(client, "doc", "$..a")
+            b"[[], [3, 4], 42]"  # Remaining elements after pop show the changes
+
+        >>> await json.set(client, "doc", "$", '[[], ["a"], ["a", "b", "c"]]')
+            b'OK'  # JSON is successfully set
+        >>> await json.arrpop(client, "doc", JsonArrPopOptions(path=".", index=-1))
+            b'["a","b","c"]'  # Pop last elements at path `.`
+    """
+    args = ["JSON.ARRPOP", key]
+    if options:
+        args.extend(options.get_options())
+
+    return cast(
+        Optional[TJsonResponse[bytes]],
+        await client.custom_command(args),
+    )
+
+
 async def delete(
     client: TGlideClient,
     key: TEncodable,

python/python/tests/tests_server_modules/test_json.py

@@ -5,7 +5,7 @@
 import pytest
 from glide.async_commands.core import ConditionalChange, InfoSection
 from glide.async_commands.server_modules import json
-from glide.async_commands.server_modules.json import JsonGetOptions
+from glide.async_commands.server_modules.json import JsonArrPopOptions, JsonGetOptions
 from glide.config import ProtocolVersion
 from glide.constants import OK
 from glide.exceptions import RequestError
@@ -196,3 +196,76 @@ async def test_json_toggle(self, glide_client: TGlideClient):
 
         with pytest.raises(RequestError):
             assert await json.toggle(glide_client, "non_exiting_key", "$")
+
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_json_arrpop(self, glide_client: TGlideClient):
+        key = get_random_string(5)
+        key2 = get_random_string(5)
+
+        json_value = '{"a": [1, 2, true], "b": {"a": [3, 4, ["value", 3, false] ,5], "c": {"a": 42}}}'
+        assert await json.set(glide_client, key, "$", json_value) == OK
+
+        assert await json.arrpop(
+            glide_client, key, JsonArrPopOptions(path="$.a", index=1)
+        ) == [b"2"]
+        assert (
+            await json.arrpop(glide_client, key, JsonArrPopOptions(path="$..a"))
+        ) == [b"true", b"5", None]
+
+        assert (
+            await json.arrpop(glide_client, key, JsonArrPopOptions(path="..a")) == b"1"
+        )
+        # Even if only one array element was returned, ensure second array at `..a` was popped
+        assert await json.get(glide_client, key, "$..a") == b"[[],[3,4],42]"
+
+        # Out of index
+        assert await json.arrpop(
+            glide_client, key, JsonArrPopOptions(path="$..a", index=10)
+        ) == [None, b"4", None]
+
+        assert (
+            await json.arrpop(
+                glide_client, key, JsonArrPopOptions(path="..a", index=-10)
+            )
+            == b"3"
+        )
+
+        # Path is not an array
+        assert await json.arrpop(glide_client, key, JsonArrPopOptions(path="$")) == [
+            None
+        ]
+        with pytest.raises(RequestError):
+            assert await json.arrpop(glide_client, key, JsonArrPopOptions(path="."))
+        with pytest.raises(RequestError):
+            assert await json.arrpop(glide_client, key)
+
+        # Non existing path
+        assert (
+            await json.arrpop(
+                glide_client, key, JsonArrPopOptions(path="$.non_existing_path")
+            )
+            == []
+        )
+        with pytest.raises(RequestError):
+            assert await json.arrpop(
+                glide_client, key, JsonArrPopOptions(path="non_existing_path")
+            )
+
+        with pytest.raises(RequestError):
+            await json.arrpop(
+                glide_client, "non_existing_key", JsonArrPopOptions(path="$.a")
+            )
+        with pytest.raises(RequestError):
+            await json.arrpop(
+                glide_client, "non_existing_key", JsonArrPopOptions(path=".a")
+            )
+
+        assert (
+            await json.set(glide_client, key2, "$", '[[], ["a"], ["a", "b", "c"]]')
+            == OK
+        )
+        assert (
+            await json.arrpop(glide_client, key2, JsonArrPopOptions(path=".", index=-1))
+            == b'["a","b","c"]'
+        )