Preserve trailing whitespace after the final semicolon

sqlparse/engine/statement_splitter.py

@@ -152,14 +152,22 @@ def process(self, stream):
         """Process the stream"""
         EOS_TTYPE = T.Whitespace, T.Comment.Single
 
+        # Tokens of the previous completed statement, held back by one step so
+        # a trailing whitespace-only run (e.g. a newline after the final ';')
+        # can be re-attached to it instead of being dropped, keeping the
+        # round-trip lossless.
+        completed = None
+
         # 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:
-                yield sql.Statement(self.tokens)
+                if completed is not None:
+                    yield sql.Statement(completed)
+                completed = self.tokens
 
                 # Reset filter and prepare to process next statement
                 self._reset()
@@ -191,6 +199,19 @@ def process(self, stream):
                 # token but not for BEGIN itself (which just set the flag)
                 self._seen_begin = False
 
-        # Yield pending statement (if any)
-        if self.tokens and not all(t.is_whitespace for t in self.tokens):
-            yield sql.Statement(self.tokens)
+        # Decide what to do with the final run of tokens.
+        if self.tokens:
+            if completed is not None and all(t.is_whitespace for t in self.tokens):
+                # A trailing whitespace-only run after a completed statement
+                # (a newline after the last ';'): re-attach it so the text is
+                # preserved, rather than emitting it as an empty statement.
+                completed.extend(self.tokens)
+            else:
+                if completed is not None:
+                    yield sql.Statement(completed)
+                completed = self.tokens
+
+        # Yield pending statement (if any); a leading whitespace-only run with
+        # no statement is dropped, as before.
+        if completed is not None and not all(t.is_whitespace for t in completed):
+            yield sql.Statement(completed)

tests/test_split.py

@@ -50,6 +50,26 @@ def test_split_dashcomments_eol(s):
     assert len(stmts) == 1
 
 
+@pytest.mark.parametrize('s', ['select 1;\n',
+                               'select 1;\n\n',
+                               'select 1; \n',
+                               'select 1;  \n    ',
+                               'select 1;\t',
+                               'select 1;   '])
+def test_split_trailing_whitespace_preserved(s):
+    # A trailing whitespace run after the final ';' must not be dropped, so
+    # concatenating the parsed statements reproduces the input exactly.
+    stmts = sqlparse.parse(s)
+    assert len(stmts) == 1
+    assert ''.join(str(stmt) for stmt in stmts) == s
+
+
+@pytest.mark.parametrize('s', ['', '   ', '\n\n', '\t'])
+def test_split_whitespace_only(s):
+    # Whitespace-only input yields no statements.
+    assert sqlparse.parse(s) == ()
+
+
 def test_split_begintag(load_file):
     sql = load_file('begintag.sql')
     stmts = sqlparse.parse(sql)