Preserve trailing newline after a statement's semicolon (fix lossless round-trip)

sqlparse/engine/statement_splitter.py

@@ -150,15 +150,18 @@ def _change_splitlevel(self, ttype, value):
 
     def process(self, stream):
         """Process the stream"""
-        EOS_TTYPE = T.Whitespace, T.Comment.Single
-
         # Run over all stream tokens
         for ttype, value in stream:
-            # Yield token if we finished a statement and there's no whitespaces
-            # It will count newline token as a non whitespace. In this context
-            # whitespace ignores newlines.
-            # why don't multi line comments also count?
-            if self.consume_ws and ttype not in EOS_TTYPE:
+            # Once a statement is finished (consume_ws), keep trailing
+            # whitespace and single-line comments with it, and only cut a
+            # new statement when real content arrives. Newlines are
+            # whitespace too, so a terminating semicolon's trailing newline
+            # stays with its statement instead of being split off into a
+            # whitespace-only remainder that gets dropped at end of stream
+            # (which broke str(parse(sql)) == sql for inputs ending in
+            # ";\n", the common end-of-file case).
+            if self.consume_ws and not (
+                    ttype in T.Whitespace or ttype in T.Comment.Single):
                 yield sql.Statement(self.tokens)
 
                 # Reset filter and prepare to process next statement

tests/test_split.py

@@ -152,6 +152,25 @@ def test_split_ignores_empty_newlines():
     assert stmts[1] == 'select bar;'
 
 
+@pytest.mark.parametrize('sql', [
+    'select 1;\n',
+    'select 1;\n\n',
+    'select 1;\r\n',
+    'select 1; ',
+    'select 1;\t',
+    'select 1;\nselect 2;\n',
+    'create table t (a int);\n\n',
+])
+def test_split_preserves_trailing_whitespace(sql):
+    # The whitespace after a terminating semicolon must stay with the
+    # statement so that the parsed statements round-trip back to the input.
+    # A trailing newline (the common end-of-file "...;\n" case) used to be
+    # split off into a whitespace-only remainder and dropped, unlike a
+    # trailing space or tab.
+    stmts = sqlparse.parse(sql)
+    assert ''.join(str(q) for q in stmts) == sql
+
+
 def test_split_quotes_with_new_line():
     stmts = sqlparse.split('select "foo\nbar"')
     assert len(stmts) == 1