2022-06-24 17:14:37 +02:00

249 lines
9.0 KiB
Python

from itertools import chain
from aioredis.util import wait_convert, wait_ok, _NOTSET
class StringCommandsMixin:
"""String commands mixin.
For commands details see: http://redis.io/commands/#string
"""
SET_IF_NOT_EXIST = 'SET_IF_NOT_EXIST' # NX
SET_IF_EXIST = 'SET_IF_EXIST' # XX
def append(self, key, value):
"""Append a value to key."""
return self.execute(b'APPEND', key, value)
def bitcount(self, key, start=None, end=None):
"""Count set bits in a string.
:raises TypeError: if only start or end specified.
"""
if start is None and end is not None:
raise TypeError("both start and stop must be specified")
elif start is not None and end is None:
raise TypeError("both start and stop must be specified")
elif start is not None and end is not None:
args = (start, end)
else:
args = ()
return self.execute(b'BITCOUNT', key, *args)
def bitfield(self):
raise NotImplementedError()
def bitop_and(self, dest, key, *keys):
"""Perform bitwise AND operations between strings."""
return self.execute(b'BITOP', b'AND', dest, key, *keys)
def bitop_or(self, dest, key, *keys):
"""Perform bitwise OR operations between strings."""
return self.execute(b'BITOP', b'OR', dest, key, *keys)
def bitop_xor(self, dest, key, *keys):
"""Perform bitwise XOR operations between strings."""
return self.execute(b'BITOP', b'XOR', dest, key, *keys)
def bitop_not(self, dest, key):
"""Perform bitwise NOT operations between strings."""
return self.execute(b'BITOP', b'NOT', dest, key)
def bitpos(self, key, bit, start=None, end=None):
"""Find first bit set or clear in a string.
:raises ValueError: if bit is not 0 or 1
"""
if bit not in (1, 0):
raise ValueError("bit argument must be either 1 or 0")
bytes_range = []
if start is not None:
bytes_range.append(start)
if end is not None:
if start is None:
bytes_range = [0, end]
else:
bytes_range.append(end)
return self.execute(b'BITPOS', key, bit, *bytes_range)
def decr(self, key):
"""Decrement the integer value of a key by one."""
return self.execute(b'DECR', key)
def decrby(self, key, decrement):
"""Decrement the integer value of a key by the given number.
:raises TypeError: if decrement is not int
"""
if not isinstance(decrement, int):
raise TypeError("decrement must be of type int")
return self.execute(b'DECRBY', key, decrement)
def get(self, key, *, encoding=_NOTSET):
"""Get the value of a key."""
return self.execute(b'GET', key, encoding=encoding)
def getbit(self, key, offset):
"""Returns the bit value at offset in the string value stored at key.
:raises TypeError: if offset is not int
:raises ValueError: if offset is less than 0
"""
if not isinstance(offset, int):
raise TypeError("offset argument must be int")
if offset < 0:
raise ValueError("offset must be greater equal 0")
return self.execute(b'GETBIT', key, offset)
def getrange(self, key, start, end, *, encoding=_NOTSET):
"""Get a substring of the string stored at a key.
:raises TypeError: if start or end is not int
"""
if not isinstance(start, int):
raise TypeError("start argument must be int")
if not isinstance(end, int):
raise TypeError("end argument must be int")
return self.execute(b'GETRANGE', key, start, end, encoding=encoding)
def getset(self, key, value, *, encoding=_NOTSET):
"""Set the string value of a key and return its old value."""
return self.execute(b'GETSET', key, value, encoding=encoding)
def incr(self, key):
"""Increment the integer value of a key by one."""
return self.execute(b'INCR', key)
def incrby(self, key, increment):
"""Increment the integer value of a key by the given amount.
:raises TypeError: if increment is not int
"""
if not isinstance(increment, int):
raise TypeError("increment must be of type int")
return self.execute(b'INCRBY', key, increment)
def incrbyfloat(self, key, increment):
"""Increment the float value of a key by the given amount.
:raises TypeError: if increment is not int
"""
if not isinstance(increment, float):
raise TypeError("increment must be of type int")
fut = self.execute(b'INCRBYFLOAT', key, increment)
return wait_convert(fut, float)
def mget(self, key, *keys, encoding=_NOTSET):
"""Get the values of all the given keys."""
return self.execute(b'MGET', key, *keys, encoding=encoding)
def mset(self, *args):
"""Set multiple keys to multiple values or unpack dict to keys & values.
:raises TypeError: if len of args is not event number
:raises TypeError: if len of args equals 1 and it is not a dict
"""
data = args
if len(args) == 1:
if not isinstance(args[0], dict):
raise TypeError("if one arg it should be a dict")
data = chain.from_iterable(args[0].items())
elif len(args) % 2 != 0:
raise TypeError("length of pairs must be even number")
fut = self.execute(b'MSET', *data)
return wait_ok(fut)
def msetnx(self, key, value, *pairs):
"""Set multiple keys to multiple values,
only if none of the keys exist.
:raises TypeError: if len of pairs is not event number
"""
if len(pairs) % 2 != 0:
raise TypeError("length of pairs must be even number")
return self.execute(b'MSETNX', key, value, *pairs)
def psetex(self, key, milliseconds, value):
"""Set the value and expiration in milliseconds of a key.
:raises TypeError: if milliseconds is not int
"""
if not isinstance(milliseconds, int):
raise TypeError("milliseconds argument must be int")
fut = self.execute(b'PSETEX', key, milliseconds, value)
return wait_ok(fut)
def set(self, key, value, *, expire=0, pexpire=0, exist=None):
"""Set the string value of a key.
:raises TypeError: if expire or pexpire is not int
"""
if expire and not isinstance(expire, int):
raise TypeError("expire argument must be int")
if pexpire and not isinstance(pexpire, int):
raise TypeError("pexpire argument must be int")
args = []
if expire:
args[:] = [b'EX', expire]
if pexpire:
args[:] = [b'PX', pexpire]
if exist is self.SET_IF_EXIST:
args.append(b'XX')
elif exist is self.SET_IF_NOT_EXIST:
args.append(b'NX')
fut = self.execute(b'SET', key, value, *args)
return wait_ok(fut)
def setbit(self, key, offset, value):
"""Sets or clears the bit at offset in the string value stored at key.
:raises TypeError: if offset is not int
:raises ValueError: if offset is less than 0 or value is not 0 or 1
"""
if not isinstance(offset, int):
raise TypeError("offset argument must be int")
if offset < 0:
raise ValueError("offset must be greater equal 0")
if value not in (0, 1):
raise ValueError("value argument must be either 1 or 0")
return self.execute(b'SETBIT', key, offset, value)
def setex(self, key, seconds, value):
"""Set the value and expiration of a key.
If seconds is float it will be multiplied by 1000
coerced to int and passed to `psetex` method.
:raises TypeError: if seconds is neither int nor float
"""
if isinstance(seconds, float):
return self.psetex(key, int(seconds * 1000), value)
if not isinstance(seconds, int):
raise TypeError("milliseconds argument must be int")
fut = self.execute(b'SETEX', key, seconds, value)
return wait_ok(fut)
def setnx(self, key, value):
"""Set the value of a key, only if the key does not exist."""
fut = self.execute(b'SETNX', key, value)
return wait_convert(fut, bool)
def setrange(self, key, offset, value):
"""Overwrite part of a string at key starting at the specified offset.
:raises TypeError: if offset is not int
:raises ValueError: if offset less than 0
"""
if not isinstance(offset, int):
raise TypeError("offset argument must be int")
if offset < 0:
raise ValueError("offset must be greater equal 0")
return self.execute(b'SETRANGE', key, offset, value)
def strlen(self, key):
"""Get the length of the value stored in a key."""
return self.execute(b'STRLEN', key)