using dispatch to implement force_bytes and force_text

Author: crazyeng

wolframclient/utils/encoding.py

@@ -1,63 +1,87 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 from wolframclient.utils import six
+from wolframclient.utils.dispatch import Dispatch
 from wolframclient.utils.functional import map
 
+force_text = Dispatch()
 
-def force_text(s, encoding="utf-8", errors="strict"):
-    """
-    Similar to smart_text, except that lazy instances are resolved to
-    strings, rather than kept as lazy objects.
-
-    If strings_only is True, don't convert (some) non-string-like objects.
-    """
-    # Handle the common case first for performance reasons.
-    if isinstance(s, six.text_type):
-        return s
-    if not isinstance(s, six.string_types):
-        if six.PY3:
-            if isinstance(s, bytes):
-                s = six.text_type(s, encoding, errors)
-            else:
-                s = six.text_type(s)
-        elif hasattr(s, "__unicode__"):
-            s = six.text_type(s)
+
+@force_text.dispatch(six.text_type)
+def encode(s, encoding="utf-8", errors="strict"):
+    return s
+
+
+@force_text.dispatch(six.binary_type, replace_existing=True)
+def encode(s, encoding="utf-8", errors="strict"):
+    return s.decode(encoding, errors)
+
+
+if not six.PY2:
+
+    @force_text.dispatch(object)
+    def encode(s, encoding="utf-8", errors="strict"):
+        return six.text_type(s)
+
+
+else:
+
+    @force_text.dispatch(object)
+    def encode(s, encoding="utf-8", errors="strict"):
+        if hasattr(s, "__unicode__"):
+            return six.text_type(s)
         else:
-            s = six.text_type(bytes(s), encoding, errors)
-    else:
-        # Note: We use .decode() here, instead of six.text_type(s, encoding,
-        # errors), so that if s is a SafeBytes, it ends up being a
-        # SafeText at the end.
-        s = s.decode(encoding, errors)
+            return six.text_type(bytes(s), encoding, errors)
+
+
+force_bytes = Dispatch()
+
+
+@force_bytes.dispatch(six.string_types)
+def encode(s, encoding="utf-8", errors="strict"):
+    return s.encode(encoding, errors)
+
+
+@force_bytes.dispatch(six.binary_type, replace_existing=True)
+def encode(s, encoding="utf-8", errors="strict"):
     return s
 
 
-def force_bytes(s, encoding="utf-8", errors="strict"):
-    """
-    If strings_only is True, don't convert (some) non-string-like objects.
-    """
-    # Handle the common case first for performance reasons.
-    if isinstance(s, bytes):
-        return s
-
-    if isinstance(s, six.buffer_types):
-        return bytes(s)
-
-    if not isinstance(s, six.string_types):
-        try:
-            if six.PY3:
-                return six.text_type(s).encode(encoding)
-            else:
-                return bytes(s)
-        except UnicodeEncodeError:
-            if isinstance(s, Exception):
-                # An Exception subclass containing non-ASCII data that doesn't
-                # know how to print itself properly. We shouldn't raise a
-                # further exception.
-                return b" ".join(force_bytes(arg, encoding, errors=errors) for arg in s)
-            return six.text_type(s).encode(encoding, errors)
-    else:
-        return s.encode(encoding, errors)
+@force_bytes.dispatch(six.buffer_types, replace_existing=True)
+def encode(s, encoding="utf-8", errors="strict"):
+    return six.binary_type(s)
+
+
+if six.PY2:
+
+    @force_bytes.dispatch(memoryview, replace_existing=True)
+    def encode(s, encoding="utf-8", errors="strict"):
+        return s.tobytes()
+
+
+if not six.PY2:
+
+    def encode_default(s, encoding="utf-8", errors="strict"):
+        return six.text_type(s).encode(encoding)
+
+
+else:
+
+    def encode_default(s, encoding="utf-8", errors="strict"):
+        return six.binary_type(s)
+
+
+@force_bytes.dispatch(object)
+def encode(s, encoding="utf-8", errors="strict"):
+    try:
+        return encode_default(s, encoding, errors)
+    except UnicodeEncodeError:
+        if isinstance(s, Exception):
+            # An Exception subclass containing non-ASCII data that doesn't
+            # know how to print itself properly. We shouldn't raise a
+            # further exception.
+            return b" ".join(force_bytes(arg, encoding, errors=errors) for arg in s)
+        return six.text_type(s).encode(encoding, errors)
 
 
 def safe_force_text(obj):