Merge pull request #41 from berkerpeksag/pm_develop

Stop triple quoting strings embedded in function calls

astor/code_gen.py

@@ -479,10 +479,12 @@ class SourceGenerator(ExplicitNodeVisitor):
 
     def visit_Str(self, node):
         result = self.result
+        embedded = self.get__pp(node) > Precedence.Expr
+
         # Cheesy way to force a flush
         self.write('foo')
         result.pop()
-        result.append(self.pretty_string(node.s, result))
+        result.append(self.pretty_string(node.s, embedded, result))
 
     def visit_Bytes(self, node):
         self.write(repr(node.s))

astor/string_repr.py

@@ -76,7 +76,7 @@ def _prep_triple_quotes(s, mysplit=mysplit, replacements=replacements):
     return ''.join(s)
 
 
-def pretty_string(s, current_output, min_trip_str=20, max_line=100):
+def pretty_string(s, embedded, current_output, min_trip_str=20, max_line=100):
     """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.
@@ -91,6 +91,9 @@ def pretty_string(s, current_output, min_trip_str=20, max_line=100):
     len_s = len(default)
     current_line = _get_line(current_output)
     if current_line.strip():
+        if embedded and '\n' not in s:
+            return default
+
         if len_s < min_trip_str:
             return default
 

tests/test_code_gen.py

@@ -323,6 +323,17 @@ class CodegenTestCase(unittest.TestCase):
         """
         self.assertAstEqualIfAtLeastVersion(source, (2, 7))
 
+    def test_output_formatting(self):
+        source = """
+            __all__ = ['ArgumentParser', 'ArgumentError', 'ArgumentTypeError',
+                'FileType', 'HelpFormatter', 'ArgumentDefaultsHelpFormatter',
+                'RawDescriptionHelpFormatter', 'RawTextHelpFormatter', 'Namespace',
+                'Action', 'ONE_OR_MORE', 'OPTIONAL', 'PARSER', 'REMAINDER', 'SUPPRESS',
+                'ZERO_OR_MORE']
+        """
+        self.maxDiff=2000
+        self.assertAstSourceEqual(source)
+
 
 if __name__ == '__main__':
     unittest.main()