coredis 5.5.0__cp313-cp313-macosx_11_0_arm64.whl

coredis/response/_callbacks/streams.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from typing import Any
+
+from coredis._utils import EncodingInsensitiveDict
+from coredis.response._callbacks import ResponseCallback
+from coredis.response._utils import flat_pairs_to_dict, flat_pairs_to_ordered_dict
+from coredis.response.types import (
+    StreamEntry,
+    StreamInfo,
+    StreamPending,
+    StreamPendingExt,
+)
+from coredis.typing import (
+    AnyStr,
+    OrderedDict,
+    ResponseType,
+    StringT,
+)
+
+
+class StreamRangeCallback(ResponseCallback[ResponseType, ResponseType, tuple[StreamEntry, ...]]):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> tuple[StreamEntry, ...]:
+        return tuple(StreamEntry(r[0], flat_pairs_to_ordered_dict(r[1])) for r in response)
+
+
+class ClaimCallback(
+    ResponseCallback[ResponseType, ResponseType, tuple[AnyStr, ...] | tuple[StreamEntry, ...]]
+):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> tuple[AnyStr, ...] | tuple[StreamEntry, ...]:
+        if self.options.get("justid") is not None:
+            return tuple(response)
+        else:
+            return StreamRangeCallback()(response)
+
+
+class AutoClaimCallback(
+    ResponseCallback[
+        ResponseType,
+        ResponseType,
+        tuple[AnyStr, tuple[AnyStr, ...]]
+        | tuple[AnyStr, tuple[StreamEntry, ...], tuple[AnyStr, ...]],
+    ]
+):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> (
+        tuple[AnyStr, tuple[AnyStr, ...]]
+        | tuple[AnyStr, tuple[StreamEntry, ...], tuple[AnyStr, ...]]
+    ):
+        if self.options.get("justid") is not None:
+            return response[0], tuple(response[1])
+        else:
+            return (
+                response[0],
+                StreamRangeCallback()(response[1]),
+                tuple(response[2]) if len(response) > 2 else (),
+            )
+
+
+class MultiStreamRangeCallback(
+    ResponseCallback[ResponseType, ResponseType, dict[AnyStr, tuple[StreamEntry, ...]] | None]
+):
+    def transform_3(
+        self,
+        response: ResponseType,
+    ) -> dict[AnyStr, tuple[StreamEntry, ...]] | None:
+        if response:
+            mapping: dict[AnyStr, tuple[StreamEntry, ...]] = {}
+
+            for stream_id, entries in response.items():
+                mapping[stream_id] = tuple(
+                    StreamEntry(r[0], flat_pairs_to_ordered_dict(r[1])) for r in entries
+                )
+
+            return mapping
+        return None
+
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> dict[AnyStr, tuple[StreamEntry, ...]] | None:
+        if response:
+            mapping: dict[AnyStr, tuple[StreamEntry, ...]] = {}
+
+            for stream_id, entries in response:
+                mapping[stream_id] = tuple(
+                    StreamEntry(r[0], flat_pairs_to_ordered_dict(r[1])) for r in entries
+                )
+
+            return mapping
+        return None
+
+
+class PendingCallback(
+    ResponseCallback[ResponseType, ResponseType, StreamPending | tuple[StreamPendingExt, ...]]
+):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> StreamPending | tuple[StreamPendingExt, ...]:
+        if not self.options.get("count"):
+            return StreamPending(
+                response[0],
+                response[1],
+                response[2],
+                OrderedDict((r[0], int(r[1])) for r in response[3] or []),
+            )
+        else:
+            return tuple(StreamPendingExt(sub[0], sub[1], sub[2], sub[3]) for sub in response)
+
+
+class XInfoCallback(ResponseCallback[ResponseType, ResponseType, tuple[dict[AnyStr, AnyStr], ...]]):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> tuple[dict[AnyStr, AnyStr], ...]:
+        return tuple(flat_pairs_to_dict(row) for row in response)
+
+
+class StreamInfoCallback(ResponseCallback[ResponseType, ResponseType, StreamInfo]):
+    def transform(
+        self,
+        response: ResponseType,
+    ) -> StreamInfo:
+        res: dict[StringT, Any] = EncodingInsensitiveDict(flat_pairs_to_dict(response))
+        if not self.options.get("full"):
+            k1 = "first-entry"
+            kn = "last-entry"
+            e1: StreamEntry | None = None
+            en: StreamEntry | None = None
+
+            if len(res.get(k1, [])) > 0:
+                v = res.get(k1)
+                e1 = StreamEntry(v[0], flat_pairs_to_ordered_dict(v[1]))
+                res.pop(k1)
+
+            if len(res.get(kn, [])) > 0:
+                v = res.get(kn)
+                en = StreamEntry(v[0], flat_pairs_to_ordered_dict(v[1]))
+                res.pop(kn)
+            res.update({"first-entry": e1, "last-entry": en})
+        else:
+            groups = res.get("groups")
+            if groups:
+                normalized_groups = []
+                for group in groups:
+                    g = EncodingInsensitiveDict(flat_pairs_to_dict(group))
+                    consumers = g["consumers"]
+                    normalized_consumers = []
+                    for consumer in consumers:
+                        normalized_consumers.append(flat_pairs_to_dict(consumer))
+                    g["consumers"] = normalized_consumers
+                    normalized_groups.append(g)
+                res["groups"] = normalized_groups
+            res.update(
+                {
+                    "entries": tuple(
+                        StreamEntry(k[0], flat_pairs_to_ordered_dict(k[1]))
+                        for k in res.get("entries", [])
+                    )
+                }
+            )
+        stream_info: StreamInfo = {
+            "first-entry": res.get("first-entry"),
+            "last-entry": res.get("last-entry"),
+            "length": res["length"],
+            "radix-tree-keys": res["radix-tree-keys"],
+            "radix-tree-nodes": res["radix-tree-nodes"],
+            "groups": res["groups"],
+            "last-generated-id": res["last-generated-id"],
+            "max-deleted-entry-id": str(res.get("max-deleted-entry-id")),
+            "entries-added": int(res.get("entries-added", 0)),
+            "recorded-first-entry-id": str(res.get("recorded-first-entry-id")),
+            "entries-read": int(res.get("entries-read", 0)),
+            "entries": res.get("entries"),
+        }
+        return stream_info

coredis/response/_callbacks/strings.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from typing import Any
+
+from coredis._utils import EncodingInsensitiveDict
+from coredis.response._callbacks import ResponseCallback, SimpleStringCallback
+from coredis.response.types import LCSMatch, LCSResult
+from coredis.typing import (
+    AnyStr,
+    ResponsePrimitive,
+    ResponseType,
+)
+
+
+class StringSetCallback(ResponseCallback[AnyStr | None, AnyStr | None, AnyStr | bool | None]):
+    def transform(self, response: AnyStr | None, **options: Any) -> AnyStr | bool | None:
+        if self.options.get("get"):
+            return response
+        else:
+            return SimpleStringCallback()(response)
+
+
+class LCSCallback(
+    ResponseCallback[
+        list[ResponseType],
+        dict[ResponsePrimitive, ResponseType],
+        LCSResult,
+    ]
+):
+    def transform(
+        self,
+        response: (list[ResponseType] | dict[ResponsePrimitive, ResponseType]),
+        **options: Any,
+    ) -> LCSResult:
+        assert (
+            isinstance(response, list)
+            and isinstance(response[-1], int)
+            and isinstance(response[1], list)
+        )
+
+        return LCSResult(
+            tuple(
+                LCSMatch(
+                    (int(k[0][0]), int(k[0][1])),
+                    (int(k[1][0]), int(k[1][1])),
+                    k[2] if len(k) > 2 else None,
+                )
+                for k in response[1]
+            ),
+            response[-1],
+        )
+
+    def transform_3(
+        self,
+        response: dict[ResponsePrimitive, ResponseType],
+        **options: Any,
+    ) -> LCSResult:
+        proxy = EncodingInsensitiveDict(response)
+
+        return LCSResult(
+            tuple(
+                LCSMatch(
+                    (int(k[0][0]), int(k[0][1])),
+                    (int(k[1][0]), int(k[1][1])),
+                    k[2] if len(k) > 2 else None,
+                )
+                for k in proxy["matches"]
+            ),
+            proxy["len"],
+        )

coredis/response/_callbacks/vector_sets.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from coredis._json import json
+from coredis._utils import nativestr
+from coredis.response._callbacks import ResponseCallback
+from coredis.response._utils import flat_pairs_to_dict
+from coredis.response.types import VectorData
+from coredis.typing import AnyStr, JsonType, ResponsePrimitive, StringT
+
+
+class VSimCallback(
+    ResponseCallback[
+        list[AnyStr],
+        list[AnyStr] | dict[AnyStr, float | list[float | JsonType]],
+        tuple[AnyStr, ...]
+        | dict[AnyStr, float]
+        | dict[AnyStr, JsonType]
+        | dict[AnyStr, tuple[float, JsonType]],
+    ],
+):
+    def transform(
+        self,
+        response: list[AnyStr],
+    ) -> (
+        tuple[AnyStr, ...]
+        | dict[AnyStr, float]
+        | dict[AnyStr, JsonType]
+        | dict[AnyStr, tuple[float, JsonType]]
+    ):
+        withscores, withattribs = self.options.get("withscores"), self.options.get("withattribs")
+        if withscores or withattribs:
+            it = iter(response)
+            match withscores, withattribs:
+                case True, None | False:
+                    return dict(list(zip(it, map(float, it))))
+                case None | False, True:
+                    return dict(list(zip(it, map(json.loads, it))))
+                case True, True:
+                    return dict(
+                        list(zip(it, map(lambda x: (float(x[0]), json.loads(x[1])), zip(it, it))))
+                    )
+        else:
+            return self.transform_3(response)
+
+    def transform_3(
+        self,
+        response: list[AnyStr]
+        | dict[AnyStr, float]
+        | dict[AnyStr, AnyStr]
+        | dict[AnyStr, list[float | AnyStr]],
+    ) -> (
+        tuple[AnyStr, ...]
+        | dict[AnyStr, float]
+        | dict[AnyStr, JsonType]
+        | dict[AnyStr, tuple[float, JsonType]]
+    ):
+        withscores, withattribs = self.options.get("withscores"), self.options.get("withattribs")
+        if withscores or withattribs:
+            assert isinstance(response, dict)
+            match withscores, withattribs:
+                case None | False, True:
+                    return {k: json.loads(v) for k, v in response.items()}
+                case True, True:
+                    return {k: (v[0], json.loads(v[1])) for k, v in response.items()}
+                case _:
+                    return response
+        else:
+            return tuple(response)
+
+
+class VLinksCallback(
+    ResponseCallback[
+        list[list[AnyStr]] | None,
+        list[list[AnyStr] | dict[AnyStr, float]] | None,
+        tuple[tuple[AnyStr, ...] | dict[AnyStr, float], ...] | None,
+    ],
+):
+    def transform(
+        self,
+        response: list[list[AnyStr]] | None,
+    ) -> tuple[tuple[AnyStr, ...] | dict[AnyStr, float], ...] | None:
+        if response:
+            if self.options.get("withscores"):
+                return tuple(dict(zip(it := iter(layer), map(float, it))) for layer in response)
+            else:
+                return tuple(tuple(layer) for layer in response)
+        return None
+
+    def transform_3(
+        self,
+        response: list[list[AnyStr] | dict[AnyStr, float]] | None,
+    ) -> tuple[tuple[AnyStr, ...] | dict[AnyStr, float], ...] | None:
+        if response:
+            if self.options.get("withscores"):
+                assert isinstance(response[0], dict)
+                return tuple(response)
+            else:
+                return tuple(tuple(layer) for layer in response)
+        return None
+
+
+class VEmbCallback(
+    ResponseCallback[
+        list[StringT] | list[ResponsePrimitive],
+        list[float] | list[ResponsePrimitive],
+        tuple[float, ...] | VectorData | None,
+    ]
+):
+    def transform(
+        self,
+        response: list[StringT] | list[ResponsePrimitive] | None,
+    ) -> tuple[float, ...] | VectorData | None:
+        if response:
+            if self.options.get("raw"):
+                return VectorData(
+                    quantization=nativestr(response[0]),
+                    blob=response[1],
+                    l2_norm=float(response[2]),
+                    quantization_range=float(response[3]) if len(response) == 4 else None,
+                )
+            else:
+                return tuple(map(float, response))
+        return None
+
+    def transform_3(
+        self,
+        response: list[float] | list[ResponsePrimitive] | None,
+    ) -> tuple[float, ...] | VectorData | None:
+        if response:
+            if self.options.get("raw"):
+                return VectorData(
+                    quantization=nativestr(response[0]),
+                    blob=response[1],
+                    l2_norm=response[2],
+                    quantization_range=response[3] if len(response) == 4 else None,
+                )
+            else:
+                return tuple(response)
+        return None
+
+
+class VInfoCallback(
+    ResponseCallback[
+        list[AnyStr | int] | None,
+        dict[AnyStr, AnyStr | int] | None,
+        dict[AnyStr, AnyStr | int] | None,
+    ]
+):
+    def transform(
+        self,
+        response: list[AnyStr | int] | None,
+    ) -> dict[AnyStr, AnyStr | int] | None:
+        return flat_pairs_to_dict(response) if response else None
+
+    def transform_3(
+        self,
+        response: dict[AnyStr, AnyStr | int] | None,
+    ) -> dict[AnyStr, AnyStr | int] | None:
+        return response

coredis/response/_utils.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from typing import cast
+
+from coredis.typing import (
+    Callable,
+    Iterable,
+    OrderedDict,
+    StringT,
+    TypeVar,
+)
+
+T_co = TypeVar("T_co")
+
+
+def flat_pairs_to_dict(
+    response: tuple[T_co, ...] | list[T_co],
+    value_transform: Callable[..., T_co] | None = None,
+) -> dict[T_co, T_co]:
+    """Creates a dict given a flat list of key/value pairs"""
+    if isinstance(response, dict):
+        return response
+    it = iter(response)
+    if value_transform:
+        return dict(zip(it, map(value_transform, it)))
+    else:
+        return dict(zip(it, it))
+
+
+def flat_pairs_to_ordered_dict(response: Iterable[T_co]) -> OrderedDict[StringT, T_co]:
+    """Creates a dict given a flat list of key/value pairs"""
+    it = iter(response)
+    return cast(OrderedDict[StringT, T_co], OrderedDict(zip(it, it)))