[py2.py3] add function to coerce strings to six.binary_type

autopxd/__init__.py

@@ -43,6 +43,24 @@
 ))
 
 
+# TODO use upstream six.ensure_binary once next version of six is released
+def ensure_binary(s, encoding='utf-8', errors='strict'):
+    """Coerce **s** to six.binary_type.
+    For Python 2:
+      - `unicode` -> encoded to `str`
+      - `str` -> `str`
+    For Python 3:
+      - `str` -> encoded to `bytes`
+      - `bytes` -> `bytes`
+    """
+    if isinstance(s, six.text_type):
+        return s.encode(encoding, errors)
+    elif isinstance(s, six.binary_type):
+        return s
+    else:
+        raise TypeError("not expecting type '%s'" % type(s))
+
+
 class PxdNode(object):
     indent = '    '
 
@@ -311,7 +329,8 @@ def preprocess(code, extra_cpp_args=[]):
         'cpp', '-nostdinc', '-D__attribute__(x)=', '-I', BUILTIN_HEADERS_DIR,
     ] + extra_cpp_args + ['-'], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
     result = []
-    result.append(proc.communicate(input=code.encode('utf-8'))[0])
+    code = ensure_binary(code)
+    result.append(proc.communicate(input=code)[0])
     while proc.poll() is None:
         result.append(proc.communicate()[0])
     if proc.returncode:
@@ -355,4 +374,7 @@ def translate(code, hdrname, extra_cpp_args=[], whitelist=None):
 @click.argument('infile', type=click.File('rb'), default=sys.stdin)
 @click.argument('outfile', type=click.File('wb'), default=sys.stdout)
 def cli(infile, outfile):
-    outfile.write(translate(infile.read(), infile.name))
+    output = translate(infile.read(), infile.name)
+    if outfile is not sys.stdout:
+        output = ensure_binary(output)
+    outfile.write(output)