Merge pull request #67 from berkerpeksag/unicode_literals

Add support for __future__ unicode_literals

astor/code_gen.py

@@ -125,6 +125,8 @@ class SourceGenerator(ExplicitNodeVisitor):
 
     """
 
+    using_unicode_literals = False
+
     def __init__(self, indent_with, add_line_information=False,
                  pretty_string=pretty_string):
         self.result = []
@@ -264,6 +266,10 @@ class SourceGenerator(ExplicitNodeVisitor):
         self.statement(node, 'from ', node.level * '.',
                        node.module or '', ' import ')
         self.comma_list(node.names)
+        # Goofy stuff for Python 2.7 _pyio module
+        if node.module == '__future__' and 'unicode_literals' in (
+            x.name for x in node.names):
+            self.using_unicode_literals = True
 
     def visit_Import(self, node):
         self.statement(node, 'import ')
@@ -497,7 +503,8 @@ class SourceGenerator(ExplicitNodeVisitor):
         # Cheesy way to force a flush
         self.write('foo')
         result.pop()
-        result.append(self.pretty_string(node.s, embedded, result))
+        result.append(self.pretty_string(node.s, embedded, result,
+                                         uni_lit=self.using_unicode_literals))
 
     def visit_JoinedStr(self, node,
                   # constants

astor/string_repr.py

@@ -76,7 +76,8 @@ def _prep_triple_quotes(s, mysplit=mysplit, replacements=replacements):
     return ''.join(s)
 
 
-def pretty_string(s, embedded, current_output, min_trip_str=20, max_line=100):
+def pretty_string(s, embedded, current_output, min_trip_str=20,
+                  max_line=100, uni_lit=False):
     """There are a lot of reasons why we might not want to or
        be able to return a triple-quoted string.  We can always
        punt back to the default normal string.
@@ -87,6 +88,8 @@ def pretty_string(s, embedded, current_output, min_trip_str=20, max_line=100):
     # Punt on abnormal strings
     if (isinstance(s, special_unicode) or not isinstance(s, basestring)):
         return default
+    if uni_lit and isinstance(s, bytes):
+        return 'b' + default
 
     len_s = len(default)
     current_line = _get_line(current_output)

tests/test_code_gen.py

@@ -417,6 +417,14 @@ class CodegenTestCase(unittest.TestCase):
         ast2 = astor.parsefile(__file__)
         self.assertEqual(astor.to_source(ast1), astor.codegen.to_source(ast2))
 
+    def test_unicode_literals(self):
+        source = """
+        from __future__ import (print_function, unicode_literals)
+        x = b'abc'
+        y = u'abc'
+        """
+        self.assertAstEqual(source)
+
 
 if __name__ == '__main__':
     unittest.main()