Merge pull request #1056 from andymccurdy/bitfield

Bitfields

3 files changed, 183 insertions, 0 deletions

diff --git a/redis/client.py b/redis/client.py
index 8173ba8..aaa7433 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -1005,6 +1005,13 @@ class StrictRedis(object):
             raise RedisError("Both start and end must be specified")
         return self.execute_command('BITCOUNT', *params)
 
+    def bitfield(self, key, default_overflow=None):
+        """
+        Return a BitFieldOperation instance to conveniently construct one or
+        more bitfield operations on ``key``.
+        """
+        return BitFieldOperation(self, key, default_overflow=default_overflow)
+
     def bitop(self, operation, dest, *keys):
         """
         Perform a bitwise operation using ``operation`` between ``keys`` and
@@ -3448,3 +3455,100 @@ class Script(object):
             # Overwrite the sha just in case there was a discrepancy.
             self.sha = client.script_load(self.script)
             return client.evalsha(self.sha, len(keys), *args)
+
+
+class BitFieldOperation(object):
+    """
+    Command builder for BITFIELD commands.
+    """
+    def __init__(self, client, key, default_overflow=None):
+        self.client = client
+        self.key = key
+        self._default_overflow = default_overflow
+        self.reset()
+
+    def reset(self):
+        """
+        Reset the state of the instance to when it was constructed
+        """
+        self.operations = []
+        self._last_overflow = 'WRAP'
+        self.overflow(self._default_overflow or self._last_overflow)
+
+    def overflow(self, overflow):
+        """
+        Update the overflow algorithm of successive INCRBY operations
+        :param overflow: Overflow algorithm, one of WRAP, SAT, FAIL. See the
+            Redis docs for descriptions of these algorithmsself.
+        :returns: a :py:class:`BitFieldOperation` instance.
+        """
+        overflow = overflow.upper()
+        if overflow != self._last_overflow:
+            self._last_overflow = overflow
+            self.operations.append(('OVERFLOW', overflow))
+        return self
+
+    def incrby(self, fmt, offset, increment, overflow=None):
+        """
+        Increment a bitfield by a given amount.
+        :param fmt: format-string for the bitfield being updated, e.g. 'u8'
+            for an unsigned 8-bit integer.
+        :param offset: offset (in number of bits). If prefixed with a
+            '#', this is an offset multiplier, e.g. given the arguments
+            fmt='u8', offset='#2', the offset will be 16.
+        :param int increment: value to increment the bitfield by.
+        :param str overflow: overflow algorithm. Defaults to WRAP, but other
+            acceptable values are SAT and FAIL. See the Redis docs for
+            descriptions of these algorithms.
+        :returns: a :py:class:`BitFieldOperation` instance.
+        """
+        if overflow is not None:
+            self.overflow(overflow)
+
+        self.operations.append(('INCRBY', fmt, offset, increment))
+        return self
+
+    def get(self, fmt, offset):
+        """
+        Get the value of a given bitfield.
+        :param fmt: format-string for the bitfield being read, e.g. 'u8' for
+            an unsigned 8-bit integer.
+        :param offset: offset (in number of bits). If prefixed with a
+            '#', this is an offset multiplier, e.g. given the arguments
+            fmt='u8', offset='#2', the offset will be 16.
+        :returns: a :py:class:`BitFieldOperation` instance.
+        """
+        self.operations.append(('GET', fmt, offset))
+        return self
+
+    def set(self, fmt, offset, value):
+        """
+        Set the value of a given bitfield.
+        :param fmt: format-string for the bitfield being read, e.g. 'u8' for
+            an unsigned 8-bit integer.
+        :param offset: offset (in number of bits). If prefixed with a
+            '#', this is an offset multiplier, e.g. given the arguments
+            fmt='u8', offset='#2', the offset will be 16.
+        :param int value: value to set at the given position.
+        :returns: a :py:class:`BitFieldOperation` instance.
+        """
+        self.operations.append(('SET', fmt, offset, value))
+        return self
+
+    @property
+    def command(self):
+        cmd = ['BITFIELD', self.key]
+        for ops in self.operations:
+            cmd.extend(ops)
+        return cmd
+
+    def execute(self):
+        """
+        Execute the operation(s) in a single BITFIELD command. The return value
+        is a list of values corresponding to each operation. If the client
+        used to create this instance was a pipeline, the list of values
+        will be present within the pipeline's execute.
+        """
+        command = self.command
+        self.reset()
+        return self.client.execute_command(*command)
diff --git a/tests/test_commands.py b/tests/test_commands.py
index f0394d7..15e7620 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -2115,6 +2115,68 @@ class TestRedisCommands(object):
         # 1 message is trimmed
         assert r.xtrim(stream, 3, approximate=False) == 1
 
+    def test_bitfield_operations(self, r):
+        # comments show affected bits
+        bf = r.bitfield('a')
+        resp = (bf
+                .set('u8', 8, 255)     # 00000000 11111111
+                .get('u8', 0)          # 00000000
+                .get('u4', 8)                   # 1111
+                .get('u4', 12)                      # 1111
+                .get('u4', 13)                       # 111 0
+                .execute())
+        assert resp == [0, 0, 15, 15, 14]
+
+        # .set() returns the previous value...
+        resp = (bf
+                .set('u8', 4, 1)           # 0000 0001
+                .get('u16', 0)         # 00000000 00011111
+                .set('u16', 0, 0)      # 00000000 00000000
+                .execute())
+        assert resp == [15, 31, 31]
+
+        # incrby adds to the value
+        resp = (bf
+                .incrby('u8', 8, 254)  # 00000000 11111110
+                .incrby('u8', 8, 1)    # 00000000 11111111
+                .get('u16', 0)         # 00000000 11111111
+                .execute())
+        assert resp == [254, 255, 255]
+
+        # Verify overflow protection works as a method:
+        r.delete('a')
+        resp = (bf
+                .set('u8', 8, 254)     # 00000000 11111110
+                .overflow('fail')
+                .incrby('u8', 8, 2)    # incrby 2 would overflow, None returned
+                .incrby('u8', 8, 1)    # 00000000 11111111
+                .incrby('u8', 8, 1)    # incrby 1 would overflow, None returned
+                .get('u16', 0)         # 00000000 11111111
+                .execute())
+        assert resp == [0, None, 255, None, 255]
+
+        # Verify overflow protection works as arg to incrby:
+        r.delete('a')
+        resp = (bf
+                .set('u8', 8, 255)           # 00000000 11111111
+                .incrby('u8', 8, 1)          # 00000000 00000000  wrap default
+                .set('u8', 8, 255)           # 00000000 11111111
+                .incrby('u8', 8, 1, 'FAIL')  # 00000000 11111111  fail
+                .incrby('u8', 8, 1)          # 00000000 11111111  still fail
+                .get('u16', 0)               # 00000000 11111111
+                .execute())
+        assert resp == [0, 0, 0, None, None, 255]
+
+        # test default default_overflow
+        r.delete('a')
+        bf = r.bitfield('a', default_overflow='FAIL')
+        resp = (bf
+                .set('u8', 8, 255)     # 00000000 11111111
+                .incrby('u8', 8, 1)    # 00000000 11111111  fail default
+                .get('u16', 0)         # 00000000 11111111
+                .execute())
+        assert resp == [0, None, 255]
+
 
 class TestStrictCommands(object):
     def test_strict_zadd(self, sr):
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index a8941d7..3d63d91 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -252,3 +252,20 @@ class TestPipeline(object):
             assert unicode(ex.value).startswith(expected)
 
         assert r[key] == b'1'
+
+    def test_pipeline_with_bitfield(self, r):
+        with r.pipeline() as pipe:
+            pipe.set('a', '1')
+            bf = pipe.bitfield('b')
+            pipe2 = (bf
+                     .set('u8', 8, 255)
+                     .get('u8', 0)
+                     .get('u4', 8)  # 1111
+                     .get('u4', 12)  # 1111
+                     .get('u4', 13)  # 1110
+                     .execute())
+            pipe.get('a')
+            response = pipe.execute()
+
+            assert pipe == pipe2
+            assert response == [True, [0, 0, 15, 15, 14], b'1']