Fix ArgWithDefault comments handling

Cargo.toml

@@ -50,15 +50,15 @@ toml = { version = "0.7.2" }
 # v0.0.1
 libcst = { git = "https://github.com/charliermarsh/LibCST", rev = "80e4c1399f95e5beb532fdd1e209ad2dbb470438" }
 # v0.0.3
-ruff_text_size = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "ed3b4eb72b6e497bbdb4d19dec6621074d724130" }
+ruff_text_size = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "08ebbe40d7776cac6e3ba66277d435056f2b8dca" }
 # v0.0.3
-rustpython-ast = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "ed3b4eb72b6e497bbdb4d19dec6621074d724130" , default-features = false, features = ["all-nodes-with-ranges", "num-bigint"]}
+rustpython-ast = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "08ebbe40d7776cac6e3ba66277d435056f2b8dca" , default-features = false, features = ["all-nodes-with-ranges", "num-bigint"]}
 # v0.0.3
-rustpython-format = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "ed3b4eb72b6e497bbdb4d19dec6621074d724130", default-features = false, features = ["num-bigint"] }
+rustpython-format = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "08ebbe40d7776cac6e3ba66277d435056f2b8dca", default-features = false, features = ["num-bigint"] }
 # v0.0.3
-rustpython-literal = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "ed3b4eb72b6e497bbdb4d19dec6621074d724130", default-features = false }
+rustpython-literal = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "08ebbe40d7776cac6e3ba66277d435056f2b8dca", default-features = false }
 # v0.0.3
-rustpython-parser = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "ed3b4eb72b6e497bbdb4d19dec6621074d724130" , default-features = false, features = ["full-lexer", "all-nodes-with-ranges", "num-bigint"] }
+rustpython-parser = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "08ebbe40d7776cac6e3ba66277d435056f2b8dca" , default-features = false, features = ["full-lexer", "all-nodes-with-ranges", "num-bigint"] }
 
 [profile.release]
 lto = "fat"

crates/ruff_python_ast/src/node.rs

@@ -3680,7 +3680,7 @@ impl AnyNodeRef<'_> {
 
     /// Compares two any node refs by their pointers (referential equality).
     pub fn ptr_eq(self, other: AnyNodeRef) -> bool {
-        self.as_ptr().eq(&other.as_ptr())
+        self.as_ptr().eq(&other.as_ptr()) && self.kind() == other.kind()
     }
 
     /// Returns the node's [`kind`](NodeKind) that has no data associated and is [`Copy`].

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py

@@ -90,3 +90,10 @@ def f2():
 # Regression test for https://github.com/python/cpython/blob/7199584ac8632eab57612f595a7162ab8d2ebbc0/Lib/warnings.py#L513
 def f(arg1=1, *, kwonlyarg1, kwonlyarg2=2):
     pass
+
+
+# Regression test for https://github.com/astral-sh/ruff/issues/5176#issuecomment-1598171989
+def foo(
+    b=3 + 2  # comment
+):
+    ...

crates/ruff_python_formatter/src/cli.rs

@@ -63,7 +63,7 @@ pub fn format_and_debug_print(input: &str, cli: &Cli) -> Result<String> {
     }
     if cli.print_comments {
         println!(
-            "{:?}",
+            "{:#?}",
             formatted.context().comments().debug(SourceCode::new(input))
         );
     }

crates/ruff_python_formatter/src/comments/placement.rs

@@ -627,17 +627,7 @@ fn handle_positional_only_arguments_separator_comment<'a>(
     };
 
     let is_last_positional_argument =
-        // If the preceding node is the identifier for the last positional argument (`a`).
-        // ```python
-        // def test(a, /, b): pass
-        // ```
-        are_same_optional(last_argument_or_default, arguments.posonlyargs.last().map(|arg| &arg.def))
-            // If the preceding node is the default for the last positional argument (`10`).
-            // ```python
-            // def test(a=10, /, b): pass
-            // ```
-            || are_same_optional(last_argument_or_default, arguments
-                .posonlyargs.last().and_then(|arg| arg.default.as_deref()));
+        are_same_optional(last_argument_or_default, arguments.posonlyargs.last());
 
     if !is_last_positional_argument {
         return CommentPlacement::Default(comment);

crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__non_positional_arguments_slash_on_same_line.snap

@@ -19,9 +19,9 @@ expression: comments.debug(test_case.source_code)
         "trailing": [],
     },
     Node {
-        kind: Arg,
-        range: 90..91,
-        source: `b`,
+        kind: ArgWithDefault,
+        range: 90..94,
+        source: `b=20`,
     }: {
         "leading": [
             SourceComment {

crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__non_positional_arguments_with_defaults.snap

@@ -24,9 +24,9 @@ expression: comments.debug(test_case.source_code)
         "trailing": [],
     },
     Node {
-        kind: ExprConstant,
-        range: 17..19,
-        source: `10`,
+        kind: ArgWithDefault,
+        range: 15..19,
+        source: `a=10`,
     }: {
         "leading": [],
         "dangling": [],
@@ -39,9 +39,9 @@ expression: comments.debug(test_case.source_code)
         ],
     },
     Node {
-        kind: Arg,
-        range: 173..174,
-        source: `b`,
+        kind: ArgWithDefault,
+        range: 173..177,
+        source: `b=20`,
     }: {
         "leading": [
             SourceComment {

crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__positional_argument_only_comment.snap

@@ -24,7 +24,7 @@ expression: comments.debug(test_case.source_code)
         "trailing": [],
     },
     Node {
-        kind: Arg,
+        kind: ArgWithDefault,
         range: 15..16,
         source: `a`,
     }: {
@@ -39,7 +39,7 @@ expression: comments.debug(test_case.source_code)
         ],
     },
     Node {
-        kind: Arg,
+        kind: ArgWithDefault,
         range: 166..167,
         source: `b`,
     }: {

crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__positional_argument_only_comment_without_following_node.snap

@@ -24,7 +24,7 @@ expression: comments.debug(test_case.source_code)
         "trailing": [],
     },
     Node {
-        kind: Arg,
+        kind: ArgWithDefault,
         range: 15..16,
         source: `a`,
     }: {

crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__positional_argument_only_leading_comma_comment.snap

@@ -24,7 +24,7 @@ expression: comments.debug(test_case.source_code)
         "trailing": [],
     },
     Node {
-        kind: Arg,
+        kind: ArgWithDefault,
         range: 15..16,
         source: `a`,
     }: {
@@ -39,7 +39,7 @@ expression: comments.debug(test_case.source_code)
         ],
     },
     Node {
-        kind: Arg,
+        kind: ArgWithDefault,
         range: 166..167,
         source: `b`,
     }: {

crates/ruff_python_formatter/src/comments/snapshots/ruff_python_formatter__comments__tests__trailing_after_comma.snap

@@ -4,7 +4,7 @@ expression: comments.debug(test_case.source_code)
 ---
 {
     Node {
-        kind: Arg,
+        kind: ArgWithDefault,
         range: 15..16,
         source: `a`,
     }: {

crates/ruff_python_formatter/src/comments/visitor.rs

@@ -236,6 +236,13 @@ impl<'ast> PreorderVisitor<'ast> for CommentsVisitor<'ast> {
         self.finish_node(arg);
     }
 
+    fn visit_arg_with_default(&mut self, arg_with_default: &'ast ArgWithDefault) {
+        if self.start_node(arg_with_default).is_traverse() {
+            walk_arg_with_default(self, arg_with_default);
+        }
+        self.finish_node(arg_with_default);
+    }
+
     fn visit_keyword(&mut self, keyword: &'ast Keyword) {
         if self.start_node(keyword).is_traverse() {
             walk_keyword(self, keyword);

crates/ruff_python_formatter/src/lib.rs

@@ -411,10 +411,12 @@ Formatted twice:
     #[test]
     fn quick_test() {
         let src = r#"
-def test(): ...
 
-# Comment
-def with_leading_comment(): ...
+def foo(
+    b=3
+    + 2  # comment
+):
+    ...
 "#;
         // Tokenize once
         let mut tokens = Vec::new();
@@ -437,10 +439,10 @@ def with_leading_comment(): ...
         // Use `dbg_write!(f, []) instead of `write!(f, [])` in your formatting code to print some IR
         // inside of a `Format` implementation
         // use ruff_formatter::FormatContext;
-        // formatted
+        // dbg!(formatted
         //     .document()
-        //     .display(formatted.context().source_code());
-
+        //     .display(formatted.context().source_code()));
+        //
         // dbg!(formatted
         //     .context()
         //     .comments()

crates/ruff_python_formatter/src/other/arg_with_default.rs

@@ -1,6 +1,5 @@
-use rustpython_parser::ast::ArgWithDefault;
-
 use ruff_formatter::write;
+use rustpython_parser::ast::ArgWithDefault;
 
 use crate::prelude::*;
 use crate::FormatNodeRule;
@@ -20,7 +19,7 @@ impl FormatNodeRule<ArgWithDefault> for FormatArgWithDefault {
 
         if let Some(default) = default {
             let space = def.annotation.is_some().then_some(space());
-            write!(f, [space, text("="), space, default.format()])?;
+            write!(f, [space, text("="), space, group(&default.format())])?;
         }
 
         Ok(())

crates/ruff_python_formatter/src/other/arguments.rs

@@ -34,29 +34,19 @@ impl FormatNodeRule<Arguments> for FormatArguments {
             let mut last_node: Option<AnyNodeRef> = None;
 
             for arg_with_default in posonlyargs {
-                joiner.entry(&arg_with_default.into_format());
-
-                last_node = Some(
-                    arg_with_default
-                        .default
-                        .as_deref()
-                        .map_or_else(|| (&arg_with_default.def).into(), AnyNodeRef::from),
-                );
+                joiner.entry(&arg_with_default.format());
+
+                last_node = Some(arg_with_default.into());
             }
 
             if !posonlyargs.is_empty() {
                 joiner.entry(&text("/"));
             }
 
             for arg_with_default in args {
-                joiner.entry(&arg_with_default.into_format());
-
-                last_node = Some(
-                    arg_with_default
-                        .default
-                        .as_deref()
-                        .map_or_else(|| (&arg_with_default.def).into(), AnyNodeRef::from),
-                );
+                joiner.entry(&arg_with_default.format());
+
+                last_node = Some(arg_with_default.into());
             }
 
             // kw only args need either a `*args` ahead of them capturing all var args or a `*`
@@ -74,14 +64,9 @@ impl FormatNodeRule<Arguments> for FormatArguments {
             }
 
             for arg_with_default in kwonlyargs {
-                joiner.entry(&arg_with_default.into_format());
-
-                last_node = Some(
-                    arg_with_default
-                        .default
-                        .as_deref()
-                        .map_or_else(|| (&arg_with_default.def).into(), AnyNodeRef::from),
-                );
+                joiner.entry(&arg_with_default.format());
+
+                last_node = Some(arg_with_default.into());
             }
 
             if let Some(kwarg) = kwarg {

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__ruff_test__statement__function_py.snap

@@ -96,6 +96,13 @@ else:
 # Regression test for https://github.com/python/cpython/blob/7199584ac8632eab57612f595a7162ab8d2ebbc0/Lib/warnings.py#L513
 def f(arg1=1, *, kwonlyarg1, kwonlyarg2=2):
     pass
+
+
+# Regression test for https://github.com/astral-sh/ruff/issues/5176#issuecomment-1598171989
+def foo(
+    b=3 + 2  # comment
+):
+    ...
 ```
 
 
@@ -223,6 +230,13 @@ else:
 # Regression test for https://github.com/python/cpython/blob/7199584ac8632eab57612f595a7162ab8d2ebbc0/Lib/warnings.py#L513
 def f(arg1=1, *, kwonlyarg1, kwonlyarg2=2):
     pass
+
+
+# Regression test for https://github.com/astral-sh/ruff/issues/5176#issuecomment-1598171989
+def foo(
+    b=3 + 2,  # comment
+):
+    ...
 ```