jazzband#597 add support for the set functions

django_redis/cache.py

@@ -186,6 +186,71 @@ 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, Tuple, Union, Set
 
 from django.conf import settings
 from django.core.cache.backends.base import DEFAULT_TIMEOUT, BaseCache, get_key_func
@@ -778,6 +778,246 @@ 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 type(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 type(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.decode(member) for member in members]
+        return int(client.srem(key, *members))
+
+    def sscan(
+            self,
+            key: Any,
+            match: Optional[str] = None,
+            count: Optional[int] = None,
+            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), count=count)
+        return {self.decode(value) for value in result}
+
+    def sscan_iter(
+            self,
+            key: Any,
+            match: Optional[str] = None,
+            count: Optional[int] = None,
+            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=match, 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",

django_redis/compressors/lz4.py

@@ -16,5 +16,5 @@ def compress(self, value: bytes) -> bytes:
     def decompress(self, value: bytes) -> bytes:
         try:
             return _decompress(value)
-        except Exception as e:  # noqa: BLE001
+        except Exception as e:
             raise CompressorError from e