Zunion (#1522)

* zinter

* change options in _zaggregate

* skip for previous versions

* add client function

* validate the aggregate value

* change options to get

* add more aggregate tests

* add weights guidance

Author: AvitalFineRedis

redis/client.py

@@ -615,7 +615,7 @@ class Redis:
             lambda r: r and set(r) or set()
         ),
         **string_keys_to_dict(
-            'ZPOPMAX ZPOPMIN ZINTER ZDIFF ZRANGE ZRANGEBYSCORE '
+            'ZPOPMAX ZPOPMIN ZINTER ZDIFF ZUNION ZRANGE ZRANGEBYSCORE '
             'ZREVRANGE ZREVRANGEBYSCORE', zset_score_pairs
         ),
         **string_keys_to_dict('BZPOPMIN BZPOPMAX', \
@@ -3416,6 +3416,16 @@ def zscore(self, name, value):
         "Return the score of element ``value`` in sorted set ``name``"
         return self.execute_command('ZSCORE', name, value)
 
+    def zunion(self, keys, aggregate=None, withscores=False):
+        """
+        Return the union of multiple sorted sets specified by ``keys``.
+        ``keys`` can be provided as dictionary of keys and their weights.
+        Scores will be aggregated based on the ``aggregate``, or SUM if
+        none is provided.
+        """
+        return self._zaggregate('ZUNION', None, keys, aggregate,
+                                withscores=withscores)
+
     def zunionstore(self, dest, keys, aggregate=None):
         """
         Union multiple sorted sets specified by ``keys`` into

tests/test_commands.py

@@ -1813,6 +1813,26 @@ def test_zscore(self, r):
         assert r.zscore('a', 'a2') == 2.0
         assert r.zscore('a', 'a4') is None
 
+    @skip_if_server_version_lt('6.2.0')
+    def test_zunion(self, r):
+        r.zadd('a', {'a1': 1, 'a2': 1, 'a3': 1})
+        r.zadd('b', {'a1': 2, 'a2': 2, 'a3': 2})
+        r.zadd('c', {'a1': 6, 'a3': 5, 'a4': 4})
+        # sum
+        assert r.zunion(['a', 'b', 'c']) == \
+            [b'a2', b'a4', b'a3', b'a1']
+        assert r.zunion(['a', 'b', 'c'], withscores=True) == \
+            [(b'a2', 3), (b'a4', 4), (b'a3', 8), (b'a1', 9)]
+        # max
+        assert r.zunion(['a', 'b', 'c'], aggregate='MAX', withscores=True)\
+               == [(b'a2', 2), (b'a4', 4), (b'a3', 5), (b'a1', 6)]
+        # min
+        assert r.zunion(['a', 'b', 'c'], aggregate='MIN', withscores=True)\
+               == [(b'a1', 1), (b'a2', 1), (b'a3', 1), (b'a4', 4)]
+        # with weight
+        assert r.zunion({'a': 1, 'b': 2, 'c': 3}, withscores=True)\
+               == [(b'a2', 5), (b'a4', 12), (b'a3', 20), (b'a1', 23)]
+
     def test_zunionstore_sum(self, r):
         r.zadd('a', {'a1': 1, 'a2': 1, 'a3': 1})
         r.zadd('b', {'a1': 2, 'a2': 2, 'a3': 2})