From 3ccca332bda5e4efcc9a60671771aa4032e15dd6 Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Mon, 30 Oct 2023 09:22:34 +0900
Subject: [PATCH] Preserve trailing semicolons when using `fmt: off` (#8275)

---
 .../ruff/fmt_on_off/trailing_semicolon.py     | 27 ++++++++
 crates/ruff_python_formatter/src/verbatim.rs  | 20 +++---
 ...mat@fmt_on_off__trailing_semicolon.py.snap | 69 +++++++++++++++++++
 3 files changed, 108 insertions(+), 8 deletions(-)
 create mode 100644 crates/ruff_python_formatter/resources/test/fixtures/ruff/fmt_on_off/trailing_semicolon.py
 create mode 100644 crates/ruff_python_formatter/tests/snapshots/format@fmt_on_off__trailing_semicolon.py.snap

diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/fmt_on_off/trailing_semicolon.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/fmt_on_off/trailing_semicolon.py
new file mode 100644
index 0000000000000..f7de6e457f348
--- /dev/null
+++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/fmt_on_off/trailing_semicolon.py
@@ -0,0 +1,27 @@
+def f():
+    # fmt: off
+    a = 10
+
+    if True:
+        with_semicolon = 10   \
+            ;
+
+formatted   = true;
+
+
+def f():
+    # fmt: off
+
+    if True:
+        with_semicolon = 20 \
+            ; # comment
+
+
+# fmt: off
+statement = 0 \
+    ;
+# fmt: on
+a = 10
+
+# fmt: off
+last_statement_with_semi   ;
diff --git a/crates/ruff_python_formatter/src/verbatim.rs b/crates/ruff_python_formatter/src/verbatim.rs
index 00f8e149609ec..5c29b0039d2b5 100644
--- a/crates/ruff_python_formatter/src/verbatim.rs
+++ b/crates/ruff_python_formatter/src/verbatim.rs
@@ -395,14 +395,18 @@ fn write_suppressed_statements<'a>(
             statement = SuiteChildStatement::Other(next_statement);
             leading_node_comments = comments.leading(next_statement);
         } else {
-            let mut nodes =
-                std::iter::successors(Some(AnyNodeRef::from(statement.statement())), |statement| {
-                    statement.last_child_in_body()
-                });
-
-            let end = nodes
-                .find_map(|statement| comments.trailing(statement).last().map(Ranged::end))
-                .unwrap_or(statement.end());
+            let mut current = AnyNodeRef::from(statement.statement());
+            // Expand the range of the statement to include any trailing comments or semicolons.
+            let end = loop {
+                if let Some(comment) = comments.trailing(current).last() {
+                    break comment.end();
+                } else if let Some(child) = current.last_child_in_body() {
+                    current = child;
+                } else {
+                    break trailing_semicolon(current, source)
+                        .map_or(statement.end(), TextRange::end);
+                }
+            };
 
             FormatVerbatimStatementRange {
                 verbatim_range: TextRange::new(format_off_comment.end(), end),
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@fmt_on_off__trailing_semicolon.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@fmt_on_off__trailing_semicolon.py.snap
new file mode 100644
index 0000000000000..463294c542dc0
--- /dev/null
+++ b/crates/ruff_python_formatter/tests/snapshots/format@fmt_on_off__trailing_semicolon.py.snap
@@ -0,0 +1,69 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/fmt_on_off/trailing_semicolon.py
+---
+## Input
+```py
+def f():
+    # fmt: off
+    a = 10
+
+    if True:
+        with_semicolon = 10   \
+            ;
+
+formatted   = true;
+
+
+def f():
+    # fmt: off
+
+    if True:
+        with_semicolon = 20 \
+            ; # comment
+
+
+# fmt: off
+statement = 0 \
+    ;
+# fmt: on
+a = 10
+
+# fmt: off
+last_statement_with_semi   ;
+```
+
+## Output
+```py
+def f():
+    # fmt: off
+    a = 10
+
+    if True:
+        with_semicolon = 10   \
+            ;
+
+
+formatted = true
+
+
+def f():
+    # fmt: off
+
+    if True:
+        with_semicolon = 20 \
+            ; # comment
+
+
+# fmt: off
+statement = 0 \
+    ;
+# fmt: on
+a = 10
+
+# fmt: off
+last_statement_with_semi   ;
+```
+
+
+