From f272a78b120714b203e3260f34aec7f44a2c570c Mon Sep 17 00:00:00 2001 From: Alexander Malev Date: Sun, 12 Nov 2023 15:01:41 +0300 Subject: [PATCH] Add expire, exists, delete, eval, z* and pf* (#33) --- redis_rs/client_async.pyi | 43 ++++++++++++ src/client.rs | 144 ++++++++++++++++++++++++++++++++++++-- tests/test_common.py | 27 +++++++ tests/test_l.py | 4 ++ tests/test_pf.py | 34 +++++++++ tests/test_script.py | 6 ++ tests/test_x.py | 4 ++ tests/test_z.py | 64 +++++++++++++++++ 8 files changed, 322 insertions(+), 4 deletions(-) create mode 100644 tests/test_common.py create mode 100644 tests/test_pf.py create mode 100644 tests/test_script.py create mode 100644 tests/test_z.py diff --git a/redis_rs/client_async.pyi b/redis_rs/client_async.pyi index 0e4211b..e8f20b7 100644 --- a/redis_rs/client_async.pyi +++ b/redis_rs/client_async.pyi @@ -12,6 +12,12 @@ class AsyncClient: async def fetch_float(self, *args: Arg) -> float: ... async def fetch_dict(self, *args: Arg, encoding: Optional[Encoding] = None) -> dict: ... async def fetch_scores(self, *args: Arg) -> Dict[str, float]: ... + async def exists(self, *keys: str) -> int: ... + async def expire(self, key: str, seconds: int, option: Optional[str] = None) -> int: ... + async def delete(self, *keys: str) -> int: ... + async def eval( + self, script: str, numkeys: int, *keys_and_args: Arg, encoding: Optional[Encoding] = None + ) -> Result: ... async def set(self, key: str, value: Arg) -> Result: ... async def get(self, key: str, *, encoding: Optional[Encoding] = None) -> Result: ... @overload @@ -37,6 +43,9 @@ class AsyncClient: encoding: Optional[Encoding] = None, ) -> List[Result]: ... async def llen(self, key: str) -> int: ... + async def pfadd(self, key: str, *elements: Arg) -> bool: ... + async def pfcount(self, *keys: str) -> int: ... + async def pfmerge(self, destkey: str, *sourcekeys: str) -> bool: ... @overload async def xadd( self, @@ -103,3 +112,37 @@ class AsyncClient: group: str, *id: Union[str, Literal["$"], Literal[0]], ) -> int: ... + @overload + async def zadd( + self, + key: str, + score: float, + value: str, + *pairs: Arg, + ) -> int: ... + @overload + async def zadd( + self, + key: str, + *args: Arg, + score: Optional[float] = None, + incr: Optional[float] = None, + ) -> int: ... + @overload + async def zrange( + self, + key: str, + start: Union[int, str] = 0, + stop: Union[int, str] = -1, + ) -> List[str]: ... + @overload + async def zrange( + self, + key: str, + start: Union[int, str] = 0, + stop: Union[int, str] = -1, + *, + withscores: Literal[True], + ) -> Dict[str, float]: ... + async def zcard(self, key: str) -> int: ... + async def zrem(self, key: str, *members: str) -> int: ... diff --git a/src/client.rs b/src/client.rs index 81b0b2a..791118d 100644 --- a/src/client.rs +++ b/src/client.rs @@ -146,10 +146,50 @@ impl Client { self.cr.fetch_int(py, cmd) } - #[pyo3(signature = (key))] - fn exists<'a>(&self, py: Python<'a>, key: types::Str) -> PyResult<&'a PyAny> { - let cmd = redis::cmd("EXISTS").arg(key).to_owned(); - self.cr.execute(py, cmd, String::default()) + #[pyo3(signature = (*keys))] + fn exists<'a>(&self, py: Python<'a>, keys: Vec) -> PyResult<&'a PyAny> { + let cmd = redis::cmd("EXISTS").arg(keys).to_owned(); + self.cr.fetch_int(py, cmd) + } + + #[pyo3(signature = (key, seconds, option = None))] + fn expire<'a>( + &self, + py: Python<'a>, + key: types::Str, + seconds: u64, + option: Option, + ) -> PyResult<&'a PyAny> { + let cmd = redis::cmd("EXPIRE") + .arg(key) + .arg(seconds) + .arg(option) + .to_owned(); + self.cr.fetch_bool(py, cmd) + } + + #[pyo3(signature = (*keys))] + fn delete<'a>(&self, py: Python<'a>, keys: Vec) -> PyResult<&'a PyAny> { + let cmd = redis::cmd("DEL").arg(keys).to_owned(); + self.cr.fetch_int(py, cmd) + } + + #[pyo3(signature = (script, numkeys, *args, **kwargs))] + fn eval<'a>( + &self, + py: Python<'a>, + script: types::Str, + numkeys: u8, + args: Vec, + kwargs: Option<&PyDict>, + ) -> PyResult<&'a PyAny> { + let cmd = redis::cmd("EVAL") + .arg(script) + .arg(numkeys) + .arg(args) + .to_owned(); + let encoding = self.get_encoding(kwargs); + self.cr.execute(py, cmd, encoding) } #[pyo3(signature = (key, value))] @@ -323,6 +363,34 @@ impl Client { self.cr.fetch_int(py, cmd) } + #[pyo3(signature = (key, *elements))] + fn pfadd<'a>( + &self, + py: Python<'a>, + key: types::Str, + elements: Vec, + ) -> PyResult<&'a PyAny> { + let cmd = redis::cmd("PFADD").arg(key).arg(elements).to_owned(); + self.cr.fetch_bool(py, cmd) + } + + #[pyo3(signature = (*keys))] + fn pfcount<'a>(&self, py: Python<'a>, keys: Vec) -> PyResult<&'a PyAny> { + let cmd = redis::cmd("PFCOUNT").arg(keys).to_owned(); + self.cr.fetch_int(py, cmd) + } + + #[pyo3(signature = (destkey, *sourcekeys))] + fn pfmerge<'a>( + &self, + py: Python<'a>, + destkey: types::Str, + sourcekeys: Vec, + ) -> PyResult<&'a PyAny> { + let cmd = redis::cmd("PFADD").arg(destkey).arg(sourcekeys).to_owned(); + self.cr.fetch_bool(py, cmd) + } + #[allow(clippy::too_many_arguments)] #[pyo3(signature = ( stream, *args, @@ -459,4 +527,72 @@ impl Client { let cmd = redis::cmd("XACK").arg(key).arg(group).arg(id).to_owned(); self.cr.fetch_int(py, cmd) } + + #[pyo3(signature = (key, *args, score = None, incr = None, encoding = None))] + fn zadd<'a>( + &self, + py: Python<'a>, + key: types::Str, + args: Vec, + score: Option, + incr: Option, + encoding: Option, + ) -> PyResult<&'a PyAny> { + let encoding = encoding.unwrap_or_default(); + let mut cmd = redis::cmd("ZADD").arg(key).to_owned(); + if let Some(incr) = incr { + cmd.arg(b"INCR").arg(incr); + } else { + cmd.arg(score); + } + cmd.arg(args); + self.cr.execute(py, cmd, encoding) + } + + #[pyo3(signature = ( + key, + start = types::Arg::Int(0), + stop = types::Arg::Int(-1), + *args, + withscores = false, + ))] + fn zrange<'a>( + &self, + py: Python<'a>, + key: types::Str, + start: types::Arg, + stop: types::Arg, + args: Vec, + withscores: bool, + ) -> PyResult<&'a PyAny> { + let mut cmd = redis::cmd("ZRANGE") + .arg(key) + .arg(start) + .arg(stop) + .arg(args) + .to_owned(); + if withscores { + cmd.arg(b"WITHSCORES"); + self.cr.fetch_dict(py, cmd, "float".to_string()) + } else { + self.cr.fetch_list(py, cmd) + } + } + + #[pyo3(signature = (key))] + fn zcard<'a>(&self, py: Python<'a>, key: types::Str) -> PyResult<&'a PyAny> { + let cmd = redis::cmd("ZCARD").arg(key).to_owned(); + self.cr.fetch_int(py, cmd) + } + + #[pyo3(signature = (key, *members))] + fn zrem<'a>( + &self, + py: Python<'a>, + key: types::Str, + members: Vec, + ) -> PyResult<&'a PyAny> { + let cmd = redis::cmd("ZREM").arg(key).arg(members).to_owned(); + self.cr.fetch_int(py, cmd) + } } diff --git a/tests/test_common.py b/tests/test_common.py new file mode 100644 index 0000000..f01321b --- /dev/null +++ b/tests/test_common.py @@ -0,0 +1,27 @@ +from uuid import uuid4 + +import redis_rs + + +async def test_exists(async_client: redis_rs.AsyncClient): + key = str(uuid4()) + await async_client.set(key, 1) + + result = await async_client.exists(key) + assert result == 1 + + +async def test_delete(async_client: redis_rs.AsyncClient): + key = str(uuid4()) + await async_client.set(key, 1) + + result = await async_client.delete(key) + assert result == 1 + + +async def test_expire(async_client: redis_rs.AsyncClient): + key = str(uuid4()) + await async_client.set(key, 1) + + result = await async_client.expire(key, 2) + assert result is True diff --git a/tests/test_l.py b/tests/test_l.py index 94f43bb..f356bfc 100644 --- a/tests/test_l.py +++ b/tests/test_l.py @@ -1,5 +1,7 @@ from uuid import uuid4 +import pytest + import redis_rs @@ -51,6 +53,7 @@ async def test_lrange(async_client: redis_rs.AsyncClient): assert result == [1, 2, 3] +@pytest.mark.redis(version=6.2) async def test_lpop(async_client: redis_rs.AsyncClient): key = str(uuid4()) @@ -73,6 +76,7 @@ async def test_lpop(async_client: redis_rs.AsyncClient): assert result == 1 +@pytest.mark.redis(version=6) async def test_blpop(async_client: redis_rs.AsyncClient): key1 = str(uuid4()) + "{a}" key2 = str(uuid4()) + "{a}" diff --git a/tests/test_pf.py b/tests/test_pf.py new file mode 100644 index 0000000..abadd4a --- /dev/null +++ b/tests/test_pf.py @@ -0,0 +1,34 @@ +from uuid import uuid4 + +import redis_rs + + +async def test_pfadd(async_client: redis_rs.AsyncClient): + key = str(uuid4()) + + result = await async_client.pfadd(key, 2) + assert result is True + + +async def test_pfcount(async_client: redis_rs.AsyncClient): + key = str(uuid4()) + + b = await async_client.pfadd(key, 2) + assert b is True + + result = await async_client.pfcount(key) + assert result == 1 + + +async def test_pfmerge(async_client: redis_rs.AsyncClient): + key1 = str(uuid4()) + key2 = str(uuid4()) + + b = await async_client.pfadd(key1, 2) + assert b is True + + b = await async_client.pfmerge(key2, key1) + assert b is True + + result = await async_client.pfcount(key2) + assert result == 1 diff --git a/tests/test_script.py b/tests/test_script.py new file mode 100644 index 0000000..27b6105 --- /dev/null +++ b/tests/test_script.py @@ -0,0 +1,6 @@ +import redis_rs + + +async def test_return(async_client: redis_rs.AsyncClient): + result = await async_client.eval("return ARGV[1]", 0, "a", encoding="utf-8") + assert result == "a" diff --git a/tests/test_x.py b/tests/test_x.py index bc027d0..b79fba3 100644 --- a/tests/test_x.py +++ b/tests/test_x.py @@ -6,6 +6,7 @@ import redis_rs +@pytest.mark.redis(version=6.2) async def test_xadd(async_client: redis_rs.AsyncClient): stream = str(uuid4()) @@ -25,6 +26,7 @@ async def test_xadd(async_client: redis_rs.AsyncClient): assert isinstance(ident, str) +@pytest.mark.redis(version=6.2) async def test_xadd_nomkstream(async_client: redis_rs.AsyncClient): stream = str(uuid4()) @@ -32,6 +34,7 @@ async def test_xadd_nomkstream(async_client: redis_rs.AsyncClient): assert ident is None +@pytest.mark.redis(version=6.2) async def test_xadd_flat(async_client: redis_rs.AsyncClient): stream = str(uuid4()) @@ -48,6 +51,7 @@ async def test_xadd_flat_id_star(id, async_client: redis_rs.AsyncClient): assert isinstance(ident, str), ident +@pytest.mark.redis(version=6.2) @pytest.mark.parametrize( "id", [ diff --git a/tests/test_z.py b/tests/test_z.py new file mode 100644 index 0000000..380b635 --- /dev/null +++ b/tests/test_z.py @@ -0,0 +1,64 @@ +from uuid import uuid4 + +import redis_rs + + +async def test_zadd(async_client: redis_rs.AsyncClient): + key = str(uuid4()) + + n = await async_client.zadd(key, 2, "a") + assert n == 1 + + n = await async_client.zadd(key, "b", score=3) + assert n == 1 + + result = await async_client.execute("ZRANGE", key, 0, -1, encoding="utf8") + assert result == ["a", "b"] + + +async def test_zrange(async_client: redis_rs.AsyncClient): + key = str(uuid4()) + + n = await async_client.zadd(key, "a", score=1) + assert n == 1 + + n = await async_client.zadd(key, "b", score=2) + assert n == 1 + + result = await async_client.zrange(key) + assert result == ["a", "b"] + + result_d = await async_client.zrange(key, withscores=True) + assert result_d == {"a": 1, "b": 2} + + +async def test_zcard(async_client: redis_rs.AsyncClient): + key = str(uuid4()) + + n = await async_client.zadd(key, "a", score=1) + assert n == 1 + + n = await async_client.zadd(key, "b", score=2) + assert n == 1 + + result = await async_client.zcard(key) + assert result == 2 + + +async def test_zrem(async_client: redis_rs.AsyncClient): + key = str(uuid4()) + + n = await async_client.zadd(key, "a", score=1) + assert n == 1 + + n = await async_client.zadd(key, "b", score=2) + assert n == 1 + + result = await async_client.zcard(key) + assert result == 2 + + result = await async_client.zrem(key, "a") + assert result == 1 + + result = await async_client.zcard(key) + assert result == 1