Add redis/redis-py@8c82e42

Zunion (redis/redis-py#1512)

Signed-off-by: Andrew-Chen-Wang <acwangpython@gmail.com>

aioredis/client.py

@@ -692,7 +692,7 @@ class Redis:
             "SDIFF SINTER SMEMBERS SUNION", lambda r: r and set(r) or set()
         ),
         **string_keys_to_dict(
-            "ZPOPMAX ZPOPMIN ZINTER ZDIFF ZRANGE ZRANGEBYSCORE ZREVRANGE "
+            "ZPOPMAX ZPOPMIN ZINTER ZDIFF ZUNION ZRANGE ZRANGEBYSCORE ZREVRANGE "
             "ZREVRANGEBYSCORE",
             zset_score_pairs,
         ),
@@ -3786,6 +3786,21 @@ def zscore(self, name: str, value: EncodableT) -> Awaitable:
         """Return the score of element ``value`` in sorted set ``name``"""
         return self.execute_command("ZSCORE", name, value)
 
+    def zunion(
+        self,
+        keys: Union[Sequence[KeyT], Mapping[AnyKeyT, float]],
+        aggregate: Optional[str] = None,
+        withscores: bool = False
+    ) -> Awaitable:
+        """
+        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: KeyT,

tests/test_commands.py

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