diff --git a/redis/client.py b/redis/client.py index 57932b9f73..8fef29bfba 100755 --- a/redis/client.py +++ b/redis/client.py @@ -595,8 +595,8 @@ class Redis: lambda r: r and set(r) or set() ), **string_keys_to_dict( - 'ZPOPMAX ZPOPMIN ZRANGE ZRANGEBYSCORE ZREVRANGE ZREVRANGEBYSCORE', - zset_score_pairs + 'ZPOPMAX ZPOPMIN ZDIFF ZRANGE ZRANGEBYSCORE ZREVRANGE ' + 'ZREVRANGEBYSCORE', zset_score_pairs ), **string_keys_to_dict('BZPOPMIN BZPOPMAX', \ lambda r: @@ -2902,6 +2902,24 @@ def zcount(self, name, min, max): """ return self.execute_command('ZCOUNT', name, min, max) + def zdiff(self, keys, withscores=False): + """ + Returns the difference between the first and all successive input + sorted sets provided in ``keys``. + """ + pieces = [len(keys), *keys] + if withscores: + pieces.append("WITHSCORES") + return self.execute_command("ZDIFF", *pieces) + + def zdiffstore(self, dest, keys): + """ + Computes the difference between the first and all successive input + sorted sets provided in ``keys`` and stores the result in ``dest``. + """ + pieces = [len(keys), *keys] + return self.execute_command("ZDIFFSTORE", dest, *pieces) + def zincrby(self, name, amount, value): "Increment the score of ``value`` in sorted set ``name`` by ``amount``" return self.execute_command('ZINCRBY', name, amount, value) diff --git a/tests/test_commands.py b/tests/test_commands.py index 27117e3188..5e5c11b469 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -1443,6 +1443,21 @@ def test_zcount(self, r): assert r.zcount('a', 1, '(' + str(2)) == 1 assert r.zcount('a', 10, 20) == 0 + @skip_if_server_version_lt('6.2.0') + def test_zdiff(self, r): + r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 3}) + r.zadd('b', {'a1': 1, 'a2': 2}) + assert r.zdiff(['a', 'b']) == [b'a3'] + assert r.zdiff(['a', 'b'], withscores=True) == [b'a3', b'3'] + + @skip_if_server_version_lt('6.2.0') + def test_zdiffstore(self, r): + r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 3}) + r.zadd('b', {'a1': 1, 'a2': 2}) + assert r.zdiffstore("out", ['a', 'b']) + assert r.zrange("out", 0, -1) == [b'a3'] + assert r.zrange("out", 0, -1, withscores=True) == [(b'a3', 3.0)] + def test_zincrby(self, r): r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 3}) assert r.zincrby('a', 1, 'a2') == 3.0