Add redis/redis-py@6f4ee2d

zinter (redis/redis-py#1520)

* Made some adjustments to typing in _zaggregate

Signed-off-by: Andrew-Chen-Wang <[email protected]>

Author: Andrew-Chen-Wang

aioredis/client.py

@@ -672,7 +672,8 @@ class Redis:
             "SDIFF SINTER SMEMBERS SUNION", lambda r: r and set(r) or set()
         ),
         **string_keys_to_dict(
-            "ZPOPMAX ZPOPMIN ZDIFF ZRANGE ZRANGEBYSCORE ZREVRANGE ZREVRANGEBYSCORE",
+            "ZPOPMAX ZPOPMIN ZINTER ZDIFF ZRANGE ZRANGEBYSCORE ZREVRANGE "
+            "ZREVRANGEBYSCORE",
             zset_score_pairs,
         ),
         **string_keys_to_dict(
@@ -3260,16 +3261,39 @@ def zincrby(self, name: KeyT, amount: float, value: EncodableT) -> Awaitable:
         """Increment the score of ``value`` in sorted set ``name`` by ``amount``"""
         return self.execute_command("ZINCRBY", name, amount, value)
 
+    def zinter(
+        self,
+        keys: KeysT,
+        aggregate: Optional[str] = None,
+        withscores: bool = False
+    ) -> Awaitable:
+        """
+        Return the intersect of multiple sorted sets specified by ``keys``.
+        With the ``aggregate`` option, it is possible to specify how the
+        results of the union are aggregated. This option defaults to SUM,
+        where the score of an element is summed across the inputs where it
+        exists. When this option is set to either MIN or MAX, the resulting
+        set will contain the minimum or maximum score of an element across
+        the inputs where it exists.
+        """
+        return self._zaggregate(
+            "ZINTER", None, keys, aggregate, withscores=withscores
+        )
+
     def zinterstore(
         self,
         dest: KeyT,
         keys: Union[Sequence[KeyT], Mapping[AnyKeyT, float]],
         aggregate: Optional[str] = None,
     ) -> Awaitable:
         """
-        Intersect multiple sorted sets specified by ``keys`` into
-        a new sorted set, ``dest``. Scores in the destination will be
-        aggregated based on the ``aggregate``, or SUM if none is provided.
+        Intersect multiple sorted sets specified by ``keys`` into a new
+        sorted set, ``dest``. Scores in the destination will be aggregated
+        based on the ``aggregate``. This option defaults to SUM, where the
+        score of an element is summed across the inputs where it exists.
+        When this option is set to either MIN or MAX, the resulting set will
+        contain the minimum or maximum score of an element across the inputs
+        where it exists.
         """
         return self._zaggregate("ZINTERSTORE", dest, keys, aggregate)
 
@@ -3593,11 +3617,15 @@ def zunionstore(
     def _zaggregate(
         self,
         command: str,
-        dest: KeyT,
+        dest: Optional[KeyT],
         keys: Union[Sequence[KeyT], Mapping[AnyKeyT, float]],
         aggregate: Optional[str] = None,
+        **options,
     ) -> Awaitable:
-        pieces: List[EncodableT] = [command, dest, len(keys)]
+        pieces: List[EncodableT] = [command]
+        if dest is not None:
+            pieces.append(dest)
+        pieces.append(len(keys))
         if isinstance(keys, dict):
             keys, weights = keys.keys(), keys.values()
         else:
@@ -3607,8 +3635,13 @@ def _zaggregate(
             pieces.append(b"WEIGHTS")
             pieces.extend(weights)
         if aggregate:
-            pieces.append(b"AGGREGATE")
-            pieces.append(aggregate)
+            if aggregate.upper() in ["SUM", "MIN", "MAX"]:
+                pieces.append(b'AGGREGATE')
+                pieces.append(aggregate)
+            else:
+                raise DataError("aggregate can be sum, min, or max")
+        if options.get("withscores", False):
+            pieces.append(b'WITHSCORES')
         return self.execute_command(*pieces)
 
     # HYPERLOGLOG COMMANDS

tests/test_commands.py

@@ -1596,6 +1596,28 @@ async def test_zlexcount(self, r: aioredis.Redis):
         assert await r.zlexcount("a", "-", "+") == 7
         assert await r.zlexcount("a", "[b", "[f") == 5
 
+    @skip_if_server_version_lt('6.2.0')
+    async def test_zinter(self, r: aioredis.Redis):
+        await r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 1})
+        await r.zadd('b', {'a1': 2, 'a2': 2, 'a3': 2})
+        await r.zadd('c', {'a1': 6, 'a3': 5, 'a4': 4})
+        assert await r.zinter(['a', 'b', 'c']) == [b'a3', b'a1']
+        # invalid aggregation
+        with pytest.raises(exceptions.DataError):
+            await r.zinter(['a', 'b', 'c'], aggregate='foo', withscores=True)
+        # aggregate with SUM
+        assert await r.zinter(['a', 'b', 'c'], withscores=True) \
+               == [(b'a3', 8), (b'a1', 9)]
+        # aggregate with MAX
+        assert await r.zinter(['a', 'b', 'c'], aggregate='MAX', withscores=True) \
+               == [(b'a3', 5), (b'a1', 6)]
+        # aggregate with MIN
+        assert await r.zinter(['a', 'b', 'c'], aggregate='MIN', withscores=True) \
+               == [(b'a1', 1), (b'a3', 1)]
+        # with weights
+        assert await r.zinter({'a': 1, 'b': 2, 'c': 3}, withscores=True) \
+               == [(b'a3', 20), (b'a1', 23)]
+
     async def test_zinterstore_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})