avoid lambda special-casing in maybe_parenthesize_expression

crates/ruff_python_formatter/src/expression/mod.rs

@@ -14,7 +14,6 @@ use ruff_text_size::Ranged;
 use crate::builders::parenthesize_if_expands;
 use crate::comments::{LeadingDanglingTrailingComments, leading_comments, trailing_comments};
 use crate::context::{NodeLevel, WithNodeLevel};
-use crate::expression::expr_lambda::ExprLambdaLayout;
 use crate::expression::parentheses::{
     NeedsParentheses, OptionalParentheses, Parentheses, Parenthesize, is_expression_parenthesized,
     optional_parentheses, parenthesized,
@@ -343,7 +342,6 @@ where
         expression,
         parent: parent.into(),
         parenthesize,
-        lambda_layout: ExprLambdaLayout::default(),
     }
 }
 
@@ -351,14 +349,6 @@ pub(crate) struct MaybeParenthesizeExpression<'a> {
     expression: &'a Expr,
     parent: AnyNodeRef<'a>,
     parenthesize: Parenthesize,
-    lambda_layout: ExprLambdaLayout,
-}
-
-impl MaybeParenthesizeExpression<'_> {
-    pub(crate) fn with_lambda_layout(mut self, layout: ExprLambdaLayout) -> Self {
-        self.lambda_layout = layout;
-        self
-    }
 }
 
 impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
@@ -367,7 +357,6 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
             expression,
             parent,
             parenthesize,
-            lambda_layout,
         } = self;
 
         let preserve_parentheses = parenthesize.is_optional()
@@ -421,19 +410,9 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
                 | Parenthesize::IfBreaksParenthesized
                 | Parenthesize::IfBreaksParenthesizedNested => {
                     if can_omit_optional_parentheses(expression, f.context()) {
-                        if let Expr::Lambda(lambda) = expression {
+                        optional_parentheses(&unparenthesized).fmt(f)
-                            optional_parentheses(&lambda.format().with_options(*lambda_layout))
-                                .fmt(f)
-                        } else {
-                            optional_parentheses(&unparenthesized).fmt(f)
-                        }
                     } else {
-                        if let Expr::Lambda(lambda) = expression {
+                        parenthesize_if_expands(&unparenthesized).fmt(f)
-                            parenthesize_if_expands(&lambda.format().with_options(*lambda_layout))
-                                .fmt(f)
-                        } else {
-                            parenthesize_if_expands(&unparenthesized).fmt(f)
-                        }
                     }
                 }
             },

crates/ruff_python_formatter/src/statement/stmt_assign.rs

@@ -304,13 +304,17 @@ impl Format<PyFormatContext<'_>> for FormatStatementsLastExpression<'_> {
                     && format_implicit_flat.is_none()
                     && format_interpolated_string.is_none()
                 {
-                    return maybe_parenthesize_expression(
+                    return if let Expr::Lambda(lambda) = value {
-                        value,
+                        let lambda = lambda.format().with_options(ExprLambdaLayout::Assignment);
-                        *statement,
+                        if can_omit_optional_parentheses(value, f.context()) {
-                        Parenthesize::IfBreaks,
+                            optional_parentheses(&lambda).fmt(f)
-                    )
+                        } else {
-                    .with_lambda_layout(ExprLambdaLayout::Assignment)
+                            parenthesize_if_expands(&lambda).fmt(f)
-                    .fmt(f);
+                        }
+                    } else {
+                        maybe_parenthesize_expression(value, *statement, Parenthesize::IfBreaks)
+                            .fmt(f)
+                    };
                 }
 
                 let comments = f.context().comments().clone();
@@ -581,20 +585,23 @@ impl Format<PyFormatContext<'_>> for FormatStatementsLastExpression<'_> {
                     && format_implicit_flat.is_none()
                     && format_interpolated_string.is_none()
                 {
+                    let formatted_value = format_with(|f| {
+                        if let Expr::Lambda(lambda) = value {
+                            let lambda = lambda.format().with_options(ExprLambdaLayout::Assignment);
+                            if can_omit_optional_parentheses(value, f.context()) {
+                                optional_parentheses(&lambda).fmt(f)
+                            } else {
+                                parenthesize_if_expands(&lambda).fmt(f)
+                            }
+                        } else {
+                            maybe_parenthesize_expression(value, *statement, Parenthesize::IfBreaks)
+                                .fmt(f)
+                        }
+                    });
+
                     return write!(
                         f,
-                        [
+                        [before_operator, space(), operator, space(), formatted_value]
-                            before_operator,
-                            space(),
-                            operator,
-                            space(),
-                            maybe_parenthesize_expression(
-                                value,
-                                *statement,
-                                Parenthesize::IfBreaks
-                            )
-                            .with_lambda_layout(ExprLambdaLayout::Assignment)
-                        ]
                     );
                 }