Indent lambda parameters if parameters wrap

MichaReiser · MichaReiser · commit d5a18a697c21 · 2023-11-03T19:02:03.000+09:00
diff --git a/crates/ruff_python_ast/src/nodes.rs b/crates/ruff_python_ast/src/nodes.rs
@@ -2313,6 +2313,14 @@ impl Parameters {
             && self.vararg.is_none()
             && self.kwarg.is_none()
     }
+
+    pub fn len(&self) -> usize {
+        self.posonlyargs.len()
+            + self.args.len()
+            + usize::from(self.vararg.is_some())
+            + self.kwonlyargs.len()
+            + usize::from(self.kwarg.is_some())
+    }
 }
 
 /// An alternative type of AST `arg`. This is used for each function argument that might have a default value.
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/lambda.options.json b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/lambda.options.json
@@ -0,0 +1,8 @@
+[
+  {
+    "preview": "disabled"
+  },
+  {
+    "preview": "enabled"
+  }
+]
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/lambda.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/lambda.py
@@ -125,6 +125,13 @@ def f(
     *x: x
 )
 
+(
+    lambda
+    # comment
+    *x,
+    **y: x
+)
+
 (
     lambda
     # comment 1
@@ -135,13 +142,32 @@ def f(
     x
 )
 
+(
+    lambda
+    # comment 1
+    *
+    # comment 2
+    x,
+    **y:
+    # comment 3
+    x
+)
+
 (
     lambda # comment 1
     * # comment 2
     x: # comment 3
     x
 )
 
+(
+    lambda # comment 1
+        * # comment 2
+        x,
+        y: # comment 3
+    x
+)
+
 lambda *x\
     :x
 
@@ -196,6 +222,17 @@ def f(
     x
 )
 
+(
+    lambda  # 1
+    # 2
+    x,  # 3
+    # 4
+    y
+    :  # 5
+    # 6
+    x
+)
+
 (
     lambda
     x,
@@ -204,6 +241,71 @@ def f(
     z
 )
 
+
+# Leading
+lambda x: (
+    lambda y: lambda z: x
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + y
+                        + z  # Trailing
+)  # Trailing
+
+
+# Leading
+lambda x: lambda y: lambda z: [
+    x,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    y,
+    z
+] # Trailing
+# Trailing
+
 lambda self, araa, kkkwargs=aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa(*args, **kwargs), e=1, f=2, g=2: d
 
 # Regression tests for https://github.com/astral-sh/ruff/issues/8179
diff --git a/crates/ruff_python_formatter/src/context.rs b/crates/ruff_python_formatter/src/context.rs
@@ -43,6 +43,14 @@ impl<'a> PyFormatContext<'a> {
     pub(crate) fn comments(&self) -> &Comments<'a> {
         &self.comments
     }
+
+    pub(crate) const fn is_preview(&self) -> bool {
+        self.options.preview().is_enabled()
+    }
+
+    pub(crate) const fn is_stable(&self) -> bool {
+        !self.is_preview()
+    }
 }
 
 impl FormatContext for PyFormatContext<'_> {
diff --git a/crates/ruff_python_formatter/src/expression/expr_lambda.rs b/crates/ruff_python_formatter/src/expression/expr_lambda.rs
@@ -1,10 +1,11 @@
-use ruff_formatter::write;
+use ruff_formatter::{format_args, write};
 use ruff_python_ast::AnyNodeRef;
 use ruff_python_ast::ExprLambda;
 use ruff_text_size::Ranged;
 
-use crate::comments::{dangling_comments, SourceComment};
-use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
+use crate::comments::{dangling_comments, leading_comments, SourceComment};
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parenthesize};
+use crate::expression::{has_own_parentheses, maybe_parenthesize_expression};
 use crate::other::parameters::ParametersParentheses;
 use crate::prelude::*;
 
@@ -25,31 +26,49 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
         write!(f, [token("lambda")])?;
 
         if let Some(parameters) = parameters {
-            // In this context, a dangling comment can either be a comment between the `lambda` the
+            // In this context, a dangling comment can either be a comment between the `lambda` and the
             // parameters, or a comment between the parameters and the body.
             let (dangling_before_parameters, dangling_after_parameters) = dangling
                 .split_at(dangling.partition_point(|comment| comment.end() < parameters.start()));
 
             if dangling_before_parameters.is_empty() {
                 write!(f, [space()])?;
-            } else {
-                write!(f, [dangling_comments(dangling_before_parameters)])?;
             }
 
-            write!(
-                f,
-                [parameters
-                    .format()
-                    .with_options(ParametersParentheses::Never)]
-            )?;
+            group(&format_with(|f: &mut PyFormatter| {
+                if f.context().node_level().is_parenthesized()
+                    && (parameters.len() > 1 || !dangling_before_parameters.is_empty())
+                {
+                    let end_of_line_start = dangling_before_parameters
+                        .partition_point(|comment| comment.line_position().is_end_of_line());
+                    let (same_line_comments, own_line_comments) =
+                        dangling_before_parameters.split_at(end_of_line_start);
 
-            write!(f, [token(":")])?;
+                    dangling_comments(same_line_comments).fmt(f)?;
 
-            if dangling_after_parameters.is_empty() {
-                write!(f, [space()])?;
-            } else {
-                write!(f, [dangling_comments(dangling_after_parameters)])?;
-            }
+                    soft_block_indent(&format_args![
+                        leading_comments(own_line_comments),
+                        parameters
+                            .format()
+                            .with_options(ParametersParentheses::Never),
+                    ])
+                    .fmt(f)
+                } else {
+                    parameters
+                        .format()
+                        .with_options(ParametersParentheses::Never)
+                        .fmt(f)
+                }?;
+
+                token(":").fmt(f)?;
+
+                if dangling_after_parameters.is_empty() {
+                    space().fmt(f)
+                } else {
+                    dangling_comments(dangling_after_parameters).fmt(f)
+                }
+            }))
+            .fmt(f)?;
         } else {
             write!(f, [token(":")])?;
 
@@ -61,7 +80,12 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
             }
         }
 
-        write!(f, [body.format()])
+        // Avoid parenthesizing lists, dictionaries, etc.
+        if f.context().is_stable() || has_own_parentheses(body, f.context()).is_some() {
+            body.format().fmt(f)
+        } else {
+            maybe_parenthesize_expression(body, item, Parenthesize::IfBreaksOrIfRequired).fmt(f)
+        }
     }
 
     fn fmt_dangling_comments(
diff --git a/crates/ruff_python_formatter/src/expression/expr_named_expr.rs b/crates/ruff_python_formatter/src/expression/expr_named_expr.rs
@@ -69,6 +69,7 @@ impl NeedsParentheses for ExprNamedExpr {
             || parent.is_stmt_delete()
             || parent.is_stmt_for()
             || parent.is_stmt_function_def()
+            || parent.is_expr_lambda()
         {
             OptionalParentheses::Always
         } else {
diff --git a/crates/ruff_python_formatter/src/options.rs b/crates/ruff_python_formatter/src/options.rs
@@ -108,7 +108,7 @@ impl PyFormatOptions {
         self.line_ending
     }
 
-    pub fn preview(&self) -> PreviewMode {
+    pub const fn preview(&self) -> PreviewMode {
         self.preview
     }
 
diff --git a/crates/ruff_python_formatter/src/other/parameters.rs b/crates/ruff_python_formatter/src/other/parameters.rs
@@ -240,14 +240,10 @@ impl FormatNodeRule<Parameters> for FormatParameters {
             Ok(())
         });
 
-        let num_parameters = posonlyargs.len()
-            + args.len()
-            + usize::from(vararg.is_some())
-            + kwonlyargs.len()
-            + usize::from(kwarg.is_some());
+        let num_parameters = item.len();
 
         if self.parentheses == ParametersParentheses::Never {
-            write!(f, [group(&format_inner), dangling_comments(dangling)])
+            write!(f, [format_inner, dangling_comments(dangling)])
         } else if num_parameters == 0 {
             let mut f = WithNodeLevel::new(NodeLevel::ParenthesizedExpression, f);
             // No parameters, format any dangling comments between `()`
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__lambda.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__lambda.py.snap

Original file line number	Diff line number	Diff line change
`@@ -2313,6 +2313,14 @@ impl Parameters {`
`2313`	`2313`	`&& self.vararg.is_none()`
`2314`	`2314`	`&& self.kwarg.is_none()`
`2315`	`2315`	`}`
	`2316`	`+`
	`2317`	`+ pub fn len(&self) -> usize {`
	`2318`	`+ self.posonlyargs.len()`
	`2319`	`+ + self.args.len()`
	`2320`	`+ + usize::from(self.vararg.is_some())`
	`2321`	`+ + self.kwonlyargs.len()`
	`2322`	`+ + usize::from(self.kwarg.is_some())`
	`2323`	`+ }`
`2316`	`2324`	`}`
`2317`	`2325`
`2318`	`2326`	/// An alternative type of AST `arg`. This is used for each function argument that might have a default value.
-Original file line number
+Diff line change
@@ @@ -0,0 +1,8 @@ @@
 +[
 +  {
 +    "preview": "disabled"
 +  },
 +  {
 +    "preview": "enabled"
 +  }
 +]
Original file line number	Diff line number	Diff line change
`@@ -43,6 +43,14 @@ impl<'a> PyFormatContext<'a> {`
`43`	`43`	`pub(crate) fn comments(&self) -> &Comments<'a> {`
`44`	`44`	`&self.comments`
`45`	`45`	`}`
	`46`	`+`
	`47`	`+ pub(crate) const fn is_preview(&self) -> bool {`
	`48`	`+ self.options.preview().is_enabled()`
	`49`	`+ }`
	`50`	`+`
	`51`	`+ pub(crate) const fn is_stable(&self) -> bool {`
	`52`	`+ !self.is_preview()`
	`53`	`+ }`
`46`	`54`	`}`
`47`	`55`
`48`	`56`	`impl FormatContext for PyFormatContext<'_> {`
Original file line number	Diff line number	Diff line change
`@@ -69,6 +69,7 @@ impl NeedsParentheses for ExprNamedExpr {`
`69`	`69`	`\|\| parent.is_stmt_delete()`
`70`	`70`	`\|\| parent.is_stmt_for()`
`71`	`71`	`\|\| parent.is_stmt_function_def()`
	`72`	`+ \|\| parent.is_expr_lambda()`
`72`	`73`	`{`
`73`	`74`	`OptionalParentheses::Always`
`74`	`75`	`} else {`
Original file line number	Diff line number	Diff line change
`@@ -108,7 +108,7 @@ impl PyFormatOptions {`
`108`	`108`	`self.line_ending`
`109`	`109`	`}`
`110`	`110`
`111`		`- pub fn preview(&self) -> PreviewMode {`
	`111`	`+ pub const fn preview(&self) -> PreviewMode {`
`112`	`112`	`self.preview`
`113`	`113`	`}`
`114`	`114`