From 4dadee89adb99d9631fa3dd3036d9274d47d63e5 Mon Sep 17 00:00:00 2001
From: konstin <konstin@mailbox.org>
Date: Mon, 9 Oct 2023 15:39:05 +0200
Subject: [PATCH] Format comment before parameter default correctly

**Handle comment before the default values of function parameters correctly by inserting a line break instead of space after the equals sign where required.

```python
def f(
    a = # parameter trailing comment; needs line break (1)
    1,
    b =
    # default leading comment; needs line break (2)
    2,
    c = ( # the default leading can only be end-of-line with parentheses; no line break
        3
    ),
    d = (
        # own line leading comment with parentheses; no line break
        4
    )
)
```

Fixes #7603

**Test Plan** Added the different cases and one more complex case as fixtures.
---
 .../test/fixtures/ruff/statement/function.py  | 26 +++++++++
 .../src/other/parameter_with_default.rs       | 48 ++++++++++++++++-
 .../format@statement__function.py.snap        | 53 +++++++++++++++++++
 3 files changed, 126 insertions(+), 1 deletion(-)

diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py
index a036d614f5eeb..98eb09eefafcb 100644
--- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py
+++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py
@@ -384,3 +384,29 @@ def test():
     #comment
 except ImportError:
     pass
+
+# https://github.com/astral-sh/ruff/issues/7603
+def default_arg_comments(
+    a: str = #a
+    "a",
+    b: str =
+    #b
+    "b",
+    c: str =( #c
+        "c"
+    ),
+    d: str =(
+        #d
+        "d"
+    )
+):
+    print(a, b, c, d)
+
+def default_arg_comments2(#
+        x: int#=
+        = #
+        #
+        123#
+        #
+):
+    print(x)
diff --git a/crates/ruff_python_formatter/src/other/parameter_with_default.rs b/crates/ruff_python_formatter/src/other/parameter_with_default.rs
index 7afcb1b31a959..a434c7586f333 100644
--- a/crates/ruff_python_formatter/src/other/parameter_with_default.rs
+++ b/crates/ruff_python_formatter/src/other/parameter_with_default.rs
@@ -1,5 +1,7 @@
 use ruff_formatter::write;
 use ruff_python_ast::ParameterWithDefault;
+use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
+use ruff_text_size::{Ranged, TextRange};
 
 use crate::prelude::*;
 
@@ -18,7 +20,51 @@ impl FormatNodeRule<ParameterWithDefault> for FormatParameterWithDefault {
 
         if let Some(default) = default {
             let space = parameter.annotation.is_some().then_some(space());
-            write!(f, [space, token("="), space, group(&default.format())])?;
+            // ```python
+            // def f(
+            //     a = # parameter trailing comment; needs line break
+            //     1,
+            //     b =
+            //     # default leading comment; needs line break
+            //     2,
+            //     c = ( # the default leading can only be end-of-line with parentheses; no line break
+            //         3
+            //     ),
+            //     d = (
+            //         # own line leading comment with parentheses; no line break
+            //         4
+            //     )
+            // )
+            // ```
+            let needs_line_break_trailing = f.context().comments().has_trailing(parameter);
+            let default_first_comment = f.context().comments().leading(default.as_ref()).first();
+            let needs_line_break_leading = default_first_comment.is_some_and(|default_leading_comment| {
+                let mut tokenizer = SimpleTokenizer::new(
+                    f.context().source(),
+                    TextRange::new(parameter.end(), default_leading_comment.start()),
+                )
+                .skip_trivia()
+                .skip_while(|token| token.kind == SimpleTokenKind::RParen);
+                let equals = tokenizer.next();
+                debug_assert!(equals.is_some_and(|token| token.kind == SimpleTokenKind::Equals));
+                let lparens = tokenizer.next();
+                debug_assert!(lparens
+                    .as_ref()
+                    .map_or(true, |token| token.kind == SimpleTokenKind::LParen));
+                lparens.is_none()
+            });
+            let needs_line_break = needs_line_break_trailing || needs_line_break_leading;
+
+            write!(
+                f,
+                [
+                    space,
+                    token("="),
+                    (!needs_line_break).then_some(space),
+                    needs_line_break.then_some(hard_line_break()),
+                    group(&default.format())
+                ]
+            )?;
         }
 
         Ok(())
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__function.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__function.py.snap
index 443e451daf5f1..391948f2ed453 100644
--- a/crates/ruff_python_formatter/tests/snapshots/format@statement__function.py.snap
+++ b/crates/ruff_python_formatter/tests/snapshots/format@statement__function.py.snap
@@ -390,6 +390,32 @@ try:
     #comment
 except ImportError:
     pass
+
+# https://github.com/astral-sh/ruff/issues/7603
+def default_arg_comments(
+    a: str = #a
+    "a",
+    b: str =
+    #b
+    "b",
+    c: str =( #c
+        "c"
+    ),
+    d: str =(
+        #d
+        "d"
+    )
+):
+    print(a, b, c, d)
+
+def default_arg_comments2(#
+        x: int#=
+        = #
+        #
+        123#
+        #
+):
+    print(x)
 ```
 
 ## Output
@@ -940,6 +966,33 @@ try:
     # comment
 except ImportError:
     pass
+
+
+# https://github.com/astral-sh/ruff/issues/7603
+def default_arg_comments(
+    a: str =  # a
+    "a",
+    b: str =
+    # b
+    "b",
+    c: str = (  # c
+        "c"
+    ),
+    d: str = (
+        # d
+        "d"
+    ),
+):
+    print(a, b, c, d)
+
+
+def default_arg_comments2(  #
+    x: int =  # =  #
+    #
+    123,  #
+    #
+):
+    print(x)
 ```