Add expire, exists, delete, eval, z* and pf* (#33)

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: ...

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<types::Str>) -> 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<types::Str>,
+    ) -> 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<types::Str>) -> 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<types::Arg>,
+        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<types::Arg>,
+    ) -> 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<types::Arg>) -> 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<types::Arg>,
+    ) -> 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<types::Arg>,
+        score: Option<f64>,
+        incr: Option<f64>,
+        encoding: Option<String>,
+    ) -> 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<types::Arg>,
+        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<types::Arg>,
+    ) -> PyResult<&'a PyAny> {
+        let cmd = redis::cmd("ZREM").arg(key).arg(members).to_owned();
+        self.cr.fetch_int(py, cmd)
+    }
 }

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

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}"

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

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"

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,13 +26,15 @@ 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())
 
     ident = await async_client.xadd(stream, {"a": "bcd"}, mkstream=False)
     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",
     [