Implement zpopmin/zpopmax

REDIS_COMMANDS.md

@@ -155,6 +155,8 @@ list of [unimplemented commands](#unimplemented-commands).
  * zincrby
  * zinterstore
  * zlexcount
+ * zpopmax
+ * zpopmin
  * zrange
  * zrangebylex
  * zrangebyscore
@@ -322,8 +324,6 @@ All of the redis commands are implemented in fakeredis with these exceptions:
  * zintercard
  * zmpop
  * zmscore
- * zpopmax
- * zpopmin
  * zrandmember
  * zrangestore
  * zunion

fakeredis/_basefakesocket.py

@@ -461,3 +461,15 @@ def _expireat(self, key, timestamp, *args):
             return 0
         key.expireat = timestamp
         return 1
+
+    def _zpop(self, key, count, reverse):
+        zset = key.value
+        members = list(zset)
+        if reverse:
+            members.reverse()
+        members = members[:count]
+        res = [(bytes(member), self._encodefloat(zset.get(member), True)) for member in members]
+        res = list(itertools.chain.from_iterable(res))
+        for item in members:
+            zset.discard(item)
+        return res

fakeredis/_fakesocket.py

@@ -1084,7 +1084,16 @@ def pfmerge(self, dest, *sources):
         return OK
 
     # Sorted set commands
-    # TODO: [b]zpopmin/zpopmax,
+
+    @command((Key(ZSet),), (Int,))
+    def zpopmin(self, key, count=1):
+        return self._zpop(key, count, False)
+
+    @command((Key(ZSet),), (Int,))
+    def zpopmax(self, key, count=1):
+        return self._zpop(key, count, True)
+
+    # TODO: bzpopmin/bzpopmax,
 
     @staticmethod
     def _limit_items(items, offset, count):

fakeredis/_helpers.py

@@ -214,8 +214,8 @@ def __eq__(self, other):
 def valid_response_type(value, nested=False):
     if isinstance(value, NoResponse) and not nested:
         return True
-    if value is not None and not isinstance(value, (bytes, SimpleString, SimpleError,
-                                                    int, list)):
+    if (value is not None
+            and not isinstance(value, (bytes, SimpleString, SimpleError, int, list))):
         return False
     if isinstance(value, list):
         if any(not valid_response_type(item, True) for item in value):

test/test_sorted_set_cmds.py

@@ -0,0 +1,37 @@
+import redis.client
+from packaging.version import Version
+
+import testtools
+
+REDIS_VERSION = Version(redis.__version__)
+
+
+def test_zpopmin(r):
+    testtools.zadd(r, 'foo', {'one': 1})
+    testtools.zadd(r, 'foo', {'two': 2})
+    testtools.zadd(r, 'foo', {'three': 3})
+    assert r.zpopmin('foo', count=2) == [(b'one', 1.0), (b'two', 2.0)]
+    assert r.zpopmin('foo', count=2) == [(b'three', 3.0)]
+
+
+def test_zpopmin_too_many(r):
+    testtools.zadd(r, 'foo', {'one': 1})
+    testtools.zadd(r, 'foo', {'two': 2})
+    testtools.zadd(r, 'foo', {'three': 3})
+    assert r.zpopmin('foo', count=5) == [(b'one', 1.0), (b'two', 2.0), (b'three', 3.0)]
+
+
+def test_zpopmax(r):
+    testtools.zadd(r, 'foo', {'one': 1})
+    testtools.zadd(r, 'foo', {'two': 2})
+    testtools.zadd(r, 'foo', {'three': 3})
+    assert r.zpopmax('foo', count=2) == [(b'three', 3.0), (b'two', 2.0)]
+    assert r.zpopmax('foo', count=2) == [(b'one', 1.0)]
+
+
+def test_zpopmax_too_many(r):
+    testtools.zadd(r, 'foo', {'one': 1})
+    testtools.zadd(r, 'foo', {'two': 2})
+    testtools.zadd(r, 'foo', {'three': 3})
+    assert r.zpopmax('foo', count=5) == [(b'three', 3.0), (b'two', 2.0), (b'one', 1.0), ]
+