Add comment fix (#102)

Author: nielsdebruin

rewrite/rewrite/python/_parser_visitor.py

@@ -15,8 +15,8 @@
     NameTree, OmitParentheses, Expression, TypeTree, TypedTree, Comment
 from rewrite.java import tree as j
 from . import tree as py
-from .support_types import PyComment
 from .markers import KeywordArguments, KeywordOnlyArguments, Quoted
+from .support_types import PyComment
 
 T = TypeVar('T')
 J2 = TypeVar('J2', bound=J)
@@ -2167,7 +2167,7 @@ def __format(self, source: str, offset: int, stop: Optional[str] = None) -> Tupl
                 whitespace.append(char)
             elif char == '#':
                 if comments:
-                    comments[-1] = comments[-1].with_suffix('\n' + ''.join(whitespace))
+                    comments[-1] = comments[-1].with_suffix(''.join(whitespace))
                 else:
                     prefix = ''.join(whitespace)
                 whitespace = []

rewrite/tests/python/all/comment_test.py

@@ -8,3 +8,61 @@ def test_comment():
 
 def test_windows_line_endings():
     rewrite_run(python("assert 1 # type: foo\r\n"))
+
+
+def test_multiline_comment():
+    # language=python
+    rewrite_run(python("""
+    '''
+    This is a
+    multiline comment
+    '''
+    assert 1
+    """))
+
+
+def test_multiline_comment_with_code():
+    # language=python
+    rewrite_run(python("""
+    # This is a
+    # multiline comment
+    assert 1
+    # This is another
+    # multiline comment
+    """))
+
+
+def test_multiline_comment_in_class():
+    # language=python
+    rewrite_run(python("""
+    class ExampleClass:
+        '''
+        This is a
+        multiline comment
+        inside a class
+        '''
+        def example_method(self):
+            def example_function():
+                '''
+                This is a
+                multiline comment
+                inside a function
+                '''
+                assert 1
+    """))
+
+
+def test_multiline_comment_with_hash():
+    # language=python
+    rewrite_run(python("""
+    class ExampleClass:
+        # This is a
+        # multiline comment
+        # inside a class
+        def example_method(self):
+            def example_function():
+                # This is a
+                # multiline comment
+                # inside a function
+                assert 1
+    """))