from typing import TypeAlias, Awaitable, AsyncGenerator from urllib import parse from datetime import datetime, timedelta import asyncio import hiredis Resp3Repr: TypeAlias = str | int | bytes | bytearray | float | bool | None class _Ellipsis: """Represents an incomplete response of Redis.""" def _encode( arg: ( Resp3Repr | list[Resp3Repr] | tuple[Resp3Repr] | dict[Resp3Repr, Resp3Repr] | set[Resp3Repr] ) ) -> bytes: """Encode a value into its RESP3 representation. This function supports str (converted to bytes), int, bytes, bytearray, float, bool, and None as basic types. It also supports lists (arrays), dicts (maps), and sets. Args: arg (Resp3Repr | list[Resp3Repr] | tuple[Resp3Repr] | dict[Resp3Repr, Resp3Repr] | set[Resp3Repr]): The argument to encode. Raises: ValueError: The value provided is invalid or not encodable. Returns: bytes: The encoded representation of the value. """ if isinstance(arg, (str, int)): arg = str(arg).encode() if isinstance(arg, (bytes, bytearray)): return f"${len(arg)}\r\n".encode() + bytes(arg) + b"\r\n" elif isinstance(arg, int): return f":{arg}\r\n".encode() elif isinstance(arg, type(None)): return b"_\r\n" elif isinstance(arg, float): result = f",{arg}\r\n".encode() elif isinstance(arg, bool): if arg: return b"#t\r\n" return b"#f\r\n" elif isinstance(arg, (list, tuple)): return f"*{len(arg)}\r\n".encode() + b"".join([_encode(i) for i in arg]) elif isinstance(arg, dict): result = f"%{len(arg.keys())}\r\n".encode() for key, value in arg.items(): result += _encode(key) + _encode(value) return result elif isinstance(arg, set): result = f"~{len(arg)}\r\n".encode() for item in arg: result += _encode(item) return result raise ValueError(f"Value of type {type(arg)} is not encodable into RESP3.") def _get_secs(time: int | timedelta) -> int: """Converts an integer or timedelta into an amount of seconds. Args: time (int | timedelta): The number or timedelta to convert. Returns: int: The amount of seconds. """ if isinstance(time, timedelta): return int(time.total_seconds()) return time def _get_millis(time: int | timedelta) -> int: """Converts an integer or timedelta into an amount of milliseconds. Args: time (int | timedelta): The number or timedelta to convert. Returns: int: The amount of milliseconds. """ if isinstance(time, timedelta): return int(time.total_seconds() * 1000) return time def _get_timestamp_secs(dtime: datetime | int) -> int: """Converts a datetime or timestamp into it's seconds timestamp. Args: dtime (datetime | int): The datetime object to convert or an integer value which will be returned as-is. Returns: int: The seconds timestamp. """ if isinstance(dtime, datetime): return int(dtime.timestamp()) return dtime def _get_timestamp_millis(dtime: datetime | int) -> int: """Converts a datetime or timestamp into it's milliseconds timestamp. Args: dtime (datetime | int): The datetime object to convert or an integer value which will be returned as-is. Returns: int: The milliseconds timestamp. """ if isinstance(dtime, datetime): return int(dtime.timestamp() * 1000) return dtime class RedisError(Exception): """Raised whenever an error is raised from this library.""" class FatalRedisError(RedisError): """Raised whenever an error occurs that does not allow the connection to stay alive.""" class Runnable: """The base class for all Redis commands. This class is inherited by Connection, Pipeline, and Transaction to implement their own functionality. """ _client: "Redis" def __init__(self, client: "Redis") -> None: """Initializes the runner with the client's configuration. Args: client (Redis): The client that initialized the runner. """ self._client = client def command(self, *args: tuple[Resp3Repr]) -> Resp3Repr | Awaitable[Resp3Repr]: """Runs a command taking it as a list of arguments. For example: ``` await Runnable.command("SET", "key", "value") ``` This method is not implemented by default, so you need a Connection, a Pipeline, or a Transaction to use it. Returns: Resp3Repr | Awaitable[Resp3Repr]: The value(s) returned by the execution of the command. """ raise NotImplementedError( "This method needs to be implemented by a subclass of Runner." ) def ping(self, arg: str = "PONG") -> Resp3Repr: """Returns PONG if no argument is provided, otherwise return a copy of the argument as a bulk. Args: arg (str): The argument to return by the command. Returns: Resp3Repr: The response of the command. """ return self.command("PING", str(arg)) class RunnableGenericMixin: """Mixin containing all Redis generic commands.""" def delete(self, *keys: tuple[str]) -> None: """Deletes one or more keys.""" return self.command("DEL", *[str(key) for key in keys]) def exists(self, *keys: tuple[str]) -> int: """Returns if key exists. The user should be aware that if the same existing key is mentioned in the arguments multiple times, it will be counted multiple times. So if somekey exists, `EXISTS somekey somekey` will return 2. Args: *keys (tuple[str]): The keys to check for existance. Returns: int: The amount of keys that exist. """ return self.command("EXISTS", *[str(key) for key in keys]) class RunnableHashMixin: """Mixin containing all Redis hash commands.""" def hget(self, key: str, field: str) -> str | None: """Returns the value associated with field in the hash stored at key. Args: key (str): The key of the hash. field (str): The field of the hash to get. Returns: str | None: The value of the field, or None if it does not exist. """ return self.command("HGET", str(key), str(field)) def hgetall(self, key: str) -> dict[Resp3Repr, Resp3Repr]: """Returns all fields and values of the hash stored at key. Args: key (str): The key to get the fields and values from. Returns: dict[Resp3Repr, Resp3Repr]: A dict containing all the fields mapped to their respective values. """ return self.command("HGETALL", str(key)) def hincrby(self, key: str, field: str, increment: int = 1) -> int: """Increments the number stored at field in the hash stored at key by increment. If key does not exist, a new key holding a hash is created. If field does not exist the value is set to 0 before the operation is performed. The range of values supported by HINCRBY is limited to 64 bit signed integers. Args: key (str): The key of the hash whose field to increment. field (str): The field to increment. increment (int, optional): The amount to increment by. Defaults to 1. Returns: int: The amount returned """ return self.command("HINCRBY", str(key), str(field), str(increment)) def hset(self, key: str, mapping: dict[Resp3Repr, Resp3Repr]) -> int: """Sets the specified fields to their respective values in the hash stored at key. This command overwrites the values of specified fields that exist in the hash. If key doesn't exist, a new key holding a hash is created. Args: key (str): The key to set the values for. mapping (dict[Resp3Repr, Resp3Repr]): The fields mapped to their values. Returns: int: The number of fields that were added. """ items = [] for k, v in mapping.items(): items.append(k) items.append(v) return self.command("HSET", str(key), *items) class RunnableSetMixin: """Mixin containing all Redis set commands.""" def sadd(self, key: str, *values: tuple[str]) -> int: """Adds the specified members to the set stored at key. Specified members that are already a member of this set are ignored. If key does not exist, a new set is created before adding the specified members. An error is raised when the value stored at key is not a set. Args: key (str): The key to add the values to. *values (tuple[str]): The values to add to the set. Returns: int: The amount of new values that were added. """ return self.command("SADD", str(key), *values) def scard(self, key: str) -> int: """Returns the set cardinality (number of elements) of the set stored at key. Args: key (str): The key of the set to get the cardinality from. Returns: int: The number of elements, or 0 if it does not exist. """ return self.command("SCARD", str(key)) def smembers(self, key: str) -> set[str]: """Returns all the members of the set value stored at key. Args: key (str): The key of the set to get the members from. Returns: set[str]: The members of the set. """ return self.command("SMEMBERS", str(key)) def sismember(self, key: str, value: str) -> bool: """Returns whether the value is a member of the set. Args: key (str): The key of the set. value (str): The value to check whether it is a member of the set. Returns: bool: Whether the value is a member of the set. """ return self.command("SISMEMBER", str(key), str(value)) def spop(self, key: str, count: int = 1) -> str | list[str]: """Removes and returns one or more random members from the set value store at key.""" return self.command("SPOP", str(key), count) def srem(self, key: str, *values: tuple[str]) -> int: """Removes the specified members from the set stored at key. Specified members that are not a member of this set are ignored. If key does not exist, it is treated as an empty set and this command returns 0. An error is returned when the value stored at key is not a set. Args: key (str): The key of the set. *values (tuple[str]): The values to remove from the set. Returns: int: The amount of values that were removed from the set. """ return self.command("SREM", str(key), *values) class RunnableStringMixin: """Mixin containing all Redis string commands.""" def append(self, key: str, value: str) -> int: """Appends a string to an existing key, or creates a new one if it does not exist. If key already exists and is a string, this command appends the value at the end of the string. If key does not exist it is created and set as an empty string, so APPEND will be similar to SET in this special case. Args: key (str): The key to append to. value (str): The string to append. Returns: int: The new length of the value. """ return self.command("APPEND", str(key), str(value)) def decr(self, key: str) -> int: """Decrements the number stored at key by one. If the key does not exist, it is set to 0 before performing the operation. An error is returned if the key contains a value of the wrong type or contains a string that can not be represented as integer. This operation is limited to 64 bit signed integers. Args: key (str): The key to decrement. Returns: int: The decremented value. """ return self.command("DECR", str(key)) def decrby(self, key: str, decrement: int) -> int: """Reduces the value stored at the specified key by the specified decrement. If the key does not exist, it is initialized with a value of 0 before performing the operation. If the key's value is not of the correct type or cannot be represented as an integer, an error is returned. This operation is limited to 64-bit signed integers. Args: key (str): The key to decrement. decrement (int): The amount to decrement by. Returns: int: The decremented value. """ return self.command("DECRBY", str(key), decrement) def get(self, key: str) -> str | None: """Returns the string value of a key. If the key does not exist None is returned. An error is returned if the value stored at key is not a string, because GET only handles string values. Args: key (str): The key to get the value from. Returns: str | None: The value of the key, of None if it does not exist. """ return self.command("GET", str(key)) def getdel(self, key: str) -> str | None: """Returns the value of key and delete the key. This command is similar to GET, except for the fact that it also deletes the key on success (if and only if the key's value type is a string). Args: key (str): The key to get and delete. Returns: str | None: The value of the key. """ return self.command("GETDEL", str(key)) def set( self, key: str, value: str, nx: bool = False, xx: bool = False, get: bool = False, ex: timedelta | int | None = None, px: timedelta | int | None = None, exat: datetime | int | None = None, pxat: datetime | int | None = None, keepttl: bool = False, ) -> str | None: """Set key to hold the string value. If key already holds a value, it is overwritten, regardless of its type. Any previous time to live associated with the key is discarded on successful SET operation. Args: key (str): The key to set. value (str): The value to set to the specified key. nx (bool, optional): Only set the key if it does not already exist. Defaults to False. xx (bool, optional): Only set the key if it already exists. Defaults to False. get (bool, optional): Return the old string stored at key, or None if key did not exist. An error is returned and SET aborted if the value stored at key is not a string. Defaults to False. ex (timedelta | int | None, optional): Set the specified expire time, in seconds (a positive integer). Defaults to None. px (timedelta | int | None, optional): Set the specified expire time, in milliseconds (a positive integer). Defaults to None. exat (datetime | int | None, optional): Set the specified time at which the key will expire, in seconds (a positive integer) or a datetime object. Defaults to None. pxat (datetime | int | None, optional): Set the specified time at which the key will expire, in milliseconds (a positive integer) or a datetime object. Defaults to None. keepttl (bool, optional): Retain the time to live associated with the key. Defaults to False. Returns: str | None: The old key in case `get` was set, or None. """ if nx and xx: raise ValueError("Only one of XX or NX can be set, but not both.") if len([i for i in (ex, px, exat, pxat, keepttl) if i]) > 1: raise ValueError( "Only one of EX, PX, EXAT, PXAT, or KEEPTTL can be set, but not multiple of them." ) command = ["SET", str(key), str(value)] if nx: command += ["NX"] elif xx: command += ["XX"] if get: command += ["GET"] if ex: command += ["EX", str(_get_secs(ex))] elif px: command += ["PX", str(_get_millis(px))] elif exat: command += ["EXAT", str(_get_timestamp_secs(exat))] elif pxat: command += ["PXAT", str(_get_timestamp_millis(pxat))] elif keepttl: command += ["KEEPTTL"] return self.command(*command) class FullRunnable( Runnable, RunnableGenericMixin, RunnableHashMixin, RunnableSetMixin, RunnableStringMixin, ): """A runnable containing all the command mixins.""" class Connection(FullRunnable): _reader: asyncio.StreamReader _writer: asyncio.StreamWriter _parser: hiredis.Reader _lock: asyncio.Lock def __init__( self, client: "Redis", reader: asyncio.StreamReader, writer: asyncio.StreamWriter, parser: hiredis.Reader = hiredis.Reader(notEnoughData=_Ellipsis), ) -> None: """Initializes the connection object. Args: client (Redis): The client that initialized the connection. reader (asyncio.StreamReader): The socket stream reader. writer (asyncio.StreamWriter): The socket stream writer. parser (hiredis.Reader, optional): The RESP3 parser. Defaults to a basic implementation if not set. """ self._client = client self._reader = reader self._writer = writer self._parser = parser self._lock = asyncio.Lock() super().__init__(client) def release(self) -> None: """Returns the connection back to the pool.""" self._client._pool.append(self) def destroy(self) -> None: """Closes the connection and frees it's space in the client's pool.""" self._client._open_connections -= 1 self._writer.close() async def __aenter__(self) -> "Connection": return self async def __aexit__(self, *args) -> None: self.release() async def command(self, *args: list[Resp3Repr]) -> Resp3Repr: """Executes a command and returns it's response. Raises: hiredis.ReplyError: Any errors returned by Redis. """ args = list(args) command = _encode(args) async def run(): async with self._lock: self._writer.write(command) await self._writer.drain() while True: self._parser.feed(await self._reader.readuntil(b"\r\n")) result = self._parser.gets() if result == _Ellipsis: continue elif isinstance(result, hiredis.ReplyError): raise result return result try: return await asyncio.wait_for(run(), self._client._command_timeout) except ( hiredis.ProtocolError, asyncio.TimeoutError, ConnectionResetError, ) as exc: self.destroy() raise FatalRedisError() from exc async def hello(self, username: str = None, password: str = None) -> dict: """Sends the initial HELLO command to the server. Args: username (str, optional): The username to authenticate with. Defaults to None, or "default" if `password` is set. password (str, optional): The password to authenticate with. Defaults to None. Returns: dict: The metadata map returned by the server. """ command = ["HELLO", "3"] if password is not None: command += ["AUTH", username or "default", password] return await self.command(*command) class Pipeline(FullRunnable): """A Redis pipeline.""" _commands: list[tuple[Resp3Repr]] def __init__(self, client: "Redis") -> None: self._commands = [] super().__init__(client) def command(self, *args: tuple[Resp3Repr]) -> None: self._commands.append(args) async def execute(self): """Executes the pipeline. Returns: tuple[Resp3Repr]: The results returned by the execution of the commands. """ command = b"".join([_encode(comm) for comm in self._commands]) async with await self._client.connection() as conn: async def run(): async with conn._lock: conn._writer.write(command) await conn._writer.drain() results = [] while True: conn._parser.feed(await conn._reader.readuntil(b"\r\n")) result = conn._parser.gets() if result == _Ellipsis: continue elif isinstance(result, hiredis.ReplyError): raise RedisError( "An error was returned from the Redis server." ) from result results.append(result) if len(results) == len(self._commands): return results try: return tuple( await asyncio.wait_for(run(), self._client._command_timeout) ) except ( hiredis.ProtocolError, asyncio.TimeoutError, ConnectionResetError, ) as exc: conn.destroy() raise FatalRedisError() from exc class Redis(FullRunnable): _pool: list[Connection] = [] _open_connections: int = 0 def __init__( self, url: str = "", max_connections: int = 100, connection_timeout: int = 3, command_timeout: int = 3, encoding: str = "utf-8", ) -> None: """Initializes the client class. Args: url (str): The redis URL to connect to. Defaults to "redis://localhost:6379" max_connections (int, optional): The maximum amount of connections that can be opened. Defaults to 100. connection_timeout (int, optional): The amount of seconds to wait for a connection to establish before raising a TimeoutError. Defaults to 3. command_timeout (int, optional): The amount of seconds to wait for a command to execute before raising a TimeoutError. Defaults to 3. encoding (str, optional): The encoding with which to decode Redis response values. """ self._client = self self._lock = asyncio.Lock() self._max_connections = max_connections self._connection_timeout = connection_timeout self._command_timeout = command_timeout self._encoding = encoding parsed_url = parse.urlparse(url) self._host = parsed_url.hostname or "localhost" self._port = parsed_url.port or 6379 self._username = parsed_url.username self._password = parsed_url.password async def connection(self) -> Connection: """Returns an available connection or creates a new one if none are available and the pool has space remaining. Raises: ConnectionError: If the pool has reached it's max size and no connections are available. """ async with self._lock: try: connection = self._pool.pop() return connection except IndexError as exc: if self._open_connections >= self._max_connections: raise ConnectionError( "The pool has reached it's limit of connections and no connections are " "available at the moment." ) from exc reader, writer = await asyncio.wait_for( asyncio.open_connection(host=self._host, port=self._port), self._connection_timeout, ) connection = Connection( client=self, reader=reader, writer=writer, parser=hiredis.Reader( notEnoughData=_Ellipsis, encoding=self._encoding ), ) await connection.hello(username=self._username, password=self._password) self._open_connections += 1 return connection async def command(self, *args: tuple[Resp3Repr]) -> Awaitable[Resp3Repr]: """Executes a command with a free connection from the pool. Returns: Awaitable[Resp3Repr]: The command's response. """ async with await self.connection() as conn: return await conn.command(*args) def pipeline(self) -> Pipeline: """Creates a new pipeline.""" return Pipeline(self)