added better overflow control and lifecycle mgmt to bitfields

- bitfields now accept a `default_overflow` argument that sets the
  default overflow behavior for incrby operations
- exposed an overflow() method that sets the overflow behavior for future
  incrby operations. this can be used in place of the overfly argument
  to incrby if someone prefers
- clean up bitfield instance upon execution, resetting the overflow behavior
  back to the default provided and clearing the list of operations

2 files changed, 79 insertions, 32 deletions

diff --git a/redis/client.py b/redis/client.py
index 0692b39..aaa7433 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -1005,12 +1005,12 @@ class StrictRedis(object):
             raise RedisError("Both start and end must be specified")
         return self.execute_command('BITCOUNT', *params)
 
-    def bitfield(self, key):
+    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)
+        return BitFieldOperation(self, key, default_overflow=default_overflow)
 
     def bitop(self, operation, dest, *keys):
         """
@@ -3461,29 +3461,49 @@ class BitFieldOperation(object):
     """
     Command builder for BITFIELD commands.
     """
-    def __init__(self, client, key):
+    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 = None  # Default is "WRAP".
+        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 int offset: offset (in number of bits). If prefixed with a
+        :param offset: offset (in number of bits). If prefixed with a
             '#', this is an offset multiplier, e.g. given the arguments
-            fmt='i8', offset='#2', the offset will be 16.
+            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 and overflow != self._last_overflow:
-            self._last_overflow = overflow
-            self.operations.append(('OVERFLOW', overflow))
+        if overflow is not None:
+            self.overflow(overflow)
 
         self.operations.append(('INCRBY', fmt, offset, increment))
         return self
@@ -3493,9 +3513,9 @@ class BitFieldOperation(object):
         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 int offset: offset (in number of bits). If prefixed with a
+        :param offset: offset (in number of bits). If prefixed with a
             '#', this is an offset multiplier, e.g. given the arguments
-            fmt='i8', offset='#2', the offset will be 16.
+            fmt='u8', offset='#2', the offset will be 16.
         :returns: a :py:class:`BitFieldOperation` instance.
         """
         self.operations.append(('GET', fmt, offset))
@@ -3506,9 +3526,9 @@ class BitFieldOperation(object):
         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 int offset: offset (in number of bits). If prefixed with a
+        :param offset: offset (in number of bits). If prefixed with a
             '#', this is an offset multiplier, e.g. given the arguments
-            fmt='i8', offset='#2', the offset will be 16.
+            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.
         """
@@ -3530,5 +3550,5 @@ class BitFieldOperation(object):
         will be present within the pipeline's execute.
         """
         command = self.command
-        self.operations = []
+        self.reset()
         return self.client.execute_command(*command)
diff --git a/tests/test_commands.py b/tests/test_commands.py
index 4667708..00472fd 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -2116,39 +2116,66 @@ class TestRedisCommands(object):
         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)
-                .get('u8', 0)
-                .get('u4', 8)  # 1111
-                .get('u4', 12)  # 1111
-                .get('u4', 13)  # 1110
+                .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)  # 00ff -> 001f (returns old val, 0x0f)
-                .get('u16', 0)  # 001f (00011111)
-                .set('u16', 0, 0)  # 001f -> 0000
+                .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)
-                .get('u16', 0)
+                .incrby('u8', 8, 254)  # 00000000 11111110
+                .incrby('u8', 8, 1)    # 00000000 11111111
+                .get('u16', 0)         # 00000000 11111111
                 .execute())
-        assert resp == [254, 254]
+        assert resp == [254, 255, 255]
 
-        # Verify overflow protection works:
+        # Verify overflow protection works as a method:
+        r.delete('a')
         resp = (bf
-                .incrby('u8', 8, 2, 'FAIL')
-                .incrby('u8', 8, 1)
-                .incrby('u8', 8, 1)  # Still "FAIL".
-                .get('u16', 0)
+                .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 == [None, 255, None, 255]
+        assert resp == [0, None, 255, None, 255]
 
-        assert r.get('a') == b'\x00\xff'
+        # 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):