From 6e9e42d343515e3a0966095ff0cf7cb419bd3cd6 Mon Sep 17 00:00:00 2001
From: Brent Westbrook <brentrwestbrook@gmail.com>
Date: Wed, 10 Dec 2025 09:04:00 -0500
Subject: [PATCH] factor out maybe_parenthesize_lambda

---
 .../src/statement/stmt_assign.rs              | 65 +++++++++++--------
 1 file changed, 37 insertions(+), 28 deletions(-)

diff --git a/crates/ruff_python_formatter/src/statement/stmt_assign.rs b/crates/ruff_python_formatter/src/statement/stmt_assign.rs
index 10fe3a6c13..b695e9ad8b 100644
--- a/crates/ruff_python_formatter/src/statement/stmt_assign.rs
+++ b/crates/ruff_python_formatter/src/statement/stmt_assign.rs
@@ -305,18 +305,7 @@ impl Format<PyFormatContext<'_>> for FormatStatementsLastExpression<'_> {
                     && format_implicit_flat.is_none()
                     && format_interpolated_string.is_none()
                 {
-                    return if is_parenthesize_lambda_bodies_enabled(f.context())
-                        && let Expr::Lambda(lambda) = value
-                        && !f.context().comments().has_leading(lambda)
-                    {
-                        parenthesize_if_expands(
-                            &lambda.format().with_options(ExprLambdaLayout::Assignment),
-                        )
-                        .fmt(f)
-                    } else {
-                        maybe_parenthesize_expression(value, *statement, Parenthesize::IfBreaks)
-                            .fmt(f)
-                    };
+                    return maybe_parenthesize_lambda(value, *statement).fmt(f);
                 }
 
                 let comments = f.context().comments().clone();
@@ -587,24 +576,15 @@ impl Format<PyFormatContext<'_>> for FormatStatementsLastExpression<'_> {
                     && format_implicit_flat.is_none()
                     && format_interpolated_string.is_none()
                 {
-                    let formatted_value = format_with(|f| {
-                        if is_parenthesize_lambda_bodies_enabled(f.context())
-                            && let Expr::Lambda(lambda) = value
-                            && !f.context().comments().has_leading(lambda)
-                        {
-                            parenthesize_if_expands(
-                                &lambda.format().with_options(ExprLambdaLayout::Assignment),
-                            )
-                            .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_lambda(value, *statement)
+                        ]
                     );
                 }
 
@@ -1383,3 +1363,32 @@ fn is_attribute_with_parenthesized_value(target: &Expr, context: &PyFormatContex
         _ => false,
     }
 }
+
+/// Like [`maybe_parenthesize_expression`] but with special handling for lambdas in preview.
+fn maybe_parenthesize_lambda<'a>(
+    expression: &'a Expr,
+    parent: AnyNodeRef<'a>,
+) -> MaybeParenthesizeLambda<'a> {
+    MaybeParenthesizeLambda { expression, parent }
+}
+
+struct MaybeParenthesizeLambda<'a> {
+    expression: &'a Expr,
+    parent: AnyNodeRef<'a>,
+}
+
+impl Format<PyFormatContext<'_>> for MaybeParenthesizeLambda<'_> {
+    fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
+        let MaybeParenthesizeLambda { expression, parent } = self;
+
+        if is_parenthesize_lambda_bodies_enabled(f.context())
+            && let Expr::Lambda(lambda) = expression
+            && !f.context().comments().has_leading(lambda)
+        {
+            parenthesize_if_expands(&lambda.format().with_options(ExprLambdaLayout::Assignment))
+                .fmt(f)
+        } else {
+            maybe_parenthesize_expression(expression, *parent, Parenthesize::IfBreaks).fmt(f)
+        }
+    }
+}