Fix ArgWithDefault comments handling

astral-sh · Jun 20, 2023 · 7657211 · 7657211
1 parent 4cc3cdb
commit 7657211
Show file tree

Hide file tree

Showing 15 changed files with 54 additions and 65 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -36,11 +36,11 @@ proc-macro2 = { version = "1.0.51" }
 quote = { version = "1.0.23" }
 regex = { version = "1.7.1" }
 rustc-hash = { version = "1.1.0" }
-ruff_text_size = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "8d74eee75031b68d2204219963fae54a3f31a394" }
-rustpython-ast = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "8d74eee75031b68d2204219963fae54a3f31a394" , default-features = false, features = ["all-nodes-with-ranges", "num-bigint"]}
-rustpython-format = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "8d74eee75031b68d2204219963fae54a3f31a394", default-features = false, features = ["num-bigint"] }
-rustpython-literal = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "8d74eee75031b68d2204219963fae54a3f31a394" }
-rustpython-parser = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "8d74eee75031b68d2204219963fae54a3f31a394" , default-features = false, features = ["full-lexer", "all-nodes-with-ranges", "num-bigint"] }
+ruff_text_size = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "db0735c95186c5daf5dfa626b724687a3a2285de" }
+rustpython-ast = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "db0735c95186c5daf5dfa626b724687a3a2285de" , default-features = false, features = ["all-nodes-with-ranges", "num-bigint"]}
+rustpython-format = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "db0735c95186c5daf5dfa626b724687a3a2285de", default-features = false, features = ["num-bigint"] }
+rustpython-literal = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "db0735c95186c5daf5dfa626b724687a3a2285de" }
+rustpython-parser = { git = "https://github.com/astral-sh/RustPython-Parser.git", rev = "db0735c95186c5daf5dfa626b724687a3a2285de" , default-features = false, features = ["full-lexer", "all-nodes-with-ranges", "num-bigint"] }
 schemars = { version = "0.8.12" }
 serde = { version = "1.0.152", features = ["derive"] }
 serde_json = { version = "1.0.93" }

diff --git a/crates/ruff_python_ast/src/node.rs b/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`].

diff --git a/crates/ruff_python_formatter/src/cli.rs b/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))
         );
     }

diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs
@@ -622,13 +622,8 @@ fn handle_positional_only_arguments_separator_comment<'a>(
         return CommentPlacement::Default(comment);
     };
 
-    let is_last_positional_argument = are_same_optional(last_argument_or_default, arguments.posonlyargs.last())
-            // If the preceding node is the default for the last positional argument
-            // ```python
-            // def test(a=10, /, b): pass
-            // ```
-            || are_same_optional(last_argument_or_default, arguments
-                .posonlyargs.last().and_then(|arg| arg.default.as_deref()));
+    let is_last_positional_argument =
+        are_same_optional(last_argument_or_default, arguments.posonlyargs.last());
 
     if !is_last_positional_argument {
         return CommentPlacement::Default(comment);

diff --git a/.../ruff_python_formatter__comments__tests__non_positional_arguments_slash_on_same_line.snap b/.../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 {

diff --git a/...shots/ruff_python_formatter__comments__tests__non_positional_arguments_with_defaults.snap b/...shots/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 {

diff --git a/...s/snapshots/ruff_python_formatter__comments__tests__positional_argument_only_comment.snap b/...s/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`,
     }: {

diff --git a/..._formatter__comments__tests__positional_argument_only_comment_without_following_node.snap b/..._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`,
     }: {

diff --git a/...ff_python_formatter__comments__tests__positional_argument_only_leading_comma_comment.snap b/...ff_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`,
     }: {

diff --git a/.../src/comments/snapshots/ruff_python_formatter__comments__tests__trailing_after_comma.snap b/.../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`,
     }: {

diff --git a/crates/ruff_python_formatter/src/comments/visitor.rs b/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);

diff --git a/crates/ruff_python_formatter/src/lib.rs b/crates/ruff_python_formatter/src/lib.rs
@@ -412,10 +412,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,11 @@ def with_leading_comment(): ...
         // Uncomment the `dbg` to print the IR.
         // 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;
         // dbg!(formatted
         //     .document()
         //     .display(formatted.context().source_code()));
-
+        //
         // dbg!(formatted
         //     .context()
         //     .comments()

diff --git a/crates/ruff_python_formatter/src/other/arg_with_default.rs b/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;

diff --git a/crates/ruff_python_formatter/src/other/arguments.rs b/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 {