Preserve line breaks when removing comments (fixes andialbrecht#484).

andialbrecht · andialbrecht · commit 1499cffcd7c4 · 2020-09-13T13:29:36.000+02:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -14,6 +14,7 @@ Notable Changes
 Bug Fixes
 
 * Improved parsing of IN(...) statements (issue566, pr567 by hurcy).
+* Preserve line breaks when removing comments (issue484).
 
 
 Release 0.3.1 (Feb 29, 2020)
diff --git a/sqlparse/filters/others.py b/sqlparse/filters/others.py
@@ -5,17 +5,29 @@
 # This module is part of python-sqlparse and is released under
 # the BSD License: https://opensource.org/licenses/BSD-3-Clause
 
+import re
+
 from sqlparse import sql, tokens as T
 from sqlparse.utils import split_unquoted_newlines
 
 
 class StripCommentsFilter:
+
     @staticmethod
     def _process(tlist):
         def get_next_comment():
             # TODO(andi) Comment types should be unified, see related issue38
             return tlist.token_next_by(i=sql.Comment, t=T.Comment)
 
+        def _get_insert_token(token):
+            """Returns either a whitespace or the line breaks from token."""
+            # See issue484 why line breaks should be preserved.
+            m = re.search(r'((\r\n|\r|\n)+) *$', token.value)
+            if m is not None:
+                return sql.Token(T.Whitespace.Newline, m.groups()[0])
+            else:
+                return sql.Token(T.Whitespace, ' ')
+
         tidx, token = get_next_comment()
         while token:
             pidx, prev_ = tlist.token_prev(tidx, skip_ws=False)
@@ -26,15 +38,12 @@ def get_next_comment():
                     or prev_.is_whitespace or prev_.match(T.Punctuation, '(')
                     or next_.is_whitespace or next_.match(T.Punctuation, ')')):
                 # Insert a whitespace to ensure the following SQL produces
-                # a valid SQL (see #425). For example:
-                #
-                # Before: select a--comment\nfrom foo
-                # After: select a from foo
-                if prev_ is not None and next_ is None:
-                    tlist.tokens.insert(tidx, sql.Token(T.Whitespace, ' '))
+                # a valid SQL (see #425).
+                if prev_ is not None and not prev_.match(T.Punctuation, '('):
+                    tlist.tokens.insert(tidx, _get_insert_token(token))
                 tlist.tokens.remove(token)
             else:
-                tlist.tokens[tidx] = sql.Token(T.Whitespace, ' ')
+                tlist.tokens[tidx] = _get_insert_token(token)
 
             tidx, token = get_next_comment()
 
diff --git a/tests/test_format.py b/tests/test_format.py
@@ -41,26 +41,26 @@ def test_identifiercase_quotes(self):
     def test_strip_comments_single(self):
         sql = 'select *-- statement starts here\nfrom foo'
         res = sqlparse.format(sql, strip_comments=True)
-        assert res == 'select * from foo'
+        assert res == 'select *\nfrom foo'
         sql = 'select * -- statement starts here\nfrom foo'
         res = sqlparse.format(sql, strip_comments=True)
-        assert res == 'select * from foo'
+        assert res == 'select *\nfrom foo'
         sql = 'select-- foo\nfrom -- bar\nwhere'
         res = sqlparse.format(sql, strip_comments=True)
-        assert res == 'select from where'
+        assert res == 'select\nfrom\nwhere'
         sql = 'select *-- statement starts here\n\nfrom foo'
         res = sqlparse.format(sql, strip_comments=True)
-        assert res == 'select * from foo'
+        assert res == 'select *\n\nfrom foo'
         sql = 'select * from foo-- statement starts here\nwhere'
         res = sqlparse.format(sql, strip_comments=True)
-        assert res == 'select * from foo where'
+        assert res == 'select * from foo\nwhere'
         sql = 'select a-- statement starts here\nfrom foo'
         res = sqlparse.format(sql, strip_comments=True)
-        assert res == 'select a from foo'
+        assert res == 'select a\nfrom foo'
         sql = '--comment\nselect a-- statement starts here\n' \
               'from foo--comment\nf'
         res = sqlparse.format(sql, strip_comments=True)
-        assert res == 'select a from foo f'
+        assert res == 'select a\nfrom foo\nf'
 
     def test_strip_comments_invalid_option(self):
         sql = 'select-- foo\nfrom -- bar\nwhere'
