Break global and nonlocal statements over continuation lines

astral-sh · Jul 29, 2023 · 6674125 · 6674125
1 parent 161ed4a
commit 6674125
Show file tree

Hide file tree

Showing 6 changed files with 114 additions and 10 deletions.
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/global.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/global.py
@@ -10,3 +10,7 @@ def f():
 
 def f():
     global analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model
+
+
+def f():
+    global analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model  # end-of-line comment
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/nonlocal.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/nonlocal.py
@@ -10,3 +10,7 @@ def f():
 
 def f():
     nonlocal analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model
+
+
+def f():
+    nonlocal analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model  # end-of-line comment
diff --git a/crates/ruff_python_formatter/src/statement/stmt_global.rs b/crates/ruff_python_formatter/src/statement/stmt_global.rs
@@ -1,4 +1,5 @@
 use ruff_formatter::{format_args, write, Buffer, FormatResult};
+use ruff_python_ast::node::AstNode;
 use ruff_python_ast::StmtGlobal;
 
 use crate::prelude::*;
@@ -9,10 +10,44 @@ pub struct FormatStmtGlobal;
 
 impl FormatNodeRule<StmtGlobal> for FormatStmtGlobal {
     fn fmt_fields(&self, item: &StmtGlobal, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [text("global"), space()])?;
+        // Join the `global` names, breaking across continuation lines if necessary, unless the
+        // `global` statement has a trailing comment, in which case, breaking the names would
+        // move the comment "off" of the `global` statement.
+        if f.context()
+            .comments()
+            .has_trailing_comments(item.as_any_node_ref())
+        {
+            let joined = format_with(|f| {
+                f.join_with(format_args![text(","), space()])
+                    .entries(item.names.iter().formatted())
+                    .finish()
+            });
 
-        f.join_with(format_args![text(","), space()])
-            .entries(item.names.iter().formatted())
-            .finish()
+            write!(f, [text("global"), space(), &joined])
+        } else {
+            let joined = format_with(|f| {
+                f.join_with(&format_args![
+                    text(","),
+                    space(),
+                    if_group_breaks(&text("\\")),
+                    soft_line_break(),
+                ])
+                .entries(item.names.iter().formatted())
+                .finish()
+            });
+
+            write!(
+                f,
+                [
+                    text("global"),
+                    space(),
+                    &group(&format_args!(
+                        if_group_breaks(&text("\\")),
+                        soft_line_break(),
+                        &soft_block_indent(&joined)
+                    ))
+                ]
+            )
+        }
     }
 }
diff --git a/crates/ruff_python_formatter/src/statement/stmt_nonlocal.rs b/crates/ruff_python_formatter/src/statement/stmt_nonlocal.rs
@@ -1,4 +1,5 @@
 use ruff_formatter::{format_args, write, Buffer, FormatResult};
+use ruff_python_ast::node::AstNode;
 use ruff_python_ast::StmtNonlocal;
 
 use crate::prelude::*;
@@ -9,10 +10,44 @@ pub struct FormatStmtNonlocal;
 
 impl FormatNodeRule<StmtNonlocal> for FormatStmtNonlocal {
     fn fmt_fields(&self, item: &StmtNonlocal, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [text("nonlocal"), space()])?;
+        // Join the `nonlocal` names, breaking across continuation lines if necessary, unless the
+        // `nonlocal` statement has a trailing comment, in which case, breaking the names would
+        // move the comment "off" of the `nonlocal` statement.
+        if f.context()
+            .comments()
+            .has_trailing_comments(item.as_any_node_ref())
+        {
+            let joined = format_with(|f| {
+                f.join_with(format_args![text(","), space()])
+                    .entries(item.names.iter().formatted())
+                    .finish()
+            });
 
-        f.join_with(format_args![text(","), space()])
-            .entries(item.names.iter().formatted())
-            .finish()
+            write!(f, [text("nonlocal"), space(), &joined])
+        } else {
+            let joined = format_with(|f| {
+                f.join_with(&format_args![
+                    text(","),
+                    space(),
+                    if_group_breaks(&text("\\")),
+                    soft_line_break(),
+                ])
+                .entries(item.names.iter().formatted())
+                .finish()
+            });
+
+            write!(
+                f,
+                [
+                    text("nonlocal"),
+                    space(),
+                    &group(&format_args!(
+                        if_group_breaks(&text("\\")),
+                        soft_line_break(),
+                        &soft_block_indent(&joined)
+                    ))
+                ]
+            )
+        }
     }
 }
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__global.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__global.py.snap
@@ -16,6 +16,10 @@ def f():
 
 def f():
     global analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model
+
+
+def f():
+    global analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model  # end-of-line comment
 ```
 
 ## Output
@@ -31,7 +35,16 @@ def f():
 
 
 def f():
-    global analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model
+    global \
+        analyze_featuremap_layer, \
+        analyze_featuremapcompression_layer, \
+        analyze_latencies_post, \
+        analyze_motions_layer, \
+        analyze_size_model
+
+
+def f():
+    global analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model  # end-of-line comment
 ```
 
 

diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__nonlocal.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__nonlocal.py.snap
@@ -16,6 +16,10 @@ def f():
 
 def f():
     nonlocal analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model
+
+
+def f():
+    nonlocal analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model  # end-of-line comment
 ```
 
 ## Output
@@ -31,7 +35,16 @@ def f():
 
 
 def f():
-    nonlocal analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model
+    nonlocal \
+        analyze_featuremap_layer, \
+        analyze_featuremapcompression_layer, \
+        analyze_latencies_post, \
+        analyze_motions_layer, \
+        analyze_size_model
+
+
+def f():
+    nonlocal analyze_featuremap_layer, analyze_featuremapcompression_layer, analyze_latencies_post, analyze_motions_layer, analyze_size_model  # end-of-line comment
 ```