Python: added zinter and zunion commands. (#1478)

* Python: added ZINTER and ZUNION commands.

CHANGELOG.md

@@ -11,6 +11,7 @@
 * Python: Added XLEN command ([#1503](https://github.com/aws/glide-for-redis/pull/1503))
 * Python: Added LASTSAVE command ([#1509](https://github.com/aws/glide-for-redis/pull/1509))
 * Python: Added GETDEL command ([#1514](https://github.com/aws/glide-for-redis/pull/1514))
+* Python: Added ZINTER, ZUNION commands ([#1478](https://github.com/aws/glide-for-redis/pull/1478))
 
 ## 0.4.1 (2024-02-06)
 

python/python/glide/async_commands/core.py

@@ -25,7 +25,7 @@
     RangeByScore,
     ScoreBoundary,
     ScoreFilter,
-    _create_z_cmd_store_args,
+    _create_zinter_zunion_cmd_args,
     _create_zrange_args,
 )
 from glide.constants import TOK, TResult
@@ -3228,6 +3228,75 @@ async def zdiffstore(self, destination: str, keys: List[str]) -> int:
             ),
         )
 
+    async def zinter(
+        self,
+        keys: List[str],
+    ) -> List[str]:
+        """
+        Computes the intersection of sorted sets given by the specified `keys` and returns a list of intersecting elements.
+        To get the scores as well, see `zinter_withscores`.
+        To store the result in a key as a sorted set, see `zinterstore`.
+
+        When in cluster mode, all keys in `keys` must map to the same hash slot.
+
+        See https://valkey.io/commands/zinter/ for more details.
+
+        Args:
+            keys (List[str]): The keys of the sorted sets.
+
+        Returns:
+            List[str]: The resulting array of intersecting elements.
+
+        Examples:
+            >>> await client.zadd("key1", {"member1": 10.5, "member2": 8.2})
+            >>> await client.zadd("key2", {"member1": 9.5})
+            >>> await client.zinter(["key1", "key2"])
+                ['member1']
+        """
+        return cast(
+            List[str],
+            await self._execute_command(RequestType.ZInter, [str(len(keys))] + keys),
+        )
+
+    async def zinter_withscores(
+        self,
+        keys: Union[List[str], List[Tuple[str, float]]],
+        aggregation_type: Optional[AggregationType] = None,
+    ) -> Mapping[str, float]:
+        """
+        Computes the intersection of sorted sets given by the specified `keys` and returns a sorted set of intersecting elements with scores.
+        To get the elements only, see `zinter`.
+        To store the result in a key as a sorted set, see `zinterstore`.
+
+        When in cluster mode, all keys in `keys` must map to the same hash slot.
+
+        See https://valkey.io/commands/zinter/ for more details.
+
+        Args:
+            keys (Union[List[str], List[Tuple[str, float]]]): The keys of the sorted sets with possible formats:
+                List[str] - for keys only.
+                List[Tuple[str, float]] - for weighted keys with score multipliers.
+            aggregation_type (Optional[AggregationType]): Specifies the aggregation strategy to apply
+                when combining the scores of elements. See `AggregationType`.
+
+        Returns:
+            Mapping[str, float]: The resulting sorted set with scores.
+
+        Examples:
+            >>> await client.zadd("key1", {"member1": 10.5, "member2": 8.2})
+            >>> await client.zadd("key2", {"member1": 9.5})
+            >>> await client.zinter_withscores(["key1", "key2"])
+                {'member1': 20}  # "member1" with score of 20 is the result
+            >>> await client.zinter_withscores(["key1", "key2"], AggregationType.MAX)
+                {'member1': 10.5}  # "member1" with score of 10.5 is the result.
+        """
+        args = _create_zinter_zunion_cmd_args(keys, aggregation_type)
+        args.append("WITHSCORES")
+        return cast(
+            Mapping[str, float],
+            await self._execute_command(RequestType.ZInter, args),
+        )
+
     async def zinterstore(
         self,
         destination: str,
@@ -3237,6 +3306,7 @@ async def zinterstore(
         """
         Computes the intersection of sorted sets given by the specified `keys` and stores the result in `destination`.
         If `destination` already exists, it is overwritten. Otherwise, a new sorted set will be created.
+        To get the result directly, see `zinter_withscores`.
 
         When in cluster mode, `destination` and all keys in `keys` must map to the same hash slot.
 
@@ -3246,7 +3316,7 @@ async def zinterstore(
             destination (str): The key of the destination sorted set.
             keys (Union[List[str], List[Tuple[str, float]]]): The keys of the sorted sets with possible formats:
                 List[str] - for keys only.
-                List[Tuple[str, float]]] - for weighted keys with score multipliers.
+                List[Tuple[str, float]] - for weighted keys with score multipliers.
             aggregation_type (Optional[AggregationType]): Specifies the aggregation strategy to apply
                 when combining the scores of elements. See `AggregationType`.
 
@@ -3259,18 +3329,87 @@ async def zinterstore(
             >>> await client.zinterstore("my_sorted_set", ["key1", "key2"])
                 1 # Indicates that the sorted set "my_sorted_set" contains one element.
             >>> await client.zrange_withscores("my_sorted_set", RangeByIndex(0, -1))
-                {'member1': 20}  # "member1"  is now stored in "my_sorted_set" with score of 20.
-            >>> await client.zinterstore("my_sorted_set", ["key1", "key2"] , AggregationType.MAX )
-                1 # Indicates that the sorted set "my_sorted_set" contains one element, and it's score is the maximum score between the sets.
+                {'member1': 20}  # "member1" is now stored in "my_sorted_set" with score of 20.
+            >>> await client.zinterstore("my_sorted_set", ["key1", "key2"], AggregationType.MAX)
+                1 # Indicates that the sorted set "my_sorted_set" contains one element, and its score is the maximum score between the sets.
             >>> await client.zrange_withscores("my_sorted_set", RangeByIndex(0, -1))
-                {'member1': 10.5}  # "member1"  is now stored in "my_sorted_set" with score of 10.5.
+                {'member1': 10.5}  # "member1" is now stored in "my_sorted_set" with score of 10.5.
         """
-        args = _create_z_cmd_store_args(destination, keys, aggregation_type)
+        args = _create_zinter_zunion_cmd_args(keys, aggregation_type, destination)
         return cast(
             int,
             await self._execute_command(RequestType.ZInterStore, args),
         )
 
+    async def zunion(
+        self,
+        keys: List[str],
+    ) -> List[str]:
+        """
+        Computes the union of sorted sets given by the specified `keys` and returns a list of union elements.
+        To get the scores as well, see `zunion_withscores`.
+        To store the result in a key as a sorted set, see `zunionstore`.
+
+        When in cluster mode, all keys in `keys` must map to the same hash slot.
+
+        See https://valkey.io/commands/zunion/ for more details.
+
+        Args:
+            keys (List[str]): The keys of the sorted sets.
+
+        Returns:
+            List[str]: The resulting array of union elements.
+
+        Examples:
+            >>> await client.zadd("key1", {"member1": 10.5, "member2": 8.2})
+            >>> await client.zadd("key2", {"member1": 9.5})
+            >>> await client.zunion(["key1", "key2"])
+                ['member1', 'member2']
+        """
+        return cast(
+            List[str],
+            await self._execute_command(RequestType.ZUnion, [str(len(keys))] + keys),
+        )
+
+    async def zunion_withscores(
+        self,
+        keys: Union[List[str], List[Tuple[str, float]]],
+        aggregation_type: Optional[AggregationType] = None,
+    ) -> Mapping[str, float]:
+        """
+        Computes the union of sorted sets given by the specified `keys` and returns a sorted set of union elements with scores.
+        To get the elements only, see `zunion`.
+        To store the result in a key as a sorted set, see `zunionstore`.
+
+        When in cluster mode, all keys in `keys` must map to the same hash slot.
+
+        See https://valkey.io/commands/zunion/ for more details.
+
+        Args:
+            keys (Union[List[str], List[Tuple[str, float]]]): The keys of the sorted sets with possible formats:
+                List[str] - for keys only.
+                List[Tuple[str, float]] - for weighted keys with score multipliers.
+            aggregation_type (Optional[AggregationType]): Specifies the aggregation strategy to apply
+                when combining the scores of elements. See `AggregationType`.
+
+        Returns:
+            Mapping[str, float]: The resulting sorted set with scores.
+
+        Examples:
+            >>> await client.zadd("key1", {"member1": 10.5, "member2": 8.2})
+            >>> await client.zadd("key2", {"member1": 9.5})
+            >>> await client.zunion_withscores(["key1", "key2"])
+                {'member1': 20, 'member2': 8.2}
+            >>> await client.zunion_withscores(["key1", "key2"], AggregationType.MAX)
+                {'member1': 10.5, 'member2': 8.2}
+        """
+        args = _create_zinter_zunion_cmd_args(keys, aggregation_type)
+        args.append("WITHSCORES")
+        return cast(
+            Mapping[str, float],
+            await self._execute_command(RequestType.ZUnion, args),
+        )
+
     async def zunionstore(
         self,
         destination: str,
@@ -3280,16 +3419,17 @@ async def zunionstore(
         """
         Computes the union of sorted sets given by the specified `keys` and stores the result in `destination`.
         If `destination` already exists, it is overwritten. Otherwise, a new sorted set will be created.
+        To get the result directly, see `zunion_withscores`.
 
         When in cluster mode, `destination` and all keys in `keys` must map to the same hash slot.
 
-        see https://valkey.io/commands/zunionstore/ for more details.
+        See https://valkey.io/commands/zunionstore/ for more details.
 
         Args:
             destination (str): The key of the destination sorted set.
             keys (Union[List[str], List[Tuple[str, float]]]): The keys of the sorted sets with possible formats:
                 List[str] - for keys only.
-                List[Tuple[str, float]]] - for weighted keys with score multipliers.
+                List[Tuple[str, float]] - for weighted keys with score multipliers.
             aggregation_type (Optional[AggregationType]): Specifies the aggregation strategy to apply
                 when combining the scores of elements. See `AggregationType`.
 
@@ -3300,15 +3440,15 @@ async def zunionstore(
             >>> await client.zadd("key1", {"member1": 10.5, "member2": 8.2})
             >>> await client.zadd("key2", {"member1": 9.5})
             >>> await client.zunionstore("my_sorted_set", ["key1", "key2"])
-                2 # Indicates that the sorted set "my_sorted_set" contains two element.
+                2 # Indicates that the sorted set "my_sorted_set" contains two elements.
             >>> await client.zrange_withscores("my_sorted_set", RangeByIndex(0, -1))
                 {'member1': 20, 'member2': 8.2}
-            >>> await client.zunionstore("my_sorted_set", ["key1", "key2"] , AggregationType.MAX )
-                2 # Indicates that the sorted set "my_sorted_set" contains two element, and each score is the maximum score between the sets.
+            >>> await client.zunionstore("my_sorted_set", ["key1", "key2"], AggregationType.MAX)
+                2 # Indicates that the sorted set "my_sorted_set" contains two elements, and each score is the maximum score between the sets.
             >>> await client.zrange_withscores("my_sorted_set", RangeByIndex(0, -1))
                 {'member1': 10.5, 'member2': 8.2}
         """
-        args = _create_z_cmd_store_args(destination, keys, aggregation_type)
+        args = _create_zinter_zunion_cmd_args(keys, aggregation_type, destination)
         return cast(
             int,
             await self._execute_command(RequestType.ZUnionStore, args),

python/python/glide/async_commands/sorted_set.py

@@ -220,12 +220,17 @@ def separate_keys(
     return key_list, weight_list
 
 
-def _create_z_cmd_store_args(
-    destination: str,
+def _create_zinter_zunion_cmd_args(
     keys: Union[List[str], List[Tuple[str, float]]],
     aggregation_type: Optional[AggregationType] = None,
+    destination: Optional[str] = None,
 ) -> List[str]:
-    args = [destination, str(len(keys))]
+    args = []
+
+    if destination:
+        args.append(destination)
+
+    args.append(str(len(keys)))
 
     only_keys, weights = separate_keys(keys)
 

python/python/glide/async_commands/transaction.py

@@ -24,7 +24,7 @@
     RangeByScore,
     ScoreBoundary,
     ScoreFilter,
-    _create_z_cmd_store_args,
+    _create_zinter_zunion_cmd_args,
     _create_zrange_args,
 )
 from glide.protobuf.redis_request_pb2 import RequestType
@@ -2272,6 +2272,47 @@ def zdiffstore(
             RequestType.ZDiffStore, [destination, str(len(keys))] + keys
         )
 
+    def zinter(
+        self: TTransaction,
+        keys: List[str],
+    ) -> TTransaction:
+        """
+        Computes the intersection of sorted sets given by the specified `keys` and returns a list of intersecting elements.
+
+        See https://valkey.io/commands/zinter/ for more details.
+
+        Args:
+            keys (List[str]): The keys of the sorted sets.
+
+        Command response:
+            List[str]: The resulting array of intersecting elements.
+        """
+        return self.append_command(RequestType.ZInter, [str(len(keys))] + keys)
+
+    def zinter_withscores(
+        self: TTransaction,
+        keys: Union[List[str], List[Tuple[str, float]]],
+        aggregation_type: Optional[AggregationType] = None,
+    ) -> TTransaction:
+        """
+        Computes the intersection of sorted sets given by the specified `keys` and returns a sorted set of intersecting elements with scores.
+
+        See https://valkey.io/commands/zinter/ for more details.
+
+        Args:
+            keys (Union[List[str], List[Tuple[str, float]]]): The keys of the sorted sets with possible formats:
+                List[str] - for keys only.
+                List[Tuple[str, float]] - for weighted keys with score multipliers.
+            aggregation_type (Optional[AggregationType]): Specifies the aggregation strategy to apply
+                when combining the scores of elements. See `AggregationType`.
+
+        Command response:
+            Mapping[str, float]: The resulting sorted set with scores.
+        """
+        args = _create_zinter_zunion_cmd_args(keys, aggregation_type)
+        args.append("WITHSCORES")
+        return self.append_command(RequestType.ZInter, args)
+
     def zinterstore(
         self: TTransaction,
         destination: str,
@@ -2297,9 +2338,50 @@ def zinterstore(
         Command response:
             int: The number of elements in the resulting sorted set stored at `destination`.
         """
-        args = _create_z_cmd_store_args(destination, keys, aggregation_type)
+        args = _create_zinter_zunion_cmd_args(keys, aggregation_type, destination)
         return self.append_command(RequestType.ZInterStore, args)
 
+    def zunion(
+        self: TTransaction,
+        keys: List[str],
+    ) -> TTransaction:
+        """
+        Computes the union of sorted sets given by the specified `keys` and returns a list of union elements.
+
+        See https://valkey.io/commands/zunion/ for more details.
+
+        Args:
+            keys (List[str]): The keys of the sorted sets.
+
+        Command response:
+            List[str]: The resulting array of union elements.
+        """
+        return self.append_command(RequestType.ZUnion, [str(len(keys))] + keys)
+
+    def zunion_withscores(
+        self: TTransaction,
+        keys: Union[List[str], List[Tuple[str, float]]],
+        aggregation_type: Optional[AggregationType] = None,
+    ) -> TTransaction:
+        """
+        Computes the union of sorted sets given by the specified `keys` and returns a sorted set of union elements with scores.
+
+        See https://valkey.io/commands/zunion/ for more details.
+
+        Args:
+            keys (Union[List[str], List[Tuple[str, float]]]): The keys of the sorted sets with possible formats:
+                List[str] - for keys only.
+                List[Tuple[str, float]] - for weighted keys with score multipliers.
+            aggregation_type (Optional[AggregationType]): Specifies the aggregation strategy to apply
+                when combining the scores of elements. See `AggregationType`.
+
+        Command response:
+            Mapping[str, float]: The resulting sorted set with scores.
+        """
+        args = _create_zinter_zunion_cmd_args(keys, aggregation_type)
+        args.append("WITHSCORES")
+        return self.append_command(RequestType.ZUnion, args)
+
     def zunionstore(
         self: TTransaction,
         destination: str,
@@ -2325,7 +2407,7 @@ def zunionstore(
         Command response:
             int: The number of elements in the resulting sorted set stored at `destination`.
         """
-        args = _create_z_cmd_store_args(destination, keys, aggregation_type)
+        args = _create_zinter_zunion_cmd_args(keys, aggregation_type, destination)
         return self.append_command(RequestType.ZUnionStore, args)
 
     def zrandmember(self: TTransaction, key: str) -> TTransaction: