diff --git a/autopxd/__init__.py b/autopxd/__init__.py index 81d0351..e59fcfb 100644 --- a/autopxd/__init__.py +++ b/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)