Add type hints for JSON commands

docs/examples/search_json_examples.ipynb

@@ -59,9 +59,9 @@
     "        \"city\": \"Tel Aviv\"\n",
     "    }\n",
     "}\n",
-    "r.json().set(\"user:1\", Path.rootPath(), user1)\n",
-    "r.json().set(\"user:2\", Path.rootPath(), user2)\n",
-    "r.json().set(\"user:3\", Path.rootPath(), user3)\n",
+    "r.json().set(\"user:1\", Path.root_path(), user1)\n",
+    "r.json().set(\"user:2\", Path.root_path(), user2)\n",
+    "r.json().set(\"user:3\", Path.root_path(), user3)\n",
     "\n",
     "schema = (TextField(\"$.user.name\", as_name=\"name\"),TagField(\"$.user.city\", as_name=\"city\"), NumericField(\"$.user.age\", as_name=\"age\"))\n",
     "r.ft().create_index(schema, definition=IndexDefinition(prefix=[\"user:\"], index_type=IndexType.JSON))"

redis/commands/json/_util.py

@@ -0,0 +1,3 @@
+from typing import Any, Dict, List, Union
+
+JsonType = Union[str, int, float, bool, None, Dict[str, Any], List[Any]]

redis/commands/json/commands.py

@@ -1,18 +1,22 @@
 import os
 from json import JSONDecodeError, loads
+from typing import Dict, List, Optional, Union
 
 from deprecated import deprecated
 
 from redis.exceptions import DataError
 
+from ._util import JsonType
 from .decoders import decode_dict_keys
 from .path import Path
 
 
 class JSONCommands:
     """json commands."""
 
-    def arrappend(self, name, path=Path.root_path(), *args):
+    def arrappend(
+        self, name: str, path: Optional[str] = Path.root_path(), *args: List[JsonType]
+    ) -> List[Union[int, None]]:
         """Append the objects ``args`` to the array under the
         ``path` in key ``name``.
 
@@ -23,7 +27,14 @@ def arrappend(self, name, path=Path.root_path(), *args):
             pieces.append(self._encode(o))
         return self.execute_command("JSON.ARRAPPEND", *pieces)
 
-    def arrindex(self, name, path, scalar, start=0, stop=-1):
+    def arrindex(
+        self,
+        name: str,
+        path: str,
+        scalar: int,
+        start: Optional[int] = 0,
+        stop: Optional[int] = -1,
+    ) -> List[Union[int, None]]:
         """
         Return the index of ``scalar`` in the JSON array under ``path`` at key
         ``name``.
@@ -37,7 +48,9 @@ def arrindex(self, name, path, scalar, start=0, stop=-1):
             "JSON.ARRINDEX", name, str(path), self._encode(scalar), start, stop
         )
 
-    def arrinsert(self, name, path, index, *args):
+    def arrinsert(
+        self, name: str, path: str, index: int, *args: List[JsonType]
+    ) -> List[Union[int, None]]:
         """Insert the objects ``args`` to the array at index ``index``
         under the ``path` in key ``name``.
 
@@ -48,61 +61,73 @@ def arrinsert(self, name, path, index, *args):
             pieces.append(self._encode(o))
         return self.execute_command("JSON.ARRINSERT", *pieces)
 
-    def arrlen(self, name, path=Path.root_path()):
+    def arrlen(
+        self, name: str, path: Optional[str] = Path.root_path()
+    ) -> List[Union[int, None]]:
         """Return the length of the array JSON value under ``path``
         at key``name``.
 
         For more information: https://oss.redis.com/redisjson/commands/#jsonarrlen
         """  # noqa
         return self.execute_command("JSON.ARRLEN", name, str(path))
 
-    def arrpop(self, name, path=Path.root_path(), index=-1):
+    def arrpop(
+        self,
+        name: str,
+        path: Optional[str] = Path.root_path(),
+        index: Optional[int] = -1,
+    ) -> List[Union[str, None]]:
+
         """Pop the element at ``index`` in the array JSON value under
         ``path`` at key ``name``.
 
         For more information: https://oss.redis.com/redisjson/commands/#jsonarrpop
         """  # noqa
         return self.execute_command("JSON.ARRPOP", name, str(path), index)
 
-    def arrtrim(self, name, path, start, stop):
+    def arrtrim(
+        self, name: str, path: str, start: int, stop: int
+    ) -> List[Union[int, None]]:
         """Trim the array JSON value under ``path`` at key ``name`` to the
         inclusive range given by ``start`` and ``stop``.
 
         For more information: https://oss.redis.com/redisjson/commands/#jsonarrtrim
         """  # noqa
         return self.execute_command("JSON.ARRTRIM", name, str(path), start, stop)
 
-    def type(self, name, path=Path.root_path()):
+    def type(self, name: str, path: Optional[str] = Path.root_path()) -> List[str]:
         """Get the type of the JSON value under ``path`` from key ``name``.
 
         For more information: https://oss.redis.com/redisjson/commands/#jsontype
         """  # noqa
         return self.execute_command("JSON.TYPE", name, str(path))
 
-    def resp(self, name, path=Path.root_path()):
+    def resp(self, name: str, path: Optional[str] = Path.root_path()) -> List:
         """Return the JSON value under ``path`` at key ``name``.
 
         For more information: https://oss.redis.com/redisjson/commands/#jsonresp
         """  # noqa
         return self.execute_command("JSON.RESP", name, str(path))
 
-    def objkeys(self, name, path=Path.root_path()):
+    def objkeys(
+        self, name: str, path: Optional[str] = Path.root_path()
+    ) -> List[Union[List[str], None]]:
         """Return the key names in the dictionary JSON value under ``path`` at
         key ``name``.
 
         For more information: https://oss.redis.com/redisjson/commands/#jsonobjkeys
         """  # noqa
         return self.execute_command("JSON.OBJKEYS", name, str(path))
 
-    def objlen(self, name, path=Path.root_path()):
+    def objlen(self, name: str, path: Optional[str] = Path.root_path()) -> int:
         """Return the length of the dictionary JSON value under ``path`` at key
         ``name``.
 
         For more information: https://oss.redis.com/redisjson/commands/#jsonobjlen
         """  # noqa
         return self.execute_command("JSON.OBJLEN", name, str(path))
 
-    def numincrby(self, name, path, number):
+    def numincrby(self, name: str, path: str, number: int) -> str:
         """Increment the numeric (integer or floating point) JSON value under
         ``path`` at key ``name`` by the provided ``number``.
 
@@ -113,7 +138,7 @@ def numincrby(self, name, path, number):
         )
 
     @deprecated(version="4.0.0", reason="deprecated since redisjson 1.0.0")
-    def nummultby(self, name, path, number):
+    def nummultby(self, name: str, path: str, number: int) -> str:
         """Multiply the numeric (integer or floating point) JSON value under
         ``path`` at key ``name`` with the provided ``number``.
 
@@ -123,7 +148,7 @@ def nummultby(self, name, path, number):
             "JSON.NUMMULTBY", name, str(path), self._encode(number)
         )
 
-    def clear(self, name, path=Path.root_path()):
+    def clear(self, name: str, path: Optional[str] = Path.root_path()) -> int:
         """
         Empty arrays and objects (to have zero slots/keys without deleting the
         array/object).
@@ -135,7 +160,7 @@ def clear(self, name, path=Path.root_path()):
         """  # noqa
         return self.execute_command("JSON.CLEAR", name, str(path))
 
-    def delete(self, key, path=Path.root_path()):
+    def delete(self, key: str, path: Optional[str] = Path.root_path()) -> int:
         """Delete the JSON value stored at key ``key`` under ``path``.
 
         For more information: https://oss.redis.com/redisjson/commands/#jsondel
@@ -145,7 +170,9 @@ def delete(self, key, path=Path.root_path()):
     # forget is an alias for delete
     forget = delete
 
-    def get(self, name, *args, no_escape=False):
+    def get(
+        self, name: str, *args, no_escape: Optional[bool] = False
+    ) -> List[JsonType]:
         """
         Get the object stored as a JSON value at key ``name``.
 
@@ -173,7 +200,7 @@ def get(self, name, *args, no_escape=False):
         except TypeError:
             return None
 
-    def mget(self, keys, path):
+    def mget(self, keys: List[str], path: str) -> List[JsonType]:
         """
         Get the objects stored as a JSON values under ``path``. ``keys``
         is a list of one or more keys.
@@ -185,7 +212,15 @@ def mget(self, keys, path):
         pieces.append(str(path))
         return self.execute_command("JSON.MGET", *pieces)
 
-    def set(self, name, path, obj, nx=False, xx=False, decode_keys=False):
+    def set(
+        self,
+        name: str,
+        path: str,
+        obj: JsonType,
+        nx: Optional[bool] = False,
+        xx: Optional[bool] = False,
+        decode_keys: Optional[bool] = False,
+    ) -> Optional[str]:
         """
         Set the JSON value at key ``name`` under the ``path`` to ``obj``.
 
@@ -216,7 +251,15 @@ def set(self, name, path, obj, nx=False, xx=False, decode_keys=False):
             pieces.append("XX")
         return self.execute_command("JSON.SET", *pieces)
 
-    def set_file(self, name, path, file_name, nx=False, xx=False, decode_keys=False):
+    def set_file(
+        self,
+        name: str,
+        path: str,
+        file_name: str,
+        nx: Optional[bool] = False,
+        xx: Optional[bool] = False,
+        decode_keys: Optional[bool] = False,
+    ) -> Optional[str]:
         """
         Set the JSON value at key ``name`` under the ``path`` to the content
         of the json file ``file_name``.
@@ -233,7 +276,14 @@ def set_file(self, name, path, file_name, nx=False, xx=False, decode_keys=False)
 
         return self.set(name, path, file_content, nx=nx, xx=xx, decode_keys=decode_keys)
 
-    def set_path(self, json_path, root_folder, nx=False, xx=False, decode_keys=False):
+    def set_path(
+        self,
+        json_path: str,
+        root_folder: str,
+        nx: Optional[bool] = False,
+        xx: Optional[bool] = False,
+        decode_keys: Optional[bool] = False,
+    ) -> List[Dict[str, bool]]:
         """
         Iterate over ``root_folder`` and set each JSON file to a value
         under ``json_path`` with the file name as the key.
@@ -264,7 +314,7 @@ def set_path(self, json_path, root_folder, nx=False, xx=False, decode_keys=False
 
         return set_files_result
 
-    def strlen(self, name, path=None):
+    def strlen(self, name: str, path: Optional[str] = None) -> List[Union[int, None]]:
         """Return the length of the string JSON value under ``path`` at key
         ``name``.
 
@@ -275,15 +325,19 @@ def strlen(self, name, path=None):
             pieces.append(str(path))
         return self.execute_command("JSON.STRLEN", *pieces)
 
-    def toggle(self, name, path=Path.root_path()):
+    def toggle(
+        self, name: str, path: Optional[str] = Path.root_path()
+    ) -> Union[bool, List[Optional[int]]]:
         """Toggle boolean value under ``path`` at key ``name``.
         returning the new value.
 
         For more information: https://oss.redis.com/redisjson/commands/#jsontoggle
         """  # noqa
         return self.execute_command("JSON.TOGGLE", name, str(path))
 
-    def strappend(self, name, value, path=Path.root_path()):
+    def strappend(
+        self, name: str, value: str, path: Optional[int] = Path.root_path()
+    ) -> Union[int, List[Optional[int]]]:
         """Append to the string JSON value. If two options are specified after
         the key name, the path is determined to be the first. If a single
         option is passed, then the root_path (i.e Path.root_path()) is used.
@@ -293,11 +347,16 @@ def strappend(self, name, value, path=Path.root_path()):
         pieces = [name, str(path), self._encode(value)]
         return self.execute_command("JSON.STRAPPEND", *pieces)
 
-    def debug(self, subcommand, key=None, path=Path.root_path()):
+    def debug(
+        self,
+        subcommand: str,
+        key: Optional[str] = None,
+        path: Optional[str] = Path.root_path(),
+    ) -> Union[int, List[str]]:
         """Return the memory usage in bytes of a value under ``path`` from
         key ``name``.
 
-        For more information: https://oss.redis.com/redisjson/commands/#jsondebg
+        For more information: https://oss.redis.com/redisjson/commands/#jsondebug
         """  # noqa
         valid_subcommands = ["MEMORY", "HELP"]
         if subcommand not in valid_subcommands: