Preserve comments on non-defaulted arguments (#3264)

crates/ruff_python_formatter/src/attachment.rs

@@ -1,6 +1,6 @@
 use crate::core::visitor;
 use crate::core::visitor::Visitor;
-use crate::cst::{Alias, Body, Excepthandler, Expr, Pattern, SliceIndex, Stmt};
+use crate::cst::{Alias, Arg, Body, Excepthandler, Expr, Pattern, SliceIndex, Stmt};
 use crate::trivia::{decorate_trivia, TriviaIndex, TriviaToken};
 
 struct AttachmentVisitor {
@@ -40,6 +40,14 @@ impl<'a> Visitor<'a> for AttachmentVisitor {
         visitor::walk_alias(self, alias);
     }
 
+    fn visit_arg(&mut self, arg: &'a mut Arg) {
+        let trivia = self.index.arg.remove(&arg.id());
+        if let Some(comments) = trivia {
+            arg.trivia.extend(comments);
+        }
+        visitor::walk_arg(self, arg);
+    }
+
     fn visit_excepthandler(&mut self, excepthandler: &'a mut Excepthandler) {
         let trivia = self.index.excepthandler.remove(&excepthandler.id());
         if let Some(comments) = trivia {

crates/ruff_python_formatter/src/format/arg.rs

@@ -4,6 +4,7 @@ use ruff_text_size::TextSize;
 
 use crate::context::ASTFormatContext;
 use crate::cst::Arg;
+use crate::format::comments::end_of_line_comments;
 use crate::shared_traits::AsFormat;
 
 pub struct FormatArg<'a> {
@@ -27,6 +28,7 @@ impl Format<ASTFormatContext<'_>> for FormatArg<'_> {
             write!(f, [text(": ")])?;
             write!(f, [annotation.format()])?;
         }
+        write!(f, [end_of_line_comments(arg)])?;
 
         Ok(())
     }

crates/ruff_python_formatter/src/format/expr.rs

@@ -71,9 +71,8 @@ fn format_subscript(
             Ok(())
         }))])]
     )?;
-
     write!(f, [text("]")])?;
-
+    write!(f, [end_of_line_comments(expr)])?;
     Ok(())
 }
 
@@ -286,6 +285,7 @@ fn format_list(
         )?;
     }
     write!(f, [text("]")])?;
+    write!(f, [end_of_line_comments(expr)])?;
     Ok(())
 }
 
@@ -613,6 +613,7 @@ fn format_joined_str(
     _values: &[Expr],
 ) -> FormatResult<()> {
     write!(f, [literal(Range::from_located(expr))])?;
+    write!(f, [end_of_line_comments(expr)])?;
     Ok(())
 }
 
@@ -639,6 +640,7 @@ fn format_constant(
         Constant::Complex { .. } => write!(f, [complex_literal(Range::from_located(expr))])?,
         Constant::Tuple(_) => unreachable!("Constant::Tuple should be handled by format_tuple"),
     }
+    write!(f, [end_of_line_comments(expr)])?;
     Ok(())
 }
 
@@ -713,9 +715,7 @@ fn format_attribute(
     write!(f, [value.format()])?;
     write!(f, [text(".")])?;
     write!(f, [dynamic_text(attr, TextSize::default())])?;
-
     write!(f, [end_of_line_comments(expr)])?;
-
     Ok(())
 }
 
@@ -729,9 +729,7 @@ fn format_named_expr(
     write!(f, [text(":=")])?;
     write!(f, [space()])?;
     write!(f, [group(&format_args![value.format()])])?;
-
     write!(f, [end_of_line_comments(expr)])?;
-
     Ok(())
 }
 
@@ -752,9 +750,7 @@ fn format_bool_op(
             write!(f, [group(&format_args![value.format()])])?;
         }
     }
-
     write!(f, [end_of_line_comments(expr)])?;
-
     Ok(())
 }
 
@@ -767,7 +763,6 @@ fn format_bin_op(
 ) -> FormatResult<()> {
     // https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#line-breaks-binary-operators
     let is_simple = matches!(op, Operator::Pow) && is_simple_power(left) && is_simple_power(right);
-
     write!(f, [left.format()])?;
     if !is_simple {
         write!(f, [soft_line_break_or_space()])?;
@@ -776,10 +771,8 @@ fn format_bin_op(
     if !is_simple {
         write!(f, [space()])?;
     }
-    write!(f, [group(&format_args![right.format()])])?;
-
+    write!(f, [group(&right.format())])?;
     write!(f, [end_of_line_comments(expr)])?;
-
     Ok(())
 }
 
@@ -808,6 +801,7 @@ fn format_unary_op(
     } else {
         write!(f, [operand.format()])?;
     }
+    write!(f, [end_of_line_comments(expr)])?;
     Ok(())
 }
 
@@ -825,6 +819,7 @@ fn format_lambda(
     write!(f, [text(":")])?;
     write!(f, [space()])?;
     write!(f, [body.format()])?;
+    write!(f, [end_of_line_comments(expr)])?;
     Ok(())
 }
 

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments2_py.snap

@@ -178,43 +178,9 @@ instruction()#comment with bad spacing
 ```diff
 --- Black
 +++ Ruff
-@@ -13,7 +13,7 @@
-     "Callable",
-     "ClassVar",
-     # ABCs (from collections.abc).
--    "AbstractSet",  # collections.abc.Set.
-+    "AbstractSet",
-     "ByteString",
-     "Container",
-     # Concrete collection types.
-@@ -24,7 +24,7 @@
-     "List",
-     "Set",
-     "FrozenSet",
--    "NamedTuple",  # Not really a type.
-+    "NamedTuple",
-     "Generator",
- ]
- 
-@@ -60,26 +60,32 @@
- # Comment before function.
- def inline_comments_in_brackets_ruin_everything():
-     if typedargslist:
--        parameters.children = [children[0], body, children[-1]]  # (1  # )1
-+        parameters.children = [children[0], body, children[-1]]
-         parameters.children = [
-             children[0],
-             body,
--            children[-1],  # type: ignore
-+            children[-1],
-         ]
-     else:
-         parameters.children = [
--            parameters.children[0],  # (2 what if this was actually long
-+            parameters.children[0],
+@@ -72,14 +72,20 @@
              body,
--            parameters.children[-1],  # )2
-+            parameters.children[-1],
+             parameters.children[-1],  # )2
          ]
 -        parameters.children = [parameters.what_if_this_was_actually_long.children[0], body, parameters.children[-1]]  # type: ignore
 +        parameters.children = [
@@ -319,7 +285,7 @@ __all__ = [
     "Callable",
     "ClassVar",
     # ABCs (from collections.abc).
-    "AbstractSet",
+    "AbstractSet",  # collections.abc.Set.
     "ByteString",
     "Container",
     # Concrete collection types.
@@ -330,7 +296,7 @@ __all__ = [
     "List",
     "Set",
     "FrozenSet",
-    "NamedTuple",
+    "NamedTuple",  # Not really a type.
     "Generator",
 ]
 
@@ -366,17 +332,17 @@ else:
 # Comment before function.
 def inline_comments_in_brackets_ruin_everything():
     if typedargslist:
-        parameters.children = [children[0], body, children[-1]]
+        parameters.children = [children[0], body, children[-1]]  # (1  # )1
         parameters.children = [
             children[0],
             body,
-            children[-1],
+            children[-1],  # type: ignore
         ]
     else:
         parameters.children = [
-            parameters.children[0],
+            parameters.children[0],  # (2 what if this was actually long
             body,
-            parameters.children[-1],
+            parameters.children[-1],  # )2
         ]
         parameters.children = [
             parameters.what_if_this_was_actually_long.children[0],