Implement BITOP (#110)

close #110

Author: cunla

REDIS_COMMANDS.md

@@ -68,6 +68,8 @@ list of [unimplemented commands](#unimplemented-commands).
 ### bitmap
  * [BITCOUNT](https://redis.io/commands/bitcount/)
    Count set bits in a string
+ * [BITOP](https://redis.io/commands/bitop/)
+   Perform bitwise operations between strings
  * [GETBIT](https://redis.io/commands/getbit/)
    Returns the bit value at offset in the string value stored at key
  * [SETBIT](https://redis.io/commands/setbit/)
@@ -613,8 +615,6 @@ All the redis commands are implemented in fakeredis with these exceptions:
    Perform arbitrary bitfield integer operations on strings
  * [BITFIELD_RO](https://redis.io/commands/bitfield_ro/)
    Perform arbitrary bitfield integer operations on strings. Read-only variant of BITFIELD
- * [BITOP](https://redis.io/commands/bitop/)
-   Perform bitwise operations between strings
  * [BITPOS](https://redis.io/commands/bitpos/)
    Find first bit set or clear in a string
 

fakeredis/_msgs.py

@@ -6,6 +6,7 @@
 INVALID_OFFSET_MSG = "ERR offset is out of range"
 INVALID_BIT_OFFSET_MSG = "ERR bit offset is not an integer or out of range"
 INVALID_BIT_VALUE_MSG = "ERR bit is not an integer or out of range"
+BITOP_NOT_ONE_KEY_ONLY = "ERR BITOP NOT must be called with a single source key"
 INVALID_DB_MSG = "ERR DB index is out of range"
 INVALID_MIN_MAX_FLOAT_MSG = "ERR min or max is not a float"
 INVALID_MIN_MAX_STR_MSG = "ERR min or max not a valid string range item"

fakeredis/commands_mixins/bitmap_mixin.py

@@ -5,7 +5,7 @@
 
 class BitmapCommandsMixin:
     # BITMAP commands
-    # TODO: bitfield, bitfield_ro, bitop, bitpos
+    # TODO: bitfield, bitfield_ro, bitpos
 
     @command((Key(bytes, 0),), (bytes,))
     def bitcount(self, key, *args):
@@ -70,3 +70,38 @@ def setbit(self, key, offset, value):
         reconstructed[byte] = new_byte
         key.update(bytes(reconstructed))
         return old_value
+
+    @staticmethod
+    def _bitop(op, *keys):
+        value = keys[0].value
+        if not isinstance(value, bytes):
+            raise SimpleError(msgs.WRONGTYPE_MSG)
+        ans = keys[0].value
+        i = 1
+        while i < len(keys):
+            value = keys[i].value if keys[i].value is not None else b''
+            if not isinstance(value, bytes):
+                raise SimpleError(msgs.WRONGTYPE_MSG)
+            ans = bytes(op(a, b) for a, b in zip(ans, value))
+            i += 1
+        return ans
+
+    @command((bytes, Key(), Key(bytes)), (Key(bytes),))
+    def bitop(self, op_name, dst, *keys):
+        if len(keys) == 0:
+            raise SimpleError(msgs.WRONG_ARGS_MSG6.format('bitop'))
+        if casematch(op_name, b'and'):
+            res = self._bitop(lambda a, b: a & b, *keys)
+        elif casematch(op_name, b'or'):
+            res = self._bitop(lambda a, b: a | b, *keys)
+        elif casematch(op_name, b'xor'):
+            res = self._bitop(lambda a, b: a ^ b, *keys)
+        elif casematch(op_name, b'not'):
+            if len(keys) != 1:
+                raise SimpleError(msgs.BITOP_NOT_ONE_KEY_ONLY)
+            val = keys[0].value
+            res = bytes([((1 << 8) - 1 - val[i]) for i in range(len(val))])
+        else:
+            raise SimpleError(msgs.WRONG_ARGS_MSG6.format('bitop'))
+        dst.value = res
+        return len(dst.value)

test/test_mixins/test_bitmap_commands.py

@@ -129,8 +129,33 @@ def test_bitcount_wrong_type(r):
     with pytest.raises(redis.ResponseError):
         r.bitcount('foo')
 
-# def test_bitop(r):
-#     r.set('key1', 'foobar')
-#     r.set('key2', 'abcdef')
-#     assert r.bitop('and', 'dest', 'key1', 'key2') == 6
-#     assert r.get('dest') == b'`bc`ab'
+
+def test_bitop(r):
+    r.set('key1', 'foobar')
+    r.set('key2', 'abcdef')
+
+    assert r.bitop('and', 'dest', 'key1', 'key2') == 6
+    assert r.get('dest') == b'`bc`ab'
+
+    assert r.bitop('not', 'dest1', 'key1') == 6
+    assert r.get('dest1') == b'\x99\x90\x90\x9d\x9e\x8d'
+
+    assert r.bitop('or', 'dest-or', 'key1', 'key2') == 6
+    assert r.get('dest-or') == b'goofev'
+
+    assert r.bitop('xor', 'dest-xor', 'key1', 'key2') == 6
+    assert r.get('dest-xor') == b'\x07\r\x0c\x06\x04\x14'
+
+
+def test_bitop_errors(r):
+    r.set('key1', 'foobar')
+    r.set('key2', 'abcdef')
+    r.sadd('key-set', 'member1')
+    with pytest.raises(redis.ResponseError):
+        r.bitop('not', 'dest', 'key1', 'key2')
+    with pytest.raises(redis.ResponseError):
+        r.bitop('badop', 'dest', 'key1', 'key2')
+    with pytest.raises(redis.ResponseError):
+        r.bitop('and', 'dest', 'key1', 'key-set')
+    with pytest.raises(redis.ResponseError):
+        r.bitop('and', 'dest')