Implement BITFIELD

docs/about/changelog.md

@@ -5,6 +5,12 @@ description: Change log of all fakeredis releases
 
 ## Next release
 
+## v2.20.0
+
+### 🚀 Features
+
+- Implement BITFIELD command
+
 ## v2.19.0
 
 ### 🚀 Features

docs/redis-commands/Redis.md

@@ -619,12 +619,16 @@ Closes the connection.
 Resets the connection.
 
 
-## `bitmap` commands (5/7 implemented)
+## `bitmap` commands (6/7 implemented)
 
 ### [BITCOUNT](https://redis.io/commands/bitcount/)
 
 Counts the number of set bits (population counting) in a string.
 
+#### [BITFIELD](https://redis.io/commands/bitfield/)
+
+Performs arbitrary bitfield integer operations on strings.
+
 ### [BITOP](https://redis.io/commands/bitop/)
 
 Performs bitwise operations on multiple strings, and stores the result.
@@ -643,11 +647,7 @@ Sets or clears the bit at offset of the string value. Creates the key if it does
 
 
 ### Unsupported bitmap commands 
-> To implement support for a command, see [here](../../guides/implement-command/) 
-
-#### [BITFIELD](https://redis.io/commands/bitfield/) <small>(not implemented)</small>
-
-Performs arbitrary bitfield integer operations on strings.
+> To implement support for a command, see [here](../../guides/implement-command/)
 
 #### [BITFIELD_RO](https://redis.io/commands/bitfield_ro/) <small>(not implemented)</small>
 

fakeredis/_msgs.py

@@ -97,3 +97,8 @@
 NONSCALING_FILTERS_CANNOT_EXPAND_MSG = "Nonscaling filters cannot expand"
 ITEM_EXISTS_MSG = "item exists"
 NOT_FOUND_MSG = "not found"
+INVALID_BITFIELD_TYPE = (
+    "ERR Invalid bitfield type. Use something like i16 u8. "
+    "Note that u64 is not supported but i64 is."
+)
+INVALID_OVERFLOW_TYPE = "ERR Invalid OVERFLOW type specified"

fakeredis/commands_mixins/bitmap_mixin.py

@@ -1,5 +1,6 @@
 from typing import Tuple
 
+import re
 from fakeredis import _msgs as msgs
 from fakeredis._commands import (
     command,
@@ -13,6 +14,22 @@
 from fakeredis._helpers import SimpleError, casematch
 
 
+class BitfieldEncoding:
+    signed: bool
+    size: int
+
+    def __init__(self, encoding):
+        match = re.match(br'^([ui])(\d+)$', encoding)
+        if match is None:
+            raise SimpleError(msgs.INVALID_BITFIELD_TYPE)
+
+        self.signed = match[1] == b'i'
+        self.size = int(match[2])
+
+        if self.size < 1 or self.size > (64 if self.signed else 63):
+            raise SimpleError(msgs.INVALID_BITFIELD_TYPE)
+
+
 class BitmapCommandsMixin:
     version: Tuple[int]
 
@@ -149,3 +166,85 @@ def bitop(self, op_name, dst, *keys):
             raise SimpleError(msgs.WRONG_ARGS_MSG6.format("bitop"))
         dst.value = res
         return len(dst.value)
+
+    def _bitfield_get(self, key, encoding, offset):
+        ans = 0
+        for i in range(0, encoding.size):
+            ans <<= 1
+            if self.getbit(key, offset + i):
+                ans += -1 if encoding.signed and i == 0 else 1
+        return ans
+
+    def _bitfield_set(self, key, encoding, offset, overflow, value=None, incr=0):
+        if encoding.signed:
+            min_value = -(1 << (encoding.size - 1))
+            max_value = (1 << (encoding.size - 1)) - 1
+        else:
+            min_value = 0
+            max_value = (1 << encoding.size) - 1
+
+        ans = self._bitfield_get(key, encoding, offset)
+        new_value = ans if value is None else value
+        if not encoding.signed:
+            new_value &= (1 << 64) - 1  # force cast to uint64_t
+
+        if overflow == b"FAIL" and not (min_value <= new_value + incr <= max_value):
+            return None  # yes, failing in this context is not writing the value
+        elif overflow == b"SAT":
+            if new_value + incr > max_value:
+                new_value, incr = max_value, 0
+            # REDIS only checks for unsigned underflow on negative incr:
+            if (encoding.signed or incr < 0) and new_value + incr < min_value:
+                new_value, incr = min_value, 0
+
+        new_value += incr
+        new_value &= (1 << encoding.size) - 1
+        # normalize signed number by changing the sign associated to higher bit:
+        if encoding.signed and new_value > max_value:
+            new_value -= 1 << encoding.size
+
+        for i in range(0, encoding.size):
+            bit = (new_value >> (encoding.size - i - 1)) & 1
+            self.setbit(key, offset + i, bit)
+        return new_value if value is None else ans
+
+    @command(fixed=(Key(bytes),), repeat=(bytes,))
+    def bitfield(self, key, *args):
+        overflow = b"WRAP"
+        results = []
+        i = 0
+        while i < len(args):
+            if casematch(args[i], b"overflow") and i + 1 < len(args):
+                overflow = args[i+1].upper()
+                if overflow not in (b"WRAP", b"SAT", b"FAIL"):
+                    raise SimpleError(msgs.INVALID_OVERFLOW_TYPE)
+                i += 2
+            elif casematch(args[i], b"get") and i + 2 < len(args):
+                encoding = BitfieldEncoding(args[i+1])
+                offset = BitOffset.decode(args[i+2])
+                results.append(self._bitfield_get(key, encoding, offset))
+                i += 3
+            elif casematch(args[i], b"set") and i + 3 < len(args):
+                old_value = self._bitfield_set(
+                    key=key,
+                    encoding=BitfieldEncoding(args[i + 1]),
+                    offset=BitOffset.decode(args[i + 2]),
+                    value=Int.decode(args[i + 3]),
+                    overflow=overflow
+                )
+                results.append(old_value)
+                i += 4
+            elif casematch(args[i], b"incrby") and i + 3 < len(args):
+                old_value = self._bitfield_set(
+                    key=key,
+                    encoding=BitfieldEncoding(args[i + 1]),
+                    offset=BitOffset.decode(args[i + 2]),
+                    incr=Int.decode(args[i + 3]),
+                    overflow=overflow
+                )
+                results.append(old_value)
+                i += 4
+            else:
+                raise SimpleError(msgs.SYNTAX_ERROR_MSG)
+
+        return results

test/test_mixins/test_bitmap_commands.py

@@ -208,3 +208,182 @@ def test_bitpos_wrong_arguments(r: redis.Redis):
         raw_command(r, 'bitpos', key, 1, '6', '5', 'BYTE', '6')
     with pytest.raises(redis.ResponseError):
         raw_command(r, 'bitpos', key)
+
+
+def test_bitfield_empty(r: redis.Redis):
+    key = "key:bitfield"
+    assert r.bitfield(key).execute() == []
+    for overflow in ('wrap', 'sat', 'fail'):
+        assert raw_command(r, 'bitfield', key, 'overflow', overflow) == []
+
+
+def test_bitfield_wrong_arguments(r: redis.Redis):
+    key = "key:bitfield:wrong:args"
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield')
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield', key, 'foo')
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield', key, 'overflow')
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield', key, 'overflow', 'foo')
+
+
+def test_bitfield_get(r: redis.Redis):
+    key = "key:bitfield_get"
+    r.set(key, b"\xff\xf0\x00")
+    for i in range(0, 12):
+        assert r.bitfield(key).get('u1', i).get('i1', i).execute() == [1, -1]
+    for i in range(12, 25):
+        for j in range(1, 63):
+            assert r.bitfield(key).get(f'u{j}', i).get(f'i{j}', i).execute() == [0, 0]
+
+    for i in range(0, 11):
+        assert r.bitfield(key).get('u2', i).get('i2', i).execute() == [3, -1]
+    assert r.bitfield(key).get('u2', 11).get('i2', 11).execute() == [2, -2]
+    assert r.bitfield(key).get('u8', 0).get('u8', 8).get('u8', 16).execute() == [0xff, 0xf0, 0]
+    assert r.bitfield(key).get('i8', 0).get('i8', 8).get('i8', 16).execute() == [~0, ~0x0f, 0]
+
+    assert r.bitfield(key).get('u32', 8).get('u8', 100).execute() == [0xf000_0000, 0]
+
+    r.set(key, b"\x01\x23\x45\x67\x89\xab\xcd\xef")
+    for enc in ('i16', 'u16'):
+        assert r.bitfield(key).get(enc, 0).execute() == [0x0123]
+        assert r.bitfield(key).get(enc, 4).execute() == [0x1234]
+        assert r.bitfield(key).get(enc, 8).execute() == [0x2345]
+
+        assert r.bitfield(key).get(enc, 1).execute() == [0x0246]
+        assert r.bitfield(key).get(enc, 5).execute() == [0x2468]
+        assert r.bitfield(key).get(enc, 9).execute() == [0x468a]
+
+        assert r.bitfield(key).get(enc, 2).execute() == [0x048d]
+        assert r.bitfield(key).get(enc, 6).execute() == [0x48d1]
+
+    assert r.bitfield(key).get('u16', 10).get('i16', 10).execute() == [0x8d15, 0xd15 - 0x8000]
+    assert r.bitfield(key).get('u32', 16).get('u48', 8).execute() == [0x456789ab, 0x2345_6789_abcd]
+    assert r.bitfield(key).get('i32', 16).get('i48', 8).execute() == [0x456789ab, 0x2345_6789_abcd]
+    assert r.bitfield(key).get('u63', 1).execute() == [0x123456789_abcdef]
+    assert r.bitfield(key).get('i63', 1).execute() == [0x123456789_abcdef]
+    assert r.bitfield(key).get('i64', 0).execute() == [0x123456789_abcdef]
+    assert raw_command(r, 'bitfield', key, 'get', 'i16', 0) == [0x0123]
+
+
+def test_bitfield_set(r: redis.Redis):
+    key = "key:bitfield_set"
+    r.set(key, b"\xff\xf0\x00")
+    assert r.bitfield(key).set('u8', 0, 0x55).set('u8', 16, 0xaa).execute() == [0xff, 0]
+    assert r.get(key) == b"\x55\xf0\xaa"
+    assert r.bitfield(key).set('u1', 0, 1).set('u1', 16, 2).execute() == [0, 1]
+    assert r.get(key) == b"\xd5\xf0\x2a"
+    assert r.bitfield(key).set('i1', 31, 1).set('i1', 30, 1).execute() == [0, 0]
+    assert r.get(key) == b"\xd5\xf0\x2a\x03"
+    assert r.bitfield(key).set('u36', 4, 0xbadc0ffe).execute() == [0x5_f02a_0300]
+    assert r.get(key) == b"\xd0\xba\xdc\x0f\xfe"
+    assert r.bitfield(key, 'WRAP').set('u12', 8, 0xfff).execute() == [0xbad]
+    assert r.get(key) == b"\xd0\xff\xfc\x0f\xfe"
+
+
+def test_bitfield_set_sat(r: redis.Redis):
+    key = "key:bitfield_set"
+    r.set(key, b"\xff\xf0\x00")
+    assert r.bitfield(key, 'SAT').set('u8', 4, 0x123).set('u8', 8, 0x55).execute() == [0xff, 0xf0]
+    assert r.get(key) == b"\xff\x55\x00"
+    assert r.bitfield(key, 'SAT').set('u12', 0, -1).set('u1', 1, 2).execute() == [0xff5, 1]
+    assert r.get(key) == b"\xff\xf5\x00"
+    assert r.bitfield(key, 'SAT').set('i4', 0, 8).set('i4', 4, 7).execute() == [-1, -1]
+    assert r.get(key) == b"\x77\xf5\x00"
+    assert r.bitfield(key, 'SAT').set('i4', 4, -8).set('i4', 0, -9).execute() == [7, 7]
+    assert r.get(key) == b"\x88\xf5\x00"
+    assert r.bitfield(key, 'SAT').set('i60', 0, -(1 << 62)+1).execute() == [0x88f5000_00000000-(1 << 60)]
+    assert r.get(key) == b"\x80" + b"\0" * 7
+    assert r.bitfield(key, 'SAT').set('u60', 0, -(1 << 63)+1).execute() == [1 << 59]
+    assert r.get(key) == b"\xff" * 7 + b"\xf0"
+
+
+def test_bitfield_set_fail(r: redis.Redis):
+    key = "key:bitfield_set"
+    r.set(key, b"\xff\xf0\x00")
+    assert r.bitfield(key, 'FAIL').set('u8', 4, 0x123).set('u8', 8, 0x55).execute() == [None, 0xf0]
+    assert r.get(key) == b"\xff\x55\x00"
+    assert r.bitfield(key, 'FAIL').set('u12', 0, -1).set('u1', 1, 2).execute() == [None, None]
+    assert r.get(key) == b"\xff\x55\x00"
+    assert r.bitfield(key, 'FAIL').set('i4', 0, 8).set('i4', 4, 7).execute() == [None, -1]
+    assert r.get(key) == b"\xf7\x55\x00"
+    assert r.bitfield(key, 'FAIL').set('i4', 4, -8).set('i4', 0, -9).execute() == [7, None]
+    assert r.get(key) == b"\xf8\x55\x00"
+
+
+def test_bitfield_incr(r: redis.Redis):
+    key = "key:bitfield_incr"
+    r.set(key, b"\xff\xf0\x00")
+    assert r.bitfield(key).incrby('u8', 0, 0x55).incrby('u8', 16, 0xaa).execute() == [0x54, 0xaa]
+    assert r.get(key) == b"\x54\xf0\xaa"
+    assert r.bitfield(key).incrby('u1', 0, 1).incrby('u1', 16, 2).execute() == [1, 1]
+    assert r.get(key) == b"\xd4\xf0\xaa"
+    assert r.bitfield(key).incrby('i1', 31, 1).incrby('i1', 30, 1).execute() == [-1, -1]
+    assert r.get(key) == b"\xd4\xf0\xaa\x03"
+    assert r.bitfield(key).incrby('u36', 4, 0xbadc0ffe).execute() == [0x5_ab86_12fe]
+    assert r.get(key) == b"\xd5\xab\x86\x12\xfe"
+    assert r.bitfield(key, 'WRAP').incrby('u12', 8, 0xfff).execute() == [0xab7]
+    assert r.get(key) == b"\xd5\xab\x76\x12\xfe"
+
+
+def test_bitfield_incr_sat(r: redis.Redis):
+    key = "key:bitfield_incr_sat"
+    r.set(key, b"\xff\xf0\x00")
+    assert r.bitfield(key, 'SAT').incrby('u8', 4, 0x123).incrby('u8', 8, 0x55).execute() == [0xff, 0xff]
+    assert r.get(key) == b"\xff\xff\x00"
+    assert r.bitfield(key, 'SAT').incrby('u12', 0, -1).incrby('u1', 1, 2).execute() == [0xffe, 1]
+    assert r.get(key) == b"\xff\xef\x00"
+    assert r.bitfield(key, 'SAT').incrby('i4', 0, 8).incrby('i4', 4, 7).execute() == [7, 6]
+    assert r.get(key) == b"\x76\xef\x00"
+    assert r.bitfield(key, 'SAT').incrby('i4', 4, -8).incrby('i4', 0, -9).execute() == [-2, -2]
+    assert r.get(key) == b"\xee\xef\x00"
+    assert r.bitfield(key, 'SAT').incrby('i60', 0, -(1 << 62)+1).execute() == [-(1 << 59)]
+    assert r.get(key) == b"\x80" + b"\0" * 7
+    assert r.bitfield(key, 'SAT').set('u60', 0, -(1 << 63)+1).execute() == [1 << 59]
+    assert r.get(key) == b"\xff" * 7 + b"\xf0"
+
+
+def test_bitfield_incr_fail(r: redis.Redis):
+    key = "key:bitfield_incr_fail"
+    r.set(key, b"\xff\xf0\x00")
+    assert r.bitfield(key, 'FAIL').incrby('u8', 4, 0x123).incrby('u8', 8, 0x55).execute() == [None, None]
+    assert r.get(key) == b"\xff\xf0\x00"
+    assert r.bitfield(key, 'FAIL').incrby('u12', 0, -1).incrby('u1', 1, 2).execute() == [0xffe, None]
+    assert r.get(key) == b"\xff\xe0\x00"
+    assert r.bitfield(key, 'FAIL').incrby('i4', 0, 8).incrby('i4', 4, 7).execute() == [7, 6]
+    assert r.get(key) == b"\x76\xe0\x00"
+    assert r.bitfield(key, 'FAIL').incrby('i4', 4, -8).incrby('i4', 0, -9).execute() == [-2, -2]
+    assert r.get(key) == b"\xee\xe0\x00"
+
+
+def test_bitfield_get_wrong_arguments(r: redis.Redis):
+    key = "key:bitfield_get:wrong:args"
+    r.set(key, b"\xff\xf0\x00")
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield', key, 'get')
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield', key, 'get', 'i16')
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield', key, 'get', 'i16', -1)
+    for encoding in ('I8', 'i-42', 'i5?', 'u0', 'i0', 'i65', 'u64', 'i 60'):
+        with pytest.raises(redis.ResponseError):
+            raw_command(r, 'bitfield', key, 'get', encoding, 0)
+
+
+def test_bitfield_set_wrong_arguments(r: redis.Redis):
+    key = "key:bitfield_set:wrong:args"
+    r.set(key, b"\xff\xf0\x00")
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield', key, 'set')
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield', key, 'set', 'i16')
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield', key, 'set', 'i16', -1)
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'bitfield', key, 'set', 'i16', 0, 'foo')
+    for encoding in ('I8', 'i-42', 'i5?', 'u0', 'i0', 'i65', 'u64', 'i 60'):
+        with pytest.raises(redis.ResponseError):
+            raw_command(r, 'bitfield', key, 'set', encoding, 0, 0)
+