Rewrite str_increment() as bytes_increment()

...and make it work on both Python 3 and 2.

3 files changed, 24 insertions, 20 deletions

diff --git a/happybase/table.py b/happybase/table.py
index 63425f6..8866109 100644
--- a/happybase/table.py
+++ b/happybase/table.py
@@ -10,7 +10,7 @@ from six import iteritems
 
 from .Hbase_thrift import TScan
 
-from .util import thrift_type_to_dict, str_increment, OrderedDict
+from .util import thrift_type_to_dict, bytes_increment, OrderedDict
 from .batch import Batch
 
 logger = logging.getLogger(__name__)
@@ -317,7 +317,7 @@ class Table(object):
                     "or 'row_stop'")
 
             row_start = row_prefix
-            row_stop = str_increment(row_prefix)
+            row_stop = bytes_increment(row_prefix)
 
         if row_start is None:
             row_start = ''
diff --git a/happybase/util.py b/happybase/util.py
index c697f90..f975464 100644
--- a/happybase/util.py
+++ b/happybase/util.py
@@ -69,7 +69,7 @@ def ensure_bytes(str_or_bytes, binary_type=six.binary_type,
         .format(type(str_or_bytes).__name__))
 
 
-def str_increment(s):
+def bytes_increment(b):
     """Increment and truncate a byte string (for sorting purposes)
 
     This functions returns the shortest string that sorts after the given
@@ -79,8 +79,10 @@ def str_increment(s):
     drops everything after it. If the string only contains ``0xFF`` bytes,
     `None` is returned.
     """
-    for i in range(len(s) - 1, -1, -1):
-        if s[i] != '\xff':
-            return s[:i] + chr(ord(s[i]) + 1)
-
+    assert isinstance(b, six.binary_type)
+    b = bytearray(b)  # Used subset of its API is the same on Python 2 and 3.
+    for i in range(len(b) - 1, -1, -1):
+        if b[i] != 0xff:
+            b[i] += 1
+            return bytes(b[:i+1])
     return None
diff --git a/tests/test_util.py b/tests/test_util.py
index 7eb63d5..ab4a1af 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -2,6 +2,8 @@
 HappyBase utility tests.
 """
 
+from codecs import decode, encode
+
 from nose.tools import assert_equal, assert_less
 
 import happybase.util as util
@@ -29,26 +31,26 @@ def test_camel_case_to_pep8():
         yield check, a, b, c
 
 
-def test_str_increment():
+def test_bytes_increment():
     def check(s_hex, expected):
-        s = s_hex.decode('hex')
-        v = util.str_increment(s)
-        v_hex = v.encode('hex')
+        s = decode(s_hex, 'hex')
+        v = util.bytes_increment(s)
+        v_hex = encode(v, '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'),
+        (b'00', b'01'),
+        (b'01', b'02'),
+        (b'fe', b'ff'),
+        (b'1234', b'1235'),
+        (b'12fe', b'12ff'),
+        (b'12ff', b'13'),
+        (b'424242ff', b'424243'),
+        (b'4242ffff', b'4243'),
     ]
 
-    assert util.str_increment('\xff\xff\xff') is None
+    assert util.bytes_increment(b'\xff\xff\xff') is None
 
     for s, expected in test_values:
         yield check, s, expected