Add ensure_bytes() helper

...to coerce unicode text into UTF-8 encoded byte strings.

1 files changed, 13 insertions, 0 deletions

diff --git a/happybase/util.py b/happybase/util.py
index 99141f3..c697f90 100644
--- a/happybase/util.py
+++ b/happybase/util.py
@@ -6,6 +6,7 @@ These functions are not part of the public API.
 
 import re
 
+import six
 from six.moves import range
 
 CAPITALS = re.compile('([A-Z])')
@@ -56,6 +57,18 @@ def thrift_type_to_dict(obj):
                 for attr in thrift_attrs(obj))
 
 
+def ensure_bytes(str_or_bytes, binary_type=six.binary_type,
+                 text_type=six.text_type):
+    """Convert text into bytes, and leaves bytes as-is."""
+    if isinstance(str_or_bytes, binary_type):
+        return str_or_bytes
+    if isinstance(str_or_bytes, text_type):
+        return str_or_bytes.encode('utf-8')
+    raise TypeError(
+        "input must be a text or byte string, got {}"
+        .format(type(str_or_bytes).__name__))
+
+
 def str_increment(s):
     """Increment and truncate a byte string (for sorting purposes)