from __future__ import annotations
import functools
import hashlib
import inspect
import itertools
from typing import TYPE_CHECKING, Any, cast, get_args, overload
from deprecated.sphinx import versionadded
from coredis._utils import b
from coredis.commands import CommandRequest, CommandResponseT
from coredis.exceptions import NoScriptError
from coredis.response._callbacks import NoopCallback
from coredis.retry import ConstantRetryPolicy
from coredis.typing import (
AnyStr,
Awaitable,
Callable,
Generic,
KeyT,
P,
Parameters,
R,
ResponseType,
StringT,
ValueT,
add_runtime_checks,
safe_beartype,
)
if TYPE_CHECKING:
import coredis.client
[docs]
class Script(Generic[AnyStr]):
"""
An executable Lua script object returned by :meth:`coredis.Redis.register_script`.
Instances of the class are callable and take arguments in the same shape
as :meth:`coredis.Redis.evalsha` or :meth:`coredis.Redis.eval`
(i.e a list of :paramref:`~__call__.keys` and :paramref:`~__call__.args`).
Example::
client = coredis.Redis()
async with client:
await client.set("test", "co")
concat = client.register_script("return redis.call('GET', KEYS[1]) + ARGV[1]")
assert await concat(['test'], ['redis']) == "coredis"
"""
#: SHA of this script once it's registered with the redis server
sha: AnyStr
def __init__(
self,
registered_client: coredis.client.Client[AnyStr] | None = None,
script: StringT | None = None,
readonly: bool = False,
):
"""
:param client: The client to use for executing the lua script. If
this is ``None`` the client will have to be provided when invoking
the script using :meth:`__call__` with the :paramref:`__call__.client`
parameter.
:param script: The lua script that will be used by :meth:`__call__`
:param readonly: If ``True`` the script will be called with
:meth:`coredis.Redis.evalsha_ro` instead of :meth:`coredis.Redis.evalsha`
"""
self.registered_client: coredis.client.Client[AnyStr] | None = registered_client
self.script: StringT
if not script:
raise RuntimeError("No script provided")
self.script = script
self.sha = hashlib.sha1(b(script)).hexdigest() # type: ignore
self.readonly = readonly
[docs]
def __call__(
self,
keys: Parameters[KeyT] | None = None,
args: Parameters[ValueT] | None = None,
client: coredis.client.Client[AnyStr] | None = None,
readonly: bool | None = None,
callback: Callable[..., CommandResponseT] = NoopCallback(),
) -> CommandRequest[CommandResponseT]:
"""
Executes the script registered in :paramref:`Script.script` using
:meth:`coredis.Redis.evalsha`. Additionally, if the script was not yet
registered on the instance, it will automatically do that as well
and cache the sha at :data:`Script.sha`
:param keys: The keys this script will reference
:param args: The arguments expected by the script
:param client: The redis client to use instead of :paramref:`Script.client`
:param readonly: If ``True`` forces the script to be called with
:meth:`coredis.Redis.evalsha_ro`
:param callback: a custom callback to call on the raw response from redis before
returning it.
"""
from coredis.patterns.pipeline import ClusterPipeline, Pipeline
if client is None:
client = self.registered_client
if not client:
raise RuntimeError(
"This instance is not bound to a redis client."
"Please provide a valid instance to execute the script with"
)
if readonly is None:
readonly = self.readonly
method = client.evalsha_ro if readonly else client.evalsha
if isinstance(client, (ClusterPipeline, Pipeline)):
cast(Pipeline[AnyStr], client).scripts.add(self)
return method(self.sha, keys=keys, args=args).transform(callback)
else:
return (
method(self.sha, keys=keys, args=args)
.retry(
ConstantRetryPolicy((NoScriptError,), retries=1, delay=0),
lambda _: client.script_load(self.script),
)
.transform(callback)
)
async def execute(
self,
keys: Parameters[KeyT] | None = None,
args: Parameters[ValueT] | None = None,
client: coredis.client.Client[AnyStr] | None = None,
readonly: bool | None = None,
) -> ResponseType:
"""
Executes the script registered in :paramref:`Script.script`
:meta private:
"""
return await self(keys, args, client, readonly)
@overload
def wraps(
self,
*,
client_arg: str | None = ...,
runtime_checks: bool = ...,
readonly: bool = ...,
) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]: ...
@overload
def wraps(
self,
*,
callback: Callable[..., CommandResponseT],
client_arg: str | None = ...,
runtime_checks: bool = ...,
readonly: bool = ...,
) -> Callable[[Callable[P, Awaitable[Any]]], Callable[P, Awaitable[CommandResponseT]]]: ...
[docs]
@versionadded(version="3.5.0")
def wraps(
self,
*,
callback: Callable[..., CommandResponseT] = NoopCallback(),
client_arg: str | None = None,
runtime_checks: bool = False,
readonly: bool = False,
) -> Any:
"""
Decorator for wrapping a regular python function, method or classmethod
signature with a :class:`~coredis.commands.script.Script`. This allows
exposing a strict signature instead of that which :meth:`__call__` provides.
The callable being decorated should **not** have an implementation as
it will never be called.
The main objective of the decorator is to allow you to have strict (and type safe)
signatures for wrappers for lua scripts. Internally the decorator separates
``keys`` from ``args`` before calling :meth:`coredis.Redis.evalsha`.
Mapping the decorated method's arguments to key providers is done by type
annotations: all parameters annotated as `KeyT` will be passed as keys, and the
rest will be passed as arguments.
By default the decorated method is bound to the :class:`coredis.client.Redis`
or :class:`coredis.client.RedisCluster` instance that the :class:`Script` instance
was instantiated with. This may however not be the instance you want to eventually
execute the method with. For such scenarios the decorated method can accept an additional
parameter which has the name declared by :paramref:`client_arg`).
The following example executes the script with the ``client`` instance used
to register the script. The ``key`` parameter is detected as a key provider
as it is annotated with the :data:`coredis.typing.KeyT` type, and ``value`` is
passed to redis as an ``arg``::
import coredis
from coredis.typing import KeyT, ValueT
from coredis.commands import Command Request
client = coredis.Redis()
@client.register_script("return {KEYS[1], ARGV[1]}").wraps()
def echo_key_value(key: KeyT, value: ValueT) -> CommandRequest[list[ValueT]]: ...
async with client:
res = await echo_key_value("co", "redis")
# [b"co", b"redis"]
Alternatively, the following example builds a class method that requires
the ``client`` to be passed in explicitly::
from coredis import Redis
from coredis.commands import Script
from coredis.typing import StringT
class ScriptProvider:
@classmethod
@Script(script="return KEYS[1]").wraps(
client_arg="client"
)
def echo_key(cls, client, key: KeyT) -> CommandRequest[StringT]: ...
@classmethod
@Script(script="return ARGS[1]").wraps(
client_arg="client"
)
def echo_arg(cls, client, value) -> CommandRequest[StringT]: ...
async with Redis() as client:
echoed = await ScriptProvider.echo_key(client, "coredis")
# b"coredis"
echoed = await ScriptProvider.echo_value(client, "coredis")
# b"coredis"
You can also pass a custom callback to execute on the return type, which will
be inferred as the return type rather than the annotation::
class MyCallback(ResponseCallback[Any, Any, int]):
def transform(self, response: ResponseType) -> int:
return sum([ord(c) for c in str(response)])
client = coredis.Redis(decode_responses=True)
async with client:
script = client.register_script("return {KEYS[1], ARGV[1]}")
# we use Any since return type will come from callback
@script.wraps(callback=MyCallback())
def echo_key_value(key: KeyT, value: ValueT) -> CommandRequest[int]: ...
res = await echo_key_value("co", "redis")
reveal_type(res) # int
# 1161
:param callback: a custom callback to execute on the returned value. When provided,
the callback's type will be inferred as the return type instead of the type from
the stub.
:param client_arg: The parameter of the decorator that will contain a client instance
to be used to execute the script.
:param runtime_checks: Whether to enable runtime type checking of input arguments
and return values. (requires :pypi:`beartype`). If :data:`False` the function will
still get runtime type checking if the environment configuration ``COREDIS_RUNTIME_CHECKS``
is set - for details see :ref:`handbook/typing:runtime type checking`.
:param readonly: If ``True`` forces this script to be called with
:meth:`coredis.Redis.evalsha_ro`
:return: A function that has a signature mirroring the decorated function.
"""
def wrapper(func: Callable[P, Awaitable[Any]]) -> Callable[P, Awaitable[CommandResponseT]]:
sig = inspect.signature(func)
first_arg = None
if args := list(sig.parameters.keys()):
first_arg = args[0]
runtime_check_wrapper = add_runtime_checks if not runtime_checks else safe_beartype
script_instance = self
key_params = [
n
for n, p in sig.parameters.items()
if p.annotation == "KeyT" or "KeyT" in get_args(p.annotation)
]
arg_fetch: dict[str, Callable[..., Parameters[Any]]] = {
n: (
(lambda v: [v])
if p.kind
in {
inspect.Parameter.POSITIONAL_ONLY,
inspect.Parameter.KEYWORD_ONLY,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
}
else (
(lambda v: list(itertools.chain.from_iterable(v.items())))
if p.kind == inspect.Parameter.VAR_KEYWORD
else lambda v: list(v)
)
)
for n, p in sig.parameters.items()
}
if first_arg in {"self", "cls"}:
arg_fetch.pop(first_arg)
if client_arg:
arg_fetch.pop(client_arg)
def split_args(
bound_arguments: inspect.BoundArguments,
) -> tuple[
Parameters[KeyT],
Parameters[ValueT],
coredis.client.Client[AnyStr] | None,
]:
bound_arguments.apply_defaults()
arguments = bound_arguments.arguments
keys: list[KeyT] = []
args: list[ValueT] = []
for name in sig.parameters:
if name not in arg_fetch:
continue
value = arg_fetch[name](arguments[name])
if name in key_params:
keys.extend(value)
else:
args.extend(value)
return (
keys,
args,
arguments.get(client_arg) if client_arg else None,
)
@runtime_check_wrapper
@functools.wraps(func)
def __inner(
*args: P.args,
**kwargs: P.kwargs,
) -> CommandRequest[CommandResponseT]:
keys, arguments, client = split_args(sig.bind(*args, **kwargs))
return script_instance(keys, arguments, client, readonly, callback=callback)
return __inner
return wrapper