redis 6.0.0b1__py3-none-any.whl → 6.1.0__py3-none-any.whl

redis/commands/helpers.py

@@ -43,19 +43,32 @@ def parse_to_list(response):
     """Optimistically parse the response to a list."""
     res = []
 
+    special_values = {"infinity", "nan", "-infinity"}
+
     if response is None:
         return res
 
     for item in response:
+        if item is None:
+            res.append(None)
+            continue
         try:
-            res.append(int(item))
-        except ValueError:
-            try:
-                res.append(float(item))
-            except ValueError:
-                res.append(nativestr(item))
+            item_str = nativestr(item)
         except TypeError:
             res.append(None)
+            continue
+
+        if isinstance(item_str, str) and item_str.lower() in special_values:
+            res.append(item_str)  # Keep as string
+        else:
+            try:
+                res.append(int(item))
+            except ValueError:
+                try:
+                    res.append(float(item))
+                except ValueError:
+                    res.append(item_str)
+
     return res
 
 

redis/commands/json/__init__.py

@@ -120,7 +120,7 @@ class JSON(JSONCommands):
                 startup_nodes=self.client.nodes_manager.startup_nodes,
                 result_callbacks=self.client.result_callbacks,
                 cluster_response_callbacks=self.client.cluster_response_callbacks,
-                cluster_error_retry_attempts=self.client.cluster_error_retry_attempts,
+                cluster_error_retry_attempts=self.client.retry.get_retries(),
                 read_from_replicas=self.client.read_from_replicas,
                 reinitialize_steps=self.client.reinitialize_steps,
                 lock=self.client._lock,

redis/commands/json/commands.py

@@ -15,7 +15,7 @@ class JSONCommands:
 
     def arrappend(
         self, name: str, path: Optional[str] = Path.root_path(), *args: List[JsonType]
-    ) -> List[Union[int, None]]:
+    ) -> List[Optional[int]]:
         """Append the objects ``args`` to the array under the
         ``path` in key ``name``.
 
@@ -33,7 +33,7 @@ class JSONCommands:
         scalar: int,
         start: Optional[int] = None,
         stop: Optional[int] = None,
-    ) -> List[Union[int, None]]:
+    ) -> List[Optional[int]]:
         """
         Return the index of ``scalar`` in the JSON array under ``path`` at key
         ``name``.
@@ -53,7 +53,7 @@ class JSONCommands:
 
     def arrinsert(
         self, name: str, path: str, index: int, *args: List[JsonType]
-    ) -> List[Union[int, None]]:
+    ) -> List[Optional[int]]:
         """Insert the objects ``args`` to the array at index ``index``
         under the ``path` in key ``name``.
 
@@ -66,7 +66,7 @@ class JSONCommands:
 
     def arrlen(
         self, name: str, path: Optional[str] = Path.root_path()
-    ) -> List[Union[int, None]]:
+    ) -> List[Optional[int]]:
         """Return the length of the array JSON value under ``path``
         at key``name``.
 
@@ -79,7 +79,7 @@ class JSONCommands:
         name: str,
         path: Optional[str] = Path.root_path(),
         index: Optional[int] = -1,
-    ) -> List[Union[str, None]]:
+    ) -> List[Optional[str]]:
         """Pop the element at ``index`` in the array JSON value under
         ``path`` at key ``name``.
 
@@ -89,7 +89,7 @@ class JSONCommands:
 
     def arrtrim(
         self, name: str, path: str, start: int, stop: int
-    ) -> List[Union[int, None]]:
+    ) -> List[Optional[int]]:
         """Trim the array JSON value under ``path`` at key ``name`` to the
         inclusive range given by ``start`` and ``stop``.
 
@@ -113,7 +113,7 @@ class JSONCommands:
 
     def objkeys(
         self, name: str, path: Optional[str] = Path.root_path()
-    ) -> List[Union[List[str], None]]:
+    ) -> List[Optional[List[str]]]:
         """Return the key names in the dictionary JSON value under ``path`` at
         key ``name``.
 
@@ -357,7 +357,7 @@ class JSONCommands:
 
         return set_files_result
 
-    def strlen(self, name: str, path: Optional[str] = None) -> List[Union[int, None]]:
+    def strlen(self, name: str, path: Optional[str] = None) -> List[Optional[int]]:
         """Return the length of the string JSON value under ``path`` at key
         ``name``.
 

redis/commands/redismodules.py

@@ -1,4 +1,14 @@
+from __future__ import annotations
+
 from json import JSONDecoder, JSONEncoder
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .bf import BFBloom, CFBloom, CMSBloom, TDigestBloom, TOPKBloom
+    from .json import JSON
+    from .search import AsyncSearch, Search
+    from .timeseries import TimeSeries
+    from .vectorset import VectorSet
 
 
 class RedisModuleCommands:
@@ -6,7 +16,7 @@ class RedisModuleCommands:
     modules into the command namespace.
     """
 
-    def json(self, encoder=JSONEncoder(), decoder=JSONDecoder()):
+    def json(self, encoder=JSONEncoder(), decoder=JSONDecoder()) -> JSON:
         """Access the json namespace, providing support for redis json."""
 
         from .json import JSON
@@ -14,7 +24,7 @@ class RedisModuleCommands:
         jj = JSON(client=self, encoder=encoder, decoder=decoder)
         return jj
 
-    def ft(self, index_name="idx"):
+    def ft(self, index_name="idx") -> Search:
         """Access the search namespace, providing support for redis search."""
 
         from .search import Search
@@ -22,7 +32,7 @@ class RedisModuleCommands:
         s = Search(client=self, index_name=index_name)
         return s
 
-    def ts(self):
+    def ts(self) -> TimeSeries:
         """Access the timeseries namespace, providing support for
         redis timeseries data.
         """
@@ -32,7 +42,7 @@ class RedisModuleCommands:
         s = TimeSeries(client=self)
         return s
 
-    def bf(self):
+    def bf(self) -> BFBloom:
         """Access the bloom namespace."""
 
         from .bf import BFBloom
@@ -40,7 +50,7 @@ class RedisModuleCommands:
         bf = BFBloom(client=self)
         return bf
 
-    def cf(self):
+    def cf(self) -> CFBloom:
         """Access the bloom namespace."""
 
         from .bf import CFBloom
@@ -48,7 +58,7 @@ class RedisModuleCommands:
         cf = CFBloom(client=self)
         return cf
 
-    def cms(self):
+    def cms(self) -> CMSBloom:
         """Access the bloom namespace."""
 
         from .bf import CMSBloom
@@ -56,7 +66,7 @@ class RedisModuleCommands:
         cms = CMSBloom(client=self)
         return cms
 
-    def topk(self):
+    def topk(self) -> TOPKBloom:
         """Access the bloom namespace."""
 
         from .bf import TOPKBloom
@@ -64,7 +74,7 @@ class RedisModuleCommands:
         topk = TOPKBloom(client=self)
         return topk
 
-    def tdigest(self):
+    def tdigest(self) -> TDigestBloom:
         """Access the bloom namespace."""
 
         from .bf import TDigestBloom
@@ -72,9 +82,17 @@ class RedisModuleCommands:
         tdigest = TDigestBloom(client=self)
         return tdigest
 
+    def vset(self) -> VectorSet:
+        """Access the VectorSet commands namespace."""
+
+        from .vectorset import VectorSet
+
+        vset = VectorSet(client=self)
+        return vset
+
 
 class AsyncRedisModuleCommands(RedisModuleCommands):
-    def ft(self, index_name="idx"):
+    def ft(self, index_name="idx") -> AsyncSearch:
         """Access the search namespace, providing support for redis search."""
 
         from .search import AsyncSearch

redis/commands/search/commands.py

@@ -23,7 +23,6 @@ ALTER_CMD = "FT.ALTER"
 SEARCH_CMD = "FT.SEARCH"
 ADD_CMD = "FT.ADD"
 ADDHASH_CMD = "FT.ADDHASH"
-DROP_CMD = "FT.DROP"
 DROPINDEX_CMD = "FT.DROPINDEX"
 EXPLAIN_CMD = "FT.EXPLAIN"
 EXPLAINCLI_CMD = "FT.EXPLAINCLI"
@@ -35,7 +34,6 @@ SPELLCHECK_CMD = "FT.SPELLCHECK"
 DICT_ADD_CMD = "FT.DICTADD"
 DICT_DEL_CMD = "FT.DICTDEL"
 DICT_DUMP_CMD = "FT.DICTDUMP"
-GET_CMD = "FT.GET"
 MGET_CMD = "FT.MGET"
 CONFIG_CMD = "FT.CONFIG"
 TAGVALS_CMD = "FT.TAGVALS"
@@ -406,6 +404,7 @@ class SearchCommands:
             doc_id, conn=None, score=score, language=language, replace=replace
         )
 
+    @deprecated_function(version="2.0.0", reason="deprecated since redisearch 2.0")
     def delete_document(self, doc_id, conn=None, delete_actual_document=False):
         """
         Delete a document from index
@@ -440,6 +439,7 @@ class SearchCommands:
 
         return Document(id=id, **fields)
 
+    @deprecated_function(version="2.0.0", reason="deprecated since redisearch 2.0")
     def get(self, *ids):
         """
         Returns the full contents of multiple documents.

redis/commands/timeseries/__init__.py

@@ -84,7 +84,7 @@ class TimeSeries(TimeSeriesCommands):
                 startup_nodes=self.client.nodes_manager.startup_nodes,
                 result_callbacks=self.client.result_callbacks,
                 cluster_response_callbacks=self.client.cluster_response_callbacks,
-                cluster_error_retry_attempts=self.client.cluster_error_retry_attempts,
+                cluster_error_retry_attempts=self.client.retry.get_retries(),
                 read_from_replicas=self.client.read_from_replicas,
                 reinitialize_steps=self.client.reinitialize_steps,
                 lock=self.client._lock,

redis/commands/vectorset/__init__.py

@@ -0,0 +1,46 @@
+import json
+
+from redis._parsers.helpers import pairs_to_dict
+from redis.commands.vectorset.utils import (
+    parse_vemb_result,
+    parse_vlinks_result,
+    parse_vsim_result,
+)
+
+from ..helpers import get_protocol_version
+from .commands import (
+    VEMB_CMD,
+    VGETATTR_CMD,
+    VINFO_CMD,
+    VLINKS_CMD,
+    VSIM_CMD,
+    VectorSetCommands,
+)
+
+
+class VectorSet(VectorSetCommands):
+    def __init__(self, client, **kwargs):
+        """Create a new VectorSet client."""
+        # Set the module commands' callbacks
+        self._MODULE_CALLBACKS = {
+            VEMB_CMD: parse_vemb_result,
+            VGETATTR_CMD: lambda r: r and json.loads(r) or None,
+        }
+
+        self._RESP2_MODULE_CALLBACKS = {
+            VINFO_CMD: lambda r: r and pairs_to_dict(r) or None,
+            VSIM_CMD: parse_vsim_result,
+            VLINKS_CMD: parse_vlinks_result,
+        }
+        self._RESP3_MODULE_CALLBACKS = {}
+
+        self.client = client
+        self.execute_command = client.execute_command
+
+        if get_protocol_version(self.client) in ["3", 3]:
+            self._MODULE_CALLBACKS.update(self._RESP3_MODULE_CALLBACKS)
+        else:
+            self._MODULE_CALLBACKS.update(self._RESP2_MODULE_CALLBACKS)
+
+        for k, v in self._MODULE_CALLBACKS.items():
+            self.client.set_response_callback(k, v)

redis/commands/vectorset/commands.py

@@ -0,0 +1,367 @@
+import json
+from enum import Enum
+from typing import Awaitable, Dict, List, Optional, Union
+
+from redis.client import NEVER_DECODE
+from redis.commands.helpers import get_protocol_version
+from redis.exceptions import DataError
+from redis.typing import CommandsProtocol, EncodableT, KeyT, Number
+
+VADD_CMD = "VADD"
+VSIM_CMD = "VSIM"
+VREM_CMD = "VREM"
+VDIM_CMD = "VDIM"
+VCARD_CMD = "VCARD"
+VEMB_CMD = "VEMB"
+VLINKS_CMD = "VLINKS"
+VINFO_CMD = "VINFO"
+VSETATTR_CMD = "VSETATTR"
+VGETATTR_CMD = "VGETATTR"
+VRANDMEMBER_CMD = "VRANDMEMBER"
+
+
+class QuantizationOptions(Enum):
+    """Quantization options for the VADD command."""
+
+    NOQUANT = "NOQUANT"
+    BIN = "BIN"
+    Q8 = "Q8"
+
+
+class CallbacksOptions(Enum):
+    """Options that can be set for the commands callbacks"""
+
+    RAW = "RAW"
+    WITHSCORES = "WITHSCORES"
+    ALLOW_DECODING = "ALLOW_DECODING"
+    RESP3 = "RESP3"
+
+
+class VectorSetCommands(CommandsProtocol):
+    """Redis VectorSet commands"""
+
+    def vadd(
+        self,
+        key: KeyT,
+        vector: Union[List[float], bytes],
+        element: str,
+        reduce_dim: Optional[int] = None,
+        cas: Optional[bool] = False,
+        quantization: Optional[QuantizationOptions] = None,
+        ef: Optional[Number] = None,
+        attributes: Optional[Union[dict, str]] = None,
+        numlinks: Optional[int] = None,
+    ) -> Union[Awaitable[int], int]:
+        """
+        Add vector ``vector`` for element ``element`` to a vector set ``key``.
+
+        ``reduce_dim`` sets the dimensions to reduce the vector to.
+                If not provided, the vector is not reduced.
+
+        ``cas`` is a boolean flag that indicates whether to use CAS (check-and-set style)
+                when adding the vector. If not provided, CAS is not used.
+
+        ``quantization`` sets the quantization type to use.
+                If not provided, int8 quantization is used.
+                The options are:
+                - NOQUANT: No quantization
+                - BIN: Binary quantization
+                - Q8: Signed 8-bit quantization
+
+        ``ef`` sets the exploration factor to use.
+                If not provided, the default exploration factor is used.
+
+        ``attributes`` is a dictionary or json string that contains the attributes to set for the vector.
+                If not provided, no attributes are set.
+
+        ``numlinks`` sets the number of links to create for the vector.
+                If not provided, the default number of links is used.
+
+        For more information see https://redis.io/commands/vadd
+        """
+        if not vector or not element:
+            raise DataError("Both vector and element must be provided")
+
+        pieces = []
+        if reduce_dim:
+            pieces.extend(["REDUCE", reduce_dim])
+
+        values_pieces = []
+        if isinstance(vector, bytes):
+            values_pieces.extend(["FP32", vector])
+        else:
+            values_pieces.extend(["VALUES", len(vector)])
+            values_pieces.extend(vector)
+        pieces.extend(values_pieces)
+
+        pieces.append(element)
+
+        if cas:
+            pieces.append("CAS")
+
+        if quantization:
+            pieces.append(quantization.value)
+
+        if ef:
+            pieces.extend(["EF", ef])
+
+        if attributes:
+            if isinstance(attributes, dict):
+                # transform attributes to json string
+                attributes_json = json.dumps(attributes)
+            else:
+                attributes_json = attributes
+            pieces.extend(["SETATTR", attributes_json])
+
+        if numlinks:
+            pieces.extend(["M", numlinks])
+
+        return self.execute_command(VADD_CMD, key, *pieces)
+
+    def vsim(
+        self,
+        key: KeyT,
+        input: Union[List[float], bytes, str],
+        with_scores: Optional[bool] = False,
+        count: Optional[int] = None,
+        ef: Optional[Number] = None,
+        filter: Optional[str] = None,
+        filter_ef: Optional[str] = None,
+        truth: Optional[bool] = False,
+        no_thread: Optional[bool] = False,
+    ) -> Union[
+        Awaitable[Optional[List[Union[List[EncodableT], Dict[EncodableT, Number]]]]],
+        Optional[List[Union[List[EncodableT], Dict[EncodableT, Number]]]],
+    ]:
+        """
+        Compare a vector or element ``input``  with the other vectors in a vector set ``key``.
+
+        ``with_scores`` sets if the results should be returned with the
+                similarity scores of the elements in the result.
+
+        ``count`` sets the number of results to return.
+
+        ``ef`` sets the exploration factor.
+
+        ``filter`` sets filter that should be applied for the search.
+
+        ``filter_ef`` sets the max filtering effort.
+
+        ``truth`` when enabled forces the command to perform linear scan.
+
+        ``no_thread`` when enabled forces the command to execute the search
+                on the data structure in the main thread.
+
+        For more information see https://redis.io/commands/vsim
+        """
+
+        if not input:
+            raise DataError("'input' should be provided")
+
+        pieces = []
+        options = {}
+
+        if isinstance(input, bytes):
+            pieces.extend(["FP32", input])
+        elif isinstance(input, list):
+            pieces.extend(["VALUES", len(input)])
+            pieces.extend(input)
+        else:
+            pieces.extend(["ELE", input])
+
+        if with_scores:
+            pieces.append("WITHSCORES")
+            options[CallbacksOptions.WITHSCORES.value] = True
+
+        if count:
+            pieces.extend(["COUNT", count])
+
+        if ef:
+            pieces.extend(["EF", ef])
+
+        if filter:
+            pieces.extend(["FILTER", filter])
+
+        if filter_ef:
+            pieces.extend(["FILTER-EF", filter_ef])
+
+        if truth:
+            pieces.append("TRUTH")
+
+        if no_thread:
+            pieces.append("NOTHREAD")
+
+        return self.execute_command(VSIM_CMD, key, *pieces, **options)
+
+    def vdim(self, key: KeyT) -> Union[Awaitable[int], int]:
+        """
+        Get the dimension of a vector set.
+
+        In the case of vectors that were populated using the `REDUCE`
+        option, for random projection, the vector set will report the size of
+        the projected (reduced) dimension.
+
+        Raises `redis.exceptions.ResponseError` if the vector set doesn't exist.
+
+        For more information see https://redis.io/commands/vdim
+        """
+        return self.execute_command(VDIM_CMD, key)
+
+    def vcard(self, key: KeyT) -> Union[Awaitable[int], int]:
+        """
+        Get the cardinality(the number of elements) of a vector set with key ``key``.
+
+        Raises `redis.exceptions.ResponseError` if the vector set doesn't exist.
+
+        For more information see https://redis.io/commands/vcard
+        """
+        return self.execute_command(VCARD_CMD, key)
+
+    def vrem(self, key: KeyT, element: str) -> Union[Awaitable[int], int]:
+        """
+        Remove an element from a vector set.
+
+        For more information see https://redis.io/commands/vrem
+        """
+        return self.execute_command(VREM_CMD, key, element)
+
+    def vemb(
+        self, key: KeyT, element: str, raw: Optional[bool] = False
+    ) -> Union[
+        Awaitable[Optional[Union[List[EncodableT], Dict[str, EncodableT]]]],
+        Optional[Union[List[EncodableT], Dict[str, EncodableT]]],
+    ]:
+        """
+        Get the approximated vector of an element ``element`` from vector set ``key``.
+
+        ``raw`` is a boolean flag that indicates whether to return the
+                interal representation used by the vector.
+
+
+        For more information see https://redis.io/commands/vembed
+        """
+        options = {}
+        pieces = []
+        pieces.extend([key, element])
+
+        if get_protocol_version(self.client) in ["3", 3]:
+            options[CallbacksOptions.RESP3.value] = True
+
+        if raw:
+            pieces.append("RAW")
+
+            options[NEVER_DECODE] = True
+            if (
+                hasattr(self.client, "connection_pool")
+                and self.client.connection_pool.connection_kwargs["decode_responses"]
+            ) or (
+                hasattr(self.client, "nodes_manager")
+                and self.client.nodes_manager.connection_kwargs["decode_responses"]
+            ):
+                # allow decoding in the postprocessing callback
+                # if the user set decode_responses=True
+                # in the connection pool
+                options[CallbacksOptions.ALLOW_DECODING.value] = True
+
+            options[CallbacksOptions.RAW.value] = True
+
+        return self.execute_command(VEMB_CMD, *pieces, **options)
+
+    def vlinks(
+        self, key: KeyT, element: str, with_scores: Optional[bool] = False
+    ) -> Union[
+        Awaitable[
+            Optional[
+                List[Union[List[Union[str, bytes]], Dict[Union[str, bytes], Number]]]
+            ]
+        ],
+        Optional[List[Union[List[Union[str, bytes]], Dict[Union[str, bytes], Number]]]],
+    ]:
+        """
+        Returns the neighbors for each level the element ``element`` exists in the vector set ``key``.
+
+        The result is a list of lists, where each list contains the neighbors for one level.
+        If the element does not exist, or if the vector set does not exist, None is returned.
+
+        If the ``WITHSCORES`` option is provided, the result is a list of dicts,
+        where each dict contains the neighbors for one level, with the scores as values.
+
+        For more information see https://redis.io/commands/vlinks
+        """
+        options = {}
+        pieces = []
+        pieces.extend([key, element])
+
+        if with_scores:
+            pieces.append("WITHSCORES")
+            options[CallbacksOptions.WITHSCORES.value] = True
+
+        return self.execute_command(VLINKS_CMD, *pieces, **options)
+
+    def vinfo(self, key: KeyT) -> Union[Awaitable[dict], dict]:
+        """
+        Get information about a vector set.
+
+        For more information see https://redis.io/commands/vinfo
+        """
+        return self.execute_command(VINFO_CMD, key)
+
+    def vsetattr(
+        self, key: KeyT, element: str, attributes: Optional[Union[dict, str]] = None
+    ) -> Union[Awaitable[int], int]:
+        """
+        Associate or remove JSON attributes ``attributes`` of element ``element``
+        for vector set ``key``.
+
+        For more information see https://redis.io/commands/vsetattr
+        """
+        if attributes is None:
+            attributes_json = "{}"
+        elif isinstance(attributes, dict):
+            # transform attributes to json string
+            attributes_json = json.dumps(attributes)
+        else:
+            attributes_json = attributes
+
+        return self.execute_command(VSETATTR_CMD, key, element, attributes_json)
+
+    def vgetattr(
+        self, key: KeyT, element: str
+    ) -> Union[Optional[Awaitable[dict]], Optional[dict]]:
+        """
+        Retrieve the JSON attributes of an element ``elemet`` for vector set ``key``.
+
+        If the element does not exist, or if the vector set does not exist, None is
+        returned.
+
+        For more information see https://redis.io/commands/vgetattr
+        """
+        return self.execute_command(VGETATTR_CMD, key, element)
+
+    def vrandmember(
+        self, key: KeyT, count: Optional[int] = None
+    ) -> Union[
+        Awaitable[Optional[Union[List[str], str]]], Optional[Union[List[str], str]]
+    ]:
+        """
+        Returns random elements from a vector set ``key``.
+
+        ``count`` is the number of elements to return.
+                If ``count`` is not provided, a single element is returned as a single string.
+                If ``count`` is positive(smaller than the number of elements
+                            in the vector set), the command returns a list with up to ``count``
+                            distinct elements from the vector set
+                If ``count`` is negative, the command returns a list with ``count`` random elements,
+                            potentially with duplicates.
+                If ``count`` is greater than the number of elements in the vector set,
+                            only the entire set is returned as a list.
+
+        If the vector set does not exist, ``None`` is returned.
+
+        For more information see https://redis.io/commands/vrandmember
+        """
+        pieces = []
+        pieces.append(key)
+        if count is not None:
+            pieces.append(count)
+        return self.execute_command(VRANDMEMBER_CMD, *pieces)