coredis 5.5.0__cp313-cp313-macosx_11_0_arm64.whl

coredis/modules/response/_callbacks/autocomplete.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from coredis.modules.response.types import AutocompleteSuggestion
+from coredis.response._callbacks import ResponseCallback
+from coredis.typing import AnyStr, ResponseType
+
+
+class AutocompleteCallback(
+    ResponseCallback[
+        list[ResponseType],
+        list[ResponseType],
+        tuple[AutocompleteSuggestion[AnyStr], ...] | tuple[()],
+    ]
+):
+    def transform(
+        self,
+        response: list[ResponseType],
+    ) -> tuple[AutocompleteSuggestion[AnyStr], ...] | tuple[()]:
+        if not response:
+            return ()
+        step = 1
+        results = []
+        score_idx = payload_idx = 0
+        if self.options.get("withscores"):
+            score_idx = 1
+            step += 1
+        if self.options.get("withpayloads"):
+            payload_idx = score_idx + 1
+            step += 1
+
+        for k in range(0, len(response), step):
+            section = response[k : k + step]
+            score = section[score_idx] if score_idx else None
+            results.append(
+                AutocompleteSuggestion(
+                    section[0],
+                    float(score) if score else None,
+                    section[payload_idx] if payload_idx else None,
+                )
+            )
+
+        return tuple(results)

coredis/modules/response/_callbacks/graph.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+import asyncio
+import enum
+from typing import TYPE_CHECKING, Any
+
+from coredis._utils import b, nativestr
+from coredis.modules.response.types import (
+    GraphNode,
+    GraphPath,
+    GraphQueryResult,
+    GraphRelation,
+    GraphSlowLogInfo,
+)
+from coredis.response._callbacks import ResponseCallback
+from coredis.typing import (
+    AnyStr,
+    Generic,
+    Literal,
+    ResponsePrimitive,
+    ResponseType,
+    StringT,
+)
+
+if TYPE_CHECKING:
+    from coredis.client import Client
+
+
+class RedisValueTypes(enum.IntEnum):
+    VALUE_UNKNOWN = 0
+    VALUE_NULL = 1
+    VALUE_STRING = 2
+    VALUE_INTEGER = 3
+    VALUE_BOOLEAN = 4
+    VALUE_DOUBLE = 5
+    VALUE_ARRAY = 6
+    VALUE_EDGE = 7
+    VALUE_NODE = 8
+    VALUE_PATH = 9
+    VALUE_MAP = 10
+    VALUE_POINT = 11
+
+
+PROCEDURE_CALLS = {
+    "labels": "db.labels()",
+    "relationships": "db.relationshipTypes()",
+    "properties": "db.propertyKeys()",
+}
+
+SCALAR_MAPPING = {
+    RedisValueTypes.VALUE_INTEGER: int,
+    RedisValueTypes.VALUE_BOOLEAN: lambda v: b(v) == b"true",
+    RedisValueTypes.VALUE_DOUBLE: float,
+    RedisValueTypes.VALUE_STRING: lambda v: v,
+    RedisValueTypes.VALUE_NULL: lambda _: None,
+}
+
+
+class QueryCallback(
+    ResponseCallback[ResponseType, ResponseType, GraphQueryResult[AnyStr]],
+    Generic[AnyStr],
+):
+    properties: dict[int, StringT]
+    relationships: dict[int, StringT]
+    labels: dict[int, StringT]
+
+    def __init__(self, graph: StringT, **options: Any):
+        self.graph = graph
+        self.properties = {}
+        self.relationships = {}
+        self.labels = {}
+        super().__init__(**options)
+
+    async def pre_process(
+        self,
+        client: Client[Any],
+        response: ResponseType,
+    ) -> None:
+        if not len(response) == 3:
+            return
+        result_set = response[1]
+        max_label_id, max_relation_id, max_property_id = -1, -1, -1
+
+        cache = client.callback_storage[self.__class__]
+        self.labels = cache.setdefault(f"{self.graph}:labels", {})
+        self.relationships = cache.setdefault(f"{self.graph}:relationships", {})
+        self.properties = cache.setdefault(f"{self.graph}:properties", {})
+        for row in result_set:
+            for entity in row:
+                max_label_id, max_relation_id, max_property_id = self.fetch_max_ids(
+                    entity, max_label_id, max_relation_id, max_property_id
+                )
+        if any(k != -1 for k in [max_label_id, max_relation_id, max_property_id]):
+            self.labels, self.relationships, self.properties = await asyncio.gather(
+                self.fetch_mapping(max_label_id, "labels", client),
+                self.fetch_mapping(max_relation_id, "relationships", client),
+                self.fetch_mapping(max_property_id, "properties", client),
+            )
+
+    def fetch_max_ids(
+        self, entity: Any, max_label_id: int, max_relation_id: int, max_property_id: int
+    ) -> tuple[int, int, int]:
+        result_type = entity[0]
+        if result_type == RedisValueTypes.VALUE_NODE:
+            for label_id in entity[1][1]:
+                max_label_id = max(max_label_id, label_id)
+            for property_id in [k[0] for k in entity[1][2]]:
+                max_property_id = max(max_property_id, property_id)
+        elif result_type == RedisValueTypes.VALUE_EDGE:
+            max_relation_id = max(max_relation_id, entity[1][1])
+            for property_id in [k[0] for k in entity[1][4]]:
+                max_property_id = max(max_property_id, property_id)
+        elif result_type == RedisValueTypes.VALUE_PATH:
+            for segment in entity[1]:
+                max_label_id, max_relation_id, max_property_id = self.fetch_max_ids(
+                    segment, max_label_id, max_relation_id, max_property_id
+                )
+        elif result_type == RedisValueTypes.VALUE_ARRAY:
+            for segment in entity[1]:
+                max_label_id, max_relation_id, max_property_id = self.fetch_max_ids(
+                    segment, max_label_id, max_relation_id, max_property_id
+                )
+        return max_label_id, max_relation_id, max_property_id
+
+    async def fetch_mapping(
+        self,
+        max_id: int,
+        type: Literal["labels", "properties", "relationships"],
+        client: Client[Any],
+    ) -> dict[int, StringT]:
+        cache = client.callback_storage[self.__class__]
+        if max_id > max(cache[f"{self.graph}:{type}"] or [-1]):
+            cache[f"{self.graph}:{type}"] = dict(
+                enumerate(
+                    [
+                        k[0]
+                        for k in (
+                            await client.graph.ro_query(self.graph, f"CALL {PROCEDURE_CALLS[type]}")
+                        ).result_set
+                    ],
+                )
+            )
+        return cache[f"{self.graph}:{type}"]
+
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> GraphQueryResult[AnyStr]:
+        result_set = []
+        headers = []
+        if len(response) == 3:
+            headers = [k[1] for k in response[0]]
+            stats = response[2]
+            for result in response[1]:
+                entities = []
+                for entity in result:
+                    entities.append(self.parse_entity(entity))
+                result_set.append(entities)
+        else:
+            stats = response[0]
+        stats_mapping = dict(
+            map(
+                lambda v: int(v) if v.isalnum() else v,
+                map(lambda v: v.strip(), nativestr(m).split(":")),
+            )
+            for m in stats
+        )
+        return GraphQueryResult(tuple(headers), tuple(result_set), stats_mapping)
+
+    def parse_entity(self, entity):
+        result_type = entity[0]
+        if result_type in [
+            RedisValueTypes.VALUE_NULL,
+            RedisValueTypes.VALUE_STRING,
+            RedisValueTypes.VALUE_INTEGER,
+            RedisValueTypes.VALUE_BOOLEAN,
+            RedisValueTypes.VALUE_DOUBLE,
+        ]:
+            return SCALAR_MAPPING[result_type](entity[1])
+        elif result_type == RedisValueTypes.VALUE_MAP:
+            it = iter(entity[1])
+            return dict(zip(it, map(self.parse_entity, it)))
+        elif result_type == RedisValueTypes.VALUE_ARRAY:
+            return [self.parse_entity(k) for k in entity[1]]
+        elif result_type == RedisValueTypes.VALUE_POINT:
+            return tuple(map(float, entity[1]))
+        elif result_type == RedisValueTypes.VALUE_EDGE:
+            return GraphRelation(
+                id=entity[1][0],
+                type=self.relationships[entity[1][1]],
+                src_node=entity[1][2],
+                destination_node=entity[1][3],
+                properties={
+                    self.properties[k[0]]: self.parse_entity((k[1], k[2])) for k in entity[1][4]
+                },
+            )
+        elif result_type == RedisValueTypes.VALUE_NODE:
+            return GraphNode(
+                id=entity[1][0],
+                labels={self.labels[k] for k in entity[1][1]},
+                properties={
+                    self.properties[k[0]]: self.parse_entity((k[1], k[2])) for k in entity[1][2]
+                },
+            )
+        elif result_type == RedisValueTypes.VALUE_PATH:
+            nodes, relations = entity[1]
+            nodes = self.parse_entity(nodes)
+            relations = self.parse_entity(relations)
+            return GraphPath(nodes, relations)
+
+
+class GraphSlowLogCallback(
+    ResponseCallback[ResponseType, ResponseType, tuple[GraphSlowLogInfo, ...]]
+):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> tuple[GraphSlowLogInfo, ...]:
+        return tuple(GraphSlowLogInfo(int(k[0]), k[1], k[2], float(k[3])) for k in response)
+
+
+class ConfigGetCallback(
+    ResponseCallback[
+        ResponseType,
+        ResponseType,
+        ResponsePrimitive | dict[AnyStr, ResponsePrimitive],
+    ]
+):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> ResponsePrimitive | dict[AnyStr, ResponsePrimitive]:
+        if isinstance(response, list):
+            if isinstance(response[0], list):
+                return dict(response)
+            else:
+                return response[1]

coredis/modules/response/_callbacks/json.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from typing import cast
+
+from coredis._json import json
+from coredis.response._callbacks import ResponseCallback
+from coredis.typing import JsonType, ResponseType
+
+
+class JsonCallback(ResponseCallback[ResponseType, ResponseType, JsonType]):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> JsonType:
+        if isinstance(response, (bytes, str)):
+            deser = json.loads(response)
+        elif isinstance(response, list):
+            deser = [json.loads(e) if isinstance(e, (bytes, str)) else e for e in response]
+        else:
+            deser = response
+        return cast(JsonType, deser)  # alas we lie.

coredis/modules/response/_callbacks/search.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+from collections import ChainMap, OrderedDict
+from functools import partial
+
+from coredis._json import json
+from coredis._utils import EncodingInsensitiveDict
+from coredis.modules.response.types import (
+    SearchAggregationResult,
+    SearchDocument,
+    SearchResult,
+)
+from coredis.response._callbacks import ResponseCallback
+from coredis.response._utils import flat_pairs_to_dict
+from coredis.typing import (
+    AnyStr,
+    ResponsePrimitive,
+    ResponseType,
+)
+
+
+class SearchConfigCallback(
+    ResponseCallback[
+        list[list[ResponsePrimitive]],
+        dict[AnyStr, ResponseType] | list[list[ResponsePrimitive]],
+        dict[AnyStr, ResponsePrimitive],
+    ]
+):
+    def transform(
+        self,
+        response: list[list[ResponsePrimitive]],
+    ) -> dict[AnyStr, ResponsePrimitive]:
+        command_arguments = []
+        for item in response:
+            try:
+                v = (item[0], json.loads(item[1]))
+            except (ValueError, TypeError):
+                v = item
+            command_arguments.append(v)
+        return dict(command_arguments)
+
+    def transform_3(
+        self,
+        response: dict[AnyStr, ResponseType] | list[list[ResponsePrimitive]],
+    ) -> dict[AnyStr, ResponsePrimitive]:
+        if isinstance(response, list):
+            return self.transform(response)
+        else:
+            config = {}
+            for item, value in response.items():
+                try:
+                    config[item] = json.loads(value)
+                except (ValueError, TypeError):
+                    config[item] = value
+            return config
+
+
+class SearchResultCallback(
+    ResponseCallback[
+        list[ResponseType],
+        list[ResponseType] | dict[AnyStr, ResponseType],
+        SearchResult[AnyStr],
+    ]
+):
+    def transform(
+        self,
+        response: list[ResponseType],
+    ) -> SearchResult[AnyStr]:
+        if self.options.get("nocontent"):
+            return SearchResult[AnyStr](
+                response[0],
+                tuple(SearchDocument(i, None, None, None, None, {}) for i in response[1:]),
+            )
+        step = 2
+        results = []
+        score_idx = payload_idx = sort_key_idx = 0
+        if self.options.get("withscores"):
+            score_idx = 1
+            step += 1
+        if self.options.get("withpayloads"):
+            payload_idx = score_idx + 1
+            step += 1
+        if self.options.get("withsortkeys"):
+            sort_key_idx = payload_idx + 1
+            step += 1
+
+        for k in range(1, len(response) - 1, step):
+            section = response[k : k + step]
+            score_explain = None
+            if self.options.get("explainscore"):
+                score = section[score_idx][0]
+                score_explain = section[score_idx][1]
+            else:
+                score = section[score_idx] if score_idx else None
+            fields = EncodingInsensitiveDict(flat_pairs_to_dict(section[-1]))
+            if "$" in fields:
+                fields = json.loads(fields.pop("$"))
+            results.append(
+                SearchDocument(
+                    section[0],
+                    float(score) if score else None,
+                    score_explain,
+                    section[payload_idx] if payload_idx else None,
+                    section[sort_key_idx] if sort_key_idx else None,
+                    fields,
+                )
+            )
+        return SearchResult[AnyStr](response[0], tuple(results))
+
+    def transform_3(
+        self,
+        response: list[ResponseType] | dict[AnyStr, ResponseType],
+    ) -> SearchResult[AnyStr]:
+        results = []
+        if isinstance(response, list):
+            return self.transform(response)
+        else:
+            response = EncodingInsensitiveDict(response)
+            for result in response["results"]:
+                result = EncodingInsensitiveDict(result)
+                score_explain = None
+                if self.options.get("explainscore"):
+                    score, score_explain = result.get("score")
+                else:
+                    score = result.get("score", None)
+                fields = EncodingInsensitiveDict(result.get("extra_attributes", {}))
+                if "$" in fields:
+                    fields = json.loads(fields.pop("$"))
+                results.append(
+                    SearchDocument(
+                        result["id"],
+                        float(score) if score else None,
+                        score_explain,
+                        result["payload"] if self.options.get("withpayloads") else None,
+                        result["sortkey"] if self.options.get("withsortkeys") else None,
+                        fields,
+                    )
+                )
+            return SearchResult[AnyStr](response["total_results"], tuple(results))
+
+
+class AggregationResultCallback(
+    ResponseCallback[
+        list[ResponseType],
+        dict[AnyStr, ResponseType] | list[ResponseType],
+        SearchAggregationResult[AnyStr],
+    ]
+):
+    def transform(
+        self,
+        response: list[ResponseType],
+    ) -> SearchAggregationResult:
+        return SearchAggregationResult[AnyStr](
+            [
+                flat_pairs_to_dict(k, partial(self.try_json, self.options))
+                for k in (response[1:] if not self.options.get("with_cursor") else response[0][1:])
+            ],
+            response[1] if self.options.get("with_cursor") else None,
+        )
+
+    def transform_3(
+        self,
+        response: dict[AnyStr, ResponseType] | list[ResponseType],
+    ) -> SearchAggregationResult:
+        if (
+            self.options.get("with_cursor")
+            and isinstance(response[0], dict)
+            or isinstance(response, dict)
+        ):
+            response, cursor = response if self.options.get("with_cursor") else (response, None)
+            response = EncodingInsensitiveDict(response)
+            return SearchAggregationResult[AnyStr](
+                [
+                    {
+                        r: self.try_json(self.options, v)
+                        for r, v in EncodingInsensitiveDict(k)["extra_attributes"].items()
+                    }
+                    for k in response["results"]
+                ],
+                cursor,
+            )
+        else:
+            return self.transform(response)
+
+    @staticmethod
+    def try_json(options, value):
+        if not options.get("dialect", None) == 3:
+            return value
+        try:
+            return json.loads(value)
+        except ValueError:
+            return value
+
+
+class SpellCheckCallback(
+    ResponseCallback[
+        list[ResponseType],
+        dict[AnyStr, ResponseType] | list[ResponseType],
+        dict[AnyStr, OrderedDict[AnyStr, float]],
+    ]
+):
+    def transform(
+        self,
+        response: list[ResponseType],
+    ) -> dict[AnyStr, OrderedDict[AnyStr, float]]:
+        return {
+            result[1]: OrderedDict(
+                (suggestion[1], float(suggestion[0])) for suggestion in result[2]
+            )
+            for result in response
+        }
+
+    def transform_3(
+        self,
+        response: dict[AnyStr, ResponseType] | list[ResponseType],
+    ) -> dict[AnyStr, OrderedDict[AnyStr, float]]:
+        # For older versions of redis search that didn't support RESP3
+        if isinstance(response, list):
+            return self.transform(response)
+        response = EncodingInsensitiveDict(response)
+        return {key: OrderedDict(ChainMap(*result)) for key, result in response["results"].items()}

coredis/modules/response/_callbacks/timeseries.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+from typing import Any, cast
+
+from coredis._utils import EncodingInsensitiveDict
+from coredis.response._callbacks import (
+    ClusterMergeMapping,
+    DictCallback,
+    ResponseCallback,
+)
+from coredis.response._utils import flat_pairs_to_dict
+from coredis.typing import (
+    AnyStr,
+    RedisValueT,
+    ResponsePrimitive,
+    ResponseType,
+    Sequence,
+)
+
+
+class SampleCallback(
+    ResponseCallback[
+        list[RedisValueT],
+        list[RedisValueT],
+        tuple[int, float] | tuple[()],
+    ]
+):
+    def transform(
+        self,
+        response: list[RedisValueT],
+    ) -> tuple[int, float] | tuple[()]:
+        return (int(response[0]), float(response[1])) if response else ()
+
+
+class SamplesCallback(
+    ResponseCallback[
+        list[list[RedisValueT]] | None,
+        list[list[RedisValueT]] | None,
+        tuple[tuple[int, float], ...] | tuple[()],
+    ],
+):
+    def transform(
+        self,
+        response: list[list[RedisValueT]] | None,
+    ) -> tuple[tuple[int, float], ...] | tuple[()]:
+        if response:
+            return tuple(cast(tuple[int, float], SampleCallback().transform(r)) for r in response)
+        return ()
+
+
+class TimeSeriesInfoCallback(DictCallback[AnyStr, ResponseType]):
+    def transform(
+        self,
+        response: Sequence[ResponseType] | dict[ResponsePrimitive, ResponseType],
+    ) -> dict[AnyStr, ResponseType]:
+        dct = EncodingInsensitiveDict(super().transform(response))
+        if "labels" in dct:
+            dct["labels"] = dict(dct["labels"])
+        if "Chunks" in dct:
+            dct["Chunks"] = [flat_pairs_to_dict(chunk) for chunk in dct["Chunks"]]
+        if "rules" in dct and not isinstance(dct["rules"], dict):
+            dct["rules"] = {rule[0]: rule[1:] for rule in dct["rules"]}
+
+        return dict(dct)
+
+
+class TimeSeriesCallback(
+    ResponseCallback[
+        ResponseType,
+        ResponseType,
+        dict[AnyStr, tuple[dict[AnyStr, AnyStr], tuple[int, float] | tuple[()]]],
+    ]
+):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> dict[AnyStr, tuple[dict[AnyStr, AnyStr], tuple[int, float] | tuple[()]]]:
+        if isinstance(response, dict):
+            return {k: (v[0], tuple(v[1])) for k, v in response.items()}
+        else:
+            return {
+                r[0]: (dict(r[1]), (r[2][0], float(r[2][1])) if r[2] else tuple()) for r in response
+            }
+
+
+class TimeSeriesMultiCallback(
+    ResponseCallback[
+        ResponseType,
+        ResponseType,
+        dict[
+            AnyStr,
+            tuple[dict[AnyStr, AnyStr], tuple[tuple[int, float], ...] | tuple[()]],
+        ],
+    ]
+):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> dict[
+        AnyStr,
+        tuple[dict[AnyStr, AnyStr], tuple[tuple[int, float], ...] | tuple[()]],
+    ]:
+        if self.options.get("grouped"):
+            return {
+                r[0]: (
+                    flat_pairs_to_dict(r[1][0]) if r[1] else {},
+                    tuple(SampleCallback().transform(t) for t in r[2]),
+                )
+                for r in cast(Any, response)
+            }
+        else:
+            return {
+                r[0]: (
+                    dict(r[1]),
+                    tuple(SampleCallback().transform(t) for t in r[2]),
+                )
+                for r in cast(Any, response)
+            }
+
+    def transform_3(
+        self,
+        response: ResponseType,
+    ) -> dict[
+        AnyStr,
+        tuple[dict[AnyStr, AnyStr], tuple[tuple[int, float], ...] | tuple[()]],
+    ]:
+        if isinstance(response, dict):
+            if self.options.get("grouped"):
+                return {
+                    k: (
+                        r[0],
+                        tuple(SampleCallback().transform(t) for t in r[-1]),
+                    )
+                    for k, r in response.items()
+                }
+            else:
+                return {
+                    k: (
+                        r[0],
+                        tuple(SampleCallback().transform(t) for t in r[-1]),
+                    )
+                    for k, r in response.items()
+                }
+        else:
+            return self.transform(response)
+
+
+class ClusterMergeTimeSeries(ClusterMergeMapping[AnyStr, tuple[Any, ...]]):
+    def __init__(self) -> None:
+        super().__init__(value_combine=self.merge)
+
+    def merge(self, values: Any) -> tuple[dict[AnyStr, AnyStr], tuple[tuple[int, float], ...]]:
+        merged_labels: dict[AnyStr, AnyStr] = {}
+        merged_series: tuple[tuple[int, float], ...] = ()
+        for value in values:
+            merged_labels.update(value[0])
+            merged_series = merged_series + value[1]
+        return merged_labels, tuple(merged_series)