Use line_suffix for end-of-line comments (#2975)

astral-sh · Feb 16, 2023 · 6088a36 · 6088a36
1 parent 66a162f
commit 6088a36
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 23 deletions.
diff --git a/crates/ruff_python_formatter/src/format/expr.rs b/crates/ruff_python_formatter/src/format/expr.rs
@@ -1,9 +1,10 @@
 #![allow(unused_variables, clippy::too_many_arguments)]
 
+use rustpython_parser::ast::Constant;
+
 use ruff_formatter::prelude::*;
 use ruff_formatter::{format_args, write};
 use ruff_text_size::TextSize;
-use rustpython_parser::ast::Constant;
 
 use crate::builders::literal;
 use crate::context::ASTFormatContext;
@@ -40,9 +41,9 @@ fn format_starred(
         }
     }) {
         if std::mem::take(&mut first) {
-            write!(f, [text("  ")])?;
+            write!(f, [line_suffix(&text("  "))])?;
         }
-        write!(f, [literal(range)])?;
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     Ok(())
@@ -69,9 +70,9 @@ fn format_name(
         }
     }) {
         if std::mem::take(&mut first) {
-            write!(f, [text("  ")])?;
+            write!(f, [line_suffix(&text("  "))])?;
         }
-        write!(f, [literal(range)])?;
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     Ok(())
@@ -283,7 +284,9 @@ fn format_set(
                             write!(f, [text(",")])?;
                             write!(f, [soft_line_break_or_space()])?;
                         } else {
-                            write!(f, [if_group_breaks(&text(","))])?;
+                            if magic_trailing_comma {
+                                write!(f, [if_group_breaks(&text(","))])?;
+                            }
                         }
                     }
                     Ok(())
@@ -306,9 +309,47 @@ fn format_call(
     if args.is_empty() && keywords.is_empty() {
         write!(f, [text("(")])?;
         write!(f, [text(")")])?;
+
+        // Apply any inline comments.
+        let mut first = true;
+        for range in expr.trivia.iter().filter_map(|trivia| {
+            if matches!(trivia.relationship, Relationship::Trailing) {
+                if let TriviaKind::InlineComment(range) = trivia.kind {
+                    Some(range)
+                } else {
+                    None
+                }
+            } else {
+                None
+            }
+        }) {
+            if std::mem::take(&mut first) {
+                write!(f, [line_suffix(&text("  "))])?;
+            }
+            write!(f, [line_suffix(&literal(range))])?;
+        }
     } else {
         write!(f, [text("(")])?;
 
+        // Apply any inline comments.
+        let mut first = true;
+        for range in expr.trivia.iter().filter_map(|trivia| {
+            if matches!(trivia.relationship, Relationship::Trailing) {
+                if let TriviaKind::InlineComment(range) = trivia.kind {
+                    Some(range)
+                } else {
+                    None
+                }
+            } else {
+                None
+            }
+        }) {
+            if std::mem::take(&mut first) {
+                write!(f, [line_suffix(&text("  "))])?;
+            }
+            write!(f, [line_suffix(&literal(range))])?;
+        }
+
         let magic_trailing_comma = expr
             .trivia
             .iter()
@@ -551,6 +592,26 @@ fn format_compare(
         write!(f, [space()])?;
         write!(f, [group(&format_args![comparators[i].format()])])?;
     }
+
+    // Apply any inline comments.
+    let mut first = true;
+    for range in expr.trivia.iter().filter_map(|trivia| {
+        if matches!(trivia.relationship, Relationship::Trailing) {
+            if let TriviaKind::InlineComment(range) = trivia.kind {
+                Some(range)
+            } else {
+                None
+            }
+        } else {
+            None
+        }
+    }) {
+        if std::mem::take(&mut first) {
+            write!(f, [line_suffix(&text("  "))])?;
+        }
+        write!(f, [line_suffix(&literal(range))])?;
+    }
+
     Ok(())
 }
 
@@ -662,9 +723,9 @@ fn format_attribute(
         }
     }) {
         if std::mem::take(&mut first) {
-            write!(f, [text("  ")])?;
+            write!(f, [line_suffix(&text("  "))])?;
         }
-        write!(f, [literal(range)])?;
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     Ok(())
@@ -702,9 +763,9 @@ fn format_bool_op(
         }
     }) {
         if std::mem::take(&mut first) {
-            write!(f, [text("  ")])?;
+            write!(f, [line_suffix(&text("  "))])?;
         }
-        write!(f, [literal(range)])?;
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     Ok(())
@@ -744,9 +805,9 @@ fn format_bin_op(
         }
     }) {
         if std::mem::take(&mut first) {
-            write!(f, [text("  ")])?;
+            write!(f, [line_suffix(&text("  "))])?;
         }
-        write!(f, [literal(range)])?;
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     Ok(())

diff --git a/crates/ruff_python_formatter/src/format/stmt.rs b/crates/ruff_python_formatter/src/format/stmt.rs
@@ -34,9 +34,9 @@ fn format_pass(f: &mut Formatter<ASTFormatContext<'_>>, stmt: &Stmt) -> FormatRe
         }
     }) {
         if std::mem::take(&mut first) {
-            write!(f, [text("  ")])?;
+            write!(f, [line_suffix(&text("  "))])?;
         }
-        write!(f, [literal(range)])?;
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     Ok(())
@@ -220,9 +220,9 @@ fn format_func_def(
         }
     }) {
         if std::mem::take(&mut first) {
-            write!(f, [text("  ")])?;
+            write!(f, [line_suffix(&text("  "))])?;
         }
-        write!(f, [literal(range)])?;
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     write!(f, [block_indent(&format_args![block(body)])])
@@ -269,9 +269,9 @@ fn format_assign(
         }
     }) {
         if std::mem::take(&mut first) {
-            write!(f, [text("  ")])?;
+            write!(f, [line_suffix(&text("  "))])?;
         }
-        write!(f, [literal(range)])?;
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     Ok(())
@@ -557,9 +557,9 @@ fn format_import_from(
         }
     }) {
         if std::mem::take(&mut first) {
-            write!(f, [text("  ")])?;
+            write!(f, [line_suffix(&text("  "))])?;
         }
-        write!(f, [literal(range)])?;
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     Ok(())
@@ -606,9 +606,9 @@ fn format_expr(
         }
     }) {
         if std::mem::take(&mut first) {
-            write!(f, [text("  ")])?;
+            write!(f, [line_suffix(&text("  "))])?;
         }
-        write!(f, [literal(range)])?;
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     Ok(())

diff --git a/crates/ruff_python_formatter/src/lib.rs b/crates/ruff_python_formatter/src/lib.rs
@@ -69,6 +69,7 @@ mod tests {
     #[test_case(Path::new("simple_cases/class_methods_new_line.py"); "class_methods_new_line")]
     #[test_case(Path::new("simple_cases/beginning_backslash.py"); "beginning_backslash")]
     #[test_case(Path::new("simple_cases/import_spacing.py"); "import_spacing")]
+    #[test_case(Path::new("simple_cases/power_op_spacing.py"); "power_op_spacing")]
     fn passing(path: &Path) -> Result<()> {
         let snapshot = format!("{}", path.display());
         let content = std::fs::read_to_string(test_resource_path(
@@ -103,7 +104,8 @@ mod tests {
     #[test_case(Path::new("simple_cases/expression.py"); "expression")]
     #[test_case(Path::new("simple_cases/function.py"); "function")]
     #[test_case(Path::new("simple_cases/function2.py"); "function2")]
-    #[test_case(Path::new("simple_cases/power_op_spacing.py"); "power_op_spacing")]
+    #[test_case(Path::new("simple_cases/function_trailing_comma.py"); "function_trailing_comma")]
+    #[test_case(Path::new("simple_cases/composition.py"); "composition")]
     fn failing(path: &Path) -> Result<()> {
         let snapshot = format!("{}", path.display());
         let content = std::fs::read_to_string(test_resource_path(

diff --git a/...le_cases__power_op_spacing.py.snap.expect → ...s__simple_cases__power_op_spacing.py.snap b/...le_cases__power_op_spacing.py.snap.expect → ...s__simple_cases__power_op_spacing.py.snap