1 files changed, 95 insertions, 0 deletions
diff --git a/redis/asyncio/parser.py b/redis/asyncio/parser.py
new file mode 100644
index 0000000..273fe03
--- /dev/null
+++ b/redis/asyncio/parser.py
@@ -0,0 +1,95 @@
+from typing import TYPE_CHECKING, List, Optional, Union
+
+from redis.exceptions import RedisError, ResponseError
+
+if TYPE_CHECKING:
+    from redis.asyncio.cluster import ClusterNode
+
+
+class CommandsParser:
+    """
+    Parses Redis commands to get command keys.
+
+    COMMAND output is used to determine key locations.
+    Commands that do not have a predefined key location are flagged with 'movablekeys',
+    and these commands' keys are determined by the command 'COMMAND GETKEYS'.
+
+    NOTE: Due to a bug in redis<7.0, this does not work properly
+    for EVAL or EVALSHA when the `numkeys` arg is 0.
+     - issue: https://github.com/redis/redis/issues/9493
+     - fix: https://github.com/redis/redis/pull/9733
+
+    So, don't use this with EVAL or EVALSHA.
+    """
+
+    __slots__ = ("commands",)
+
+    def __init__(self) -> None:
+        self.commands = {}
+
+    async def initialize(self, r: "ClusterNode") -> None:
+        commands = await r.execute_command("COMMAND")
+        for cmd, command in commands.items():
+            if "movablekeys" in command["flags"]:
+                commands[cmd] = -1
+            elif command["first_key_pos"] == 0 and command["last_key_pos"] == 0:
+                commands[cmd] = 0
+            elif command["first_key_pos"] == 1 and command["last_key_pos"] == 1:
+                commands[cmd] = 1
+        self.commands = {cmd.upper(): command for cmd, command in commands.items()}
+
+    # As soon as this PR is merged into Redis, we should reimplement
+    # our logic to use COMMAND INFO changes to determine the key positions
+    # https://github.com/redis/redis/pull/8324
+    async def get_keys(
+        self, redis_conn: "ClusterNode", *args
+    ) -> Optional[Union[List[str], List[bytes]]]:
+        if len(args) < 2:
+            # The command has no keys in it
+            return None
+
+        try:
+            command = self.commands[args[0]]
+        except KeyError:
+            # try to split the command name and to take only the main command
+            # e.g. 'memory' for 'memory usage'
+            args = args[0].split() + list(args[1:])
+            cmd_name = args[0]
+            if cmd_name not in self.commands:
+                # We'll try to reinitialize the commands cache, if the engine
+                # version has changed, the commands may not be current
+                await self.initialize(redis_conn)
+                if cmd_name not in self.commands:
+                    raise RedisError(
+                        f"{cmd_name.upper()} command doesn't exist in Redis commands"
+                    )
+
+            command = self.commands[cmd_name]
+
+        if command == 1:
+            return [args[1]]
+        if command == 0:
+            return None
+        if command == -1:
+            return await self._get_moveable_keys(redis_conn, *args)
+
+        last_key_pos = command["last_key_pos"]
+        if last_key_pos < 0:
+            last_key_pos = len(args) + last_key_pos
+        return args[command["first_key_pos"] : last_key_pos + 1 : command["step_count"]]
+
+    async def _get_moveable_keys(
+        self, redis_conn: "ClusterNode", *args
+    ) -> Optional[List[str]]:
+        try:
+            keys = await redis_conn.execute_command("COMMAND GETKEYS", *args)
+        except ResponseError as e:
+            message = e.__str__()
+            if (
+                "Invalid arguments" in message
+                or "The command has no key arguments" in message
+            ):
+                return None
+            else:
+                raise e
+        return keys