Repository: JelleZijlstra/ast_decompiler

Issue:
Docstrings are not preserved correctly.
Currently it's impossible to correctly preserve docstrings. This test shows the issue:
```
from .tests import assert_decompiles


def test_docstrings():
    assert_decompiles('''
def a():
    """
    Docstring.
    """
    return 1
''', '''def a():
 """
 Docstring.
 """
 return 1
''', indentation=1)
```

Result:
```
>>> expected
def a():
 """
 Docstring.
 """
 return 1

>>> actual
def a():
 '\n    Docstring.\n    '
 return 1
```

Thoughts: when processing the function body in `write_function_def` we should look at the very first node in the function's body. If the node is of type Expr, and contains a string, then it's most likely docstring of this function and should be processed accordingly.

Discussion:

Patch:
diff --git a/CHANGELOG b/CHANGELOG
index 62b8063..3d37564 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,4 +1,5 @@
 unreleased
+- Preserve literal newlines in docstrings
 - Fix decompilation of complex infinity
 - Add support for Python 3.10 pattern matching
 - Fix incorrect decompilation of lambdas in comprehension guards on Python 3.9 and higher
diff --git a/ast_decompiler/decompiler.py b/ast_decompiler/decompiler.py
index d21c412..e642df8 100644
--- a/ast_decompiler/decompiler.py
+++ b/ast_decompiler/decompiler.py
@@ -834,13 +834,21 @@ class Decompiler(ast.NodeVisitor):
     def write_string(self, string_value: str, kind: Optional[str] = None) -> None:
         if kind is not None:
             self.write(kind)
+        if isinstance(self.get_parent_node(), ast.Expr) and '"""' not in string_value:
+            self.write('"""')
+            s = string_value.encode("unicode-escape").decode("ascii")
+            s = s.replace("\\n", "\n").replace("\\r", "\r")
+            self.write(s)
+            self.write('"""')
+            return
         if self.has_parent_of_type(ast.FormattedValue):
             delimiter = '"'
         else:
             delimiter = "'"
         self.write(delimiter)
         s = string_value.encode("unicode-escape").decode("ascii")
-        self.write(s.replace(delimiter, "\\" + delimiter))
+        s = s.replace(delimiter, "\\" + delimiter)
+        self.write(s)
         self.write(delimiter)
 
     def visit_FormattedValue(self, node: ast.FormattedValue) -> None:


Test patch:
diff --git a/tests/test_basic.py b/tests/test_basic.py
index 5b39dbc..72a8ecb 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -1,6 +1,6 @@
 import ast
 from ast_decompiler import decompile
-from .tests import check, only_on_version
+from .tests import assert_decompiles, check, only_on_version
 
 
 def test_non_module() -> None:
@@ -380,6 +380,24 @@ def test_Str() -> None:
 b'foo'
 """
     )
+    check('"a\\nb"')
+    assert_decompiles(
+        '''def f():
+    """Doc.
+
+    String.
+
+    """
+''',
+        '''
+def f():
+    """Doc.
+
+    String.
+
+    """
+''',
+    )
 
 
 def test_Attribute() -> None: