From 6f4ee2dee63171afa4c59742a3726594894916ae Mon Sep 17 00:00:00 2001 From: Avital Fine <79420960+AvitalFineRedis@users.noreply.github.com> Date: Sun, 1 Aug 2021 13:35:55 +0300 Subject: zinter (#1520) * zinter * change options in _zaggregate * skip for previous versions * flake8 * validate the aggregate value * invalid aggregation * invalid aggregation * change options to get Co-authored-by: Chayim --- redis/client.py | 46 ++++++++++++++++++++++++++++++++++++---------- 1 file changed, 36 insertions(+), 10 deletions(-) (limited to 'redis/client.py') diff --git a/redis/client.py b/redis/client.py index 4ab1b1e..0d8a5c2 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 ZDIFF ZRANGE ZRANGEBYSCORE ZREVRANGE ' - 'ZREVRANGEBYSCORE', zset_score_pairs + 'ZPOPMAX ZPOPMIN ZINTER ZDIFF ZRANGE ZRANGEBYSCORE ' + 'ZREVRANGE ZREVRANGEBYSCORE', zset_score_pairs ), **string_keys_to_dict('BZPOPMIN BZPOPMAX', \ lambda r: @@ -2959,11 +2959,28 @@ class Redis: "Increment the score of ``value`` in sorted set ``name`` by ``amount``" return self.execute_command('ZINCRBY', name, amount, value) + def zinter(self, keys, aggregate=None, withscores=False): + """ + 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, keys, aggregate=None): """ - 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) @@ -3253,8 +3270,12 @@ class Redis: """ return self._zaggregate('ZUNIONSTORE', dest, keys, aggregate) - def _zaggregate(self, command, dest, keys, aggregate=None): - pieces = [command, dest, len(keys)] + def _zaggregate(self, command, dest, keys, aggregate=None, + **options): + pieces = [command] + if dest is not None: + pieces.append(dest) + pieces.append(len(keys)) if isinstance(keys, dict): keys, weights = keys.keys(), keys.values() else: @@ -3264,9 +3285,14 @@ class Redis: pieces.append(b'WEIGHTS') pieces.extend(weights) if aggregate: - pieces.append(b'AGGREGATE') - pieces.append(aggregate) - return self.execute_command(*pieces) + 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, **options) # HYPERLOGLOG COMMANDS def pfadd(self, name, *values): -- cgit v1.2.1