Implement a str_increment() utility function

happybase/util.py

@@ -1,5 +1,7 @@
 """
 HappyBase utility module.
+
+These functions are not part of the public API.
 """
 
 import re
@@ -37,3 +39,19 @@ def thrift_type_to_dict(obj):
     """Converts a Thrift data type to a regular dictionary."""
     return dict((camel_case_to_pep8(attr), getattr(obj, attr))
                 for attr in thrift_attrs(obj))
+
+
+def str_increment(s):
+    """
+    Returns the shortest string that sorts after the given string when compared
+    using regular string comparison semantics.
+
+    This function increments the last byte that is smaller than ``0xFF``, and
+    drops everything after it. If the string only contains ``0xFF`` bytes,
+    `None` is returned.
+    """
+    for i in xrange(len(s) - 1, -1, -1):
+        if s[i] != '\xff':
+            return s[:i] + chr(ord(s[i]) + 1)
+
+    return None

tests/test_util.py

@@ -2,21 +2,21 @@
 HappyBase utility tests.
 """
 
-from nose.tools import assert_equal
+from nose.tools import assert_equal, assert_less
 
-import happybase.util
+import happybase.util as util
 
 
 def test_camel_case_to_pep8():
     def check(lower_cc, upper_cc, correct):
 
-        x1 = happybase.util.camel_case_to_pep8(lower_cc)
-        x2 = happybase.util.camel_case_to_pep8(upper_cc)
+        x1 = util.camel_case_to_pep8(lower_cc)
+        x2 = util.camel_case_to_pep8(upper_cc)
         assert_equal(correct, x1)
         assert_equal(correct, x2)
 
-        y1 = happybase.util.pep8_to_camel_case(x1, True)
-        y2 = happybase.util.pep8_to_camel_case(x2, False)
+        y1 = util.pep8_to_camel_case(x1, True)
+        y2 = util.pep8_to_camel_case(x2, False)
         assert_equal(upper_cc, y1)
         assert_equal(lower_cc, y2)
 
@@ -27,3 +27,28 @@ def test_camel_case_to_pep8():
 
     for a, b, c in examples:
         yield check, a, b, c
+
+
+def test_str_increment():
+    def check(s_hex, expected):
+        s = s_hex.decode('hex')
+        v = util.str_increment(s)
+        v_hex = v.encode('hex')
+        assert_equal(expected, v_hex)
+        assert_less(s, v)
+
+    test_values = [
+        ('00', '01'),
+        ('01', '02'),
+        ('fe', 'ff'),
+        ('1234', '1235'),
+        ('12fe', '12ff'),
+        ('12ff', '13'),
+        ('424242ff', '424243'),
+        ('4242ffff', '4243'),
+    ]
+
+    assert util.str_increment('\xff\xff\xff') is None
+
+    for s, expected in test_values:
+        yield check, s, expected