Add support for the set functions from issue jazzband#597

Co-authored-by: Ali Rezaei <[email protected]>

changelog.d/730.feature

@@ -0,0 +1 @@
+Support for sets and support basic operations, sadd, scard, sdiff, sdiffstore, sinter, sinterstore, smismember, sismember, smembers, smove, spop, srandmember, srem, sscan, sscan_iter, sunion, sunionstore

django_redis/cache.py

@@ -185,6 +185,74 @@ def close(self, **kwargs):
     def touch(self, *args, **kwargs):
         return self.client.touch(*args, **kwargs)
 
+    @omit_exception
+    def sadd(self, *args, **kwargs):
+        return self.client.sadd(*args, **kwargs)
+
+    @omit_exception
+    def scard(self, *args, **kwargs):
+        return self.client.scard(*args, **kwargs)
+
+    @omit_exception
+    def sdiff(self, *args, **kwargs):
+        return self.client.sdiff(*args, **kwargs)
+
+    @omit_exception
+    def sdiffstore(self, *args, **kwargs):
+        return self.client.sdiffstore(*args, **kwargs)
+
+    @omit_exception
+    def sinter(self, *args, **kwargs):
+        return self.client.sinter(*args, **kwargs)
+
+    @omit_exception
+    def sinterstore(self, *args, **kwargs):
+        return self.client.sinterstore(*args, **kwargs)
+
+    @omit_exception
+    def sismember(self, *args, **kwargs):
+        return self.client.sismember(*args, **kwargs)
+
+    @omit_exception
+    def smembers(self, *args, **kwargs):
+        return self.client.smembers(*args, **kwargs)
+
+    @omit_exception
+    def smove(self, *args, **kwargs):
+        return self.client.smove(*args, **kwargs)
+
+    @omit_exception
+    def spop(self, *args, **kwargs):
+        return self.client.spop(*args, **kwargs)
+
+    @omit_exception
+    def srandmember(self, *args, **kwargs):
+        return self.client.srandmember(*args, **kwargs)
+
+    @omit_exception
+    def srem(self, *args, **kwargs):
+        return self.client.srem(*args, **kwargs)
+
+    @omit_exception
+    def sscan(self, *args, **kwargs):
+        return self.client.sscan(*args, **kwargs)
+
+    @omit_exception
+    def sscan_iter(self, *args, **kwargs):
+        return self.client.sscan_iter(*args, **kwargs)
+
+    @omit_exception
+    def smismember(self, *args, **kwargs):
+        return self.client.smismember(*args, **kwargs)
+
+    @omit_exception
+    def sunion(self, *args, **kwargs):
+        return self.client.sunion(*args, **kwargs)
+
+    @omit_exception
+    def sunionstore(self, *args, **kwargs):
+        return self.client.sunionstore(*args, **kwargs)
+
     @omit_exception
     def hset(self, *args, **kwargs):
         return self.client.hset(*args, **kwargs)

django_redis/client/default.py

@@ -3,7 +3,7 @@
 import socket
 from collections import OrderedDict
 from contextlib import suppress
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Union
 
 from django.conf import settings
 from django.core.cache.backends.base import DEFAULT_TIMEOUT, BaseCache, get_key_func
@@ -778,6 +778,249 @@ def make_pattern(
 
         return CacheKey(self._backend.key_func(pattern, prefix, version_str))
 
+    def sadd(
+        self,
+        key: Any,
+        *values: Any,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> int:
+        if client is None:
+            client = self.get_client(write=True)
+
+        key = self.make_key(key, version=version)
+        values = [self.encode(value) for value in values]
+        return int(client.sadd(key, *values))
+
+    def scard(
+        self,
+        key: Any,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> int:
+        if client is None:
+            client = self.get_client(write=False)
+
+        key = self.make_key(key, version=version)
+        return int(client.scard(key))
+
+    def sdiff(
+        self,
+        *keys,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> Set:
+        if client is None:
+            client = self.get_client(write=False)
+
+        keys = [self.make_key(key, version=version) for key in keys]
+        return {self.decode(value) for value in client.sdiff(*keys)}
+
+    def sdiffstore(
+        self,
+        dest: Any,
+        *keys,
+        version_dest: Optional[int] = None,
+        version_keys: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> int:
+        if client is None:
+            client = self.get_client(write=True)
+
+        dest = self.make_key(dest, version=version_dest)
+        keys = [self.make_key(key, version=version_keys) for key in keys]
+        return int(client.sdiffstore(dest, *keys))
+
+    def sinter(
+        self,
+        *keys,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> Set:
+        if client is None:
+            client = self.get_client(write=False)
+
+        keys = [self.make_key(key, version=version) for key in keys]
+        return {self.decode(value) for value in client.sinter(*keys)}
+
+    def sinterstore(
+        self,
+        dest: Any,
+        *keys,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> int:
+        if client is None:
+            client = self.get_client(write=True)
+
+        dest = self.make_key(dest, version=version)
+        keys = [self.make_key(key, version=version) for key in keys]
+        return int(client.sinterstore(dest, *keys))
+
+    def smismember(
+        self,
+        key: Any,
+        *members,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> bool:
+        if client is None:
+            client = self.get_client(write=False)
+
+        key = self.make_key(key, version=version)
+        encoded_members = [self.encode(member) for member in members]
+
+        return [bool(value) for value in client.smismember(key, *encoded_members)]
+
+    def sismember(
+        self,
+        key: Any,
+        member: Any,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> bool:
+        if client is None:
+            client = self.get_client(write=False)
+
+        key = self.make_key(key, version=version)
+        member = self.encode(member)
+        return bool(client.sismember(key, member))
+
+    def smembers(
+        self,
+        key: Any,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> Set:
+        if client is None:
+            client = self.get_client(write=False)
+
+        key = self.make_key(key, version=version)
+        return {self.decode(value) for value in client.smembers(key)}
+
+    def smove(
+        self,
+        source: Any,
+        destination: Any,
+        member: Any,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> bool:
+        if client is None:
+            client = self.get_client(write=True)
+
+        source = self.make_key(source, version=version)
+        destination = self.make_key(destination)
+        member = self.encode(member)
+        return bool(client.smove(source, destination, member))
+
+    def spop(
+        self,
+        key: Any,
+        count: Optional[int] = None,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> Union[Set, Any]:
+        if client is None:
+            client = self.get_client(write=True)
+
+        key = self.make_key(key, version=version)
+        result = client.spop(key, count)
+        if isinstance(result, list):
+            return {self.decode(value) for value in result}
+        return self.decode(result)
+
+    def srandmember(
+        self,
+        key: Any,
+        count: Optional[int] = None,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> Union[Set, Any]:
+        if client is None:
+            client = self.get_client(write=False)
+
+        key = self.make_key(key, version=version)
+        result = client.srandmember(key, count)
+        if isinstance(result, list):
+            return {self.decode(value) for value in result}
+        return self.decode(result)
+
+    def srem(
+        self,
+        key: Any,
+        *members,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> int:
+        if client is None:
+            client = self.get_client(write=True)
+
+        key = self.make_key(key, version=version)
+        members = [self.encode(member) for member in members]
+        return int(client.srem(key, *members))
+
+    def sscan(
+        self,
+        key: Any,
+        match: Optional[str] = None,
+        count: Optional[int] = 10,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> Set[Any]:
+        if client is None:
+            client = self.get_client(write=False)
+
+        key = self.make_key(key, version=version)
+
+        cursor, result = client.sscan(
+            key, match=self.encode(match) if match else None, count=count
+        )
+        return {self.decode(value) for value in result}
+
+    def sscan_iter(
+        self,
+        key: Any,
+        match: Optional[str] = None,
+        count: Optional[int] = 10,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> Iterator[Any]:
+        if client is None:
+            client = self.get_client(write=False)
+
+        key = self.make_key(key, version=version)
+        for value in client.sscan_iter(
+            key, match=self.encode(match) if match else None, count=count
+        ):
+            yield self.decode(value)
+
+    def sunion(
+        self,
+        *keys,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> Set:
+        if client is None:
+            client = self.get_client(write=False)
+
+        keys = [self.make_key(key, version=version) for key in keys]
+        return {self.decode(value) for value in client.sunion(*keys)}
+
+    def sunionstore(
+        self,
+        destination: Any,
+        *keys,
+        version: Optional[int] = None,
+        client: Optional[Redis] = None,
+    ) -> int:
+        if client is None:
+            client = self.get_client(write=True)
+
+        destination = self.make_key(destination, version=version)
+        keys = [self.make_key(key, version=version) for key in keys]
+        return int(client.sunionstore(destination, *keys))
+
     def close(self) -> None:
         close_flag = self._options.get(
             "CLOSE_CONNECTION",