203 lines
6.3 KiB
Python
203 lines
6.3 KiB
Python
import copy
|
|
import random
|
|
import string
|
|
from typing import (
|
|
Any,
|
|
Callable,
|
|
Dict,
|
|
Iterable,
|
|
List,
|
|
Mapping,
|
|
Optional,
|
|
)
|
|
|
|
import redis
|
|
from redis.typing import ChannelT, PubSubHandler, Subscription
|
|
|
|
|
|
def list_or_args(keys: Any, args: Iterable[Any] | None) -> List[Any]:
|
|
# returns a single new list combining keys and args
|
|
try:
|
|
iter(keys)
|
|
# a string or bytes-like instance can be iterated, but indicates
|
|
# keys wasn't passed as a list
|
|
if isinstance(keys, (bytes, str, bytearray, memoryview)):
|
|
keys = [keys]
|
|
else:
|
|
keys = list(keys)
|
|
except TypeError:
|
|
keys = [keys]
|
|
if args:
|
|
keys.extend(args)
|
|
return keys
|
|
|
|
|
|
def parse_pubsub_subscriptions(
|
|
args: tuple[Any, ...], kwargs: Mapping[str, PubSubHandler]
|
|
) -> dict[ChannelT, PubSubHandler | None]:
|
|
parsed_args = list_or_args(args[0], args[1:]) if args else []
|
|
subscriptions: dict[ChannelT, PubSubHandler | None] = {}
|
|
for arg in parsed_args:
|
|
if isinstance(arg, Subscription):
|
|
subscriptions[arg.name] = arg.handler
|
|
else:
|
|
subscriptions[arg] = None
|
|
subscriptions.update(kwargs)
|
|
return subscriptions
|
|
|
|
|
|
def pubsub_subscription_args(
|
|
subscriptions: Mapping[ChannelT, PubSubHandler | None],
|
|
) -> list[ChannelT | Subscription]:
|
|
return [
|
|
channel if handler is None else Subscription(channel, handler)
|
|
for channel, handler in subscriptions.items()
|
|
]
|
|
|
|
|
|
def nativestr(x):
|
|
"""Return the decoded binary string, or a string, depending on type."""
|
|
r = x.decode("utf-8", "replace") if isinstance(x, bytes) else x
|
|
if r == "null":
|
|
return
|
|
return r
|
|
|
|
|
|
def delist(x):
|
|
"""Given a list of binaries, return the stringified version."""
|
|
if x is None:
|
|
return x
|
|
return [nativestr(obj) for obj in x]
|
|
|
|
|
|
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
|
|
if isinstance(item, float):
|
|
res.append(item)
|
|
continue
|
|
try:
|
|
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, OverflowError, TypeError):
|
|
try:
|
|
res.append(float(item))
|
|
except (ValueError, TypeError):
|
|
res.append(item_str)
|
|
|
|
return res
|
|
|
|
|
|
def random_string(length=10):
|
|
"""
|
|
Returns a random N character long string.
|
|
"""
|
|
return "".join( # nosec
|
|
random.choice(string.ascii_lowercase) for x in range(length)
|
|
)
|
|
|
|
|
|
def decode_dict_keys(obj):
|
|
"""Decode the keys of the given dictionary with utf-8."""
|
|
newobj = copy.copy(obj)
|
|
for k in obj.keys():
|
|
if isinstance(k, bytes):
|
|
newobj[k.decode("utf-8")] = newobj[k]
|
|
newobj.pop(k)
|
|
return newobj
|
|
|
|
|
|
def get_protocol_version(client):
|
|
if isinstance(client, redis.Redis) or isinstance(client, redis.asyncio.Redis):
|
|
return client.connection_pool.connection_kwargs.get("protocol")
|
|
elif isinstance(client, redis.cluster.AbstractRedisCluster):
|
|
return client.nodes_manager.connection_kwargs.get("protocol")
|
|
|
|
|
|
def get_legacy_responses(client):
|
|
"""Return the user-supplied ``legacy_responses`` flag for ``client``.
|
|
|
|
Defaults to ``True`` when the flag is not present in the client's
|
|
``connection_kwargs``. Mirrors :func:`get_protocol_version` so module
|
|
command bases can read both the protocol and the response-shape
|
|
selection from the same place.
|
|
"""
|
|
if isinstance(client, redis.Redis) or isinstance(client, redis.asyncio.Redis):
|
|
return client.connection_pool.connection_kwargs.get("legacy_responses", True)
|
|
elif isinstance(client, redis.cluster.AbstractRedisCluster):
|
|
return client.nodes_manager.connection_kwargs.get("legacy_responses", True)
|
|
return True
|
|
|
|
|
|
def apply_module_callbacks(
|
|
user_protocol: Optional[int],
|
|
legacy_responses: bool,
|
|
*,
|
|
common: Dict[str, Callable[..., Any]],
|
|
resp2: Dict[str, Callable[..., Any]],
|
|
resp3: Dict[str, Callable[..., Any]],
|
|
resp2_unified: Optional[Dict[str, Callable[..., Any]]] = None,
|
|
resp3_unified: Optional[Dict[str, Callable[..., Any]]] = None,
|
|
resp3_to_resp2_legacy: Optional[Dict[str, Callable[..., Any]]] = None,
|
|
) -> Dict[str, Callable[..., Any]]:
|
|
"""Return the merged module-callback dict for the given (protocol,
|
|
legacy_responses) combination.
|
|
|
|
Mirrors the selection used by
|
|
:func:`redis._parsers.response_callbacks.get_response_callbacks` for
|
|
the core callbacks: ``common`` is overlaid with the protocol-specific
|
|
dict matching ``user_protocol`` and ``legacy_responses``.
|
|
``resp2_unified`` defaults to ``resp2``, ``resp3_unified`` to ``resp3``,
|
|
and ``resp3_to_resp2_legacy`` to an empty dict.
|
|
"""
|
|
callbacks: Dict[str, Callable[..., Any]] = dict(common)
|
|
if legacy_responses:
|
|
if user_protocol is None:
|
|
callbacks.update(resp3_to_resp2_legacy or {})
|
|
elif user_protocol in (3, "3"):
|
|
callbacks.update(resp3)
|
|
else:
|
|
callbacks.update(resp2)
|
|
else:
|
|
if user_protocol is None or user_protocol in (3, "3"):
|
|
callbacks.update(resp3_unified if resp3_unified is not None else resp3)
|
|
else:
|
|
callbacks.update(resp2_unified if resp2_unified is not None else resp2)
|
|
return callbacks
|
|
|
|
|
|
def at_most_one_value_set(iterable: Iterable[Any]):
|
|
"""
|
|
Checks that at most one of the values in the iterable is truthy.
|
|
|
|
Args:
|
|
iterable: An iterable of values to check.
|
|
|
|
Returns:
|
|
True if at most one value is truthy, False otherwise.
|
|
|
|
Raises:
|
|
Might raise an error if the values in iterable are not boolean-compatible.
|
|
For example if the type of the values implement
|
|
__len__ or __bool__ methods and they raise an error.
|
|
"""
|
|
values = (bool(x) for x in iterable)
|
|
return sum(values) <= 1
|