move all body formatting to FormatBody

crates/ruff_python_formatter/src/expression/expr_lambda.rs

@@ -147,13 +147,95 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
             dangling
         };
 
-        if preview {
+        FormatBody {
+            body,
+            dangling,
+            layout: self.layout,
+        }
+        .fmt(f)
+    }
+}
+
+#[derive(Debug, Default, Copy, Clone)]
+pub enum ExprLambdaLayout {
+    #[default]
+    Default,
+
+    /// The [`ExprLambda`] is the direct child of an assignment expression, so it needs to use
+    /// `fits_expanded` to prefer parenthesizing its own body before the assignment tries to
+    /// parenthesize the whole lambda. For example, we want this formatting:
+    ///
+    /// ```py
+    /// long_assignment_target = lambda x, y, z: (
+    ///     x + y + z
+    /// )
+    /// ```
+    ///
+    /// instead of either of these:
+    ///
+    /// ```py
+    /// long_assignment_target = (
+    ///     lambda x, y, z: (
+    ///         x + y + z
+    ///     )
+    /// )
+    ///
+    /// long_assignment_target = (
+    ///     lambda x, y, z: x + y + z
+    /// )
+    /// ```
+    Assignment,
+}
+
+impl FormatRuleWithOptions<ExprLambda, PyFormatContext<'_>> for FormatExprLambda {
+    type Options = ExprLambdaLayout;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.layout = options;
+        self
+    }
+}
+
+impl NeedsParentheses for ExprLambda {
+    fn needs_parentheses(
+        &self,
+        parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        if parent.is_expr_await() {
+            OptionalParentheses::Always
+        } else {
+            OptionalParentheses::Multiline
+        }
+    }
+}
+
+struct FormatBody<'a> {
+    body: &'a Expr,
+    dangling: &'a [SourceComment],
+    layout: ExprLambdaLayout,
+}
+
+impl Format<PyFormatContext<'_>> for FormatBody<'_> {
+    fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
+        let FormatBody {
+            dangling,
+            body,
+            layout,
+        } = self;
+
+        if !is_parenthesize_lambda_bodies_enabled(f.context()) {
+            return body.format().fmt(f);
+        }
+
+        let body = *body;
+        let comments = f.context().comments().clone();
         let body_comments = comments.leading_dangling_trailing(body);
+
         let fmt_body = format_with(|f: &mut PyFormatter| {
             if !dangling.is_empty() {
-                    // Can't use partition_point because there can be additional end of line
+                // Can't use partition_point because there can be additional end of line comments
-                    // comments after the initial set. All of these comments are dangling, for
+                // after the initial set. All of these comments are dangling, for example:
-                    // example:
                 //
                 // ```python
                 // (
@@ -200,11 +282,8 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
                 // )
                 // ```
                 let comments = f.context().comments();
-                    if is_expression_parenthesized(
+                if is_expression_parenthesized(body.into(), comments.ranges(), f.context().source())
-                        body.into(),
+                    && comments.has_leading(body)
-                        comments.ranges(),
-                        f.context().source(),
-                    ) && comments.has_leading(body)
                 {
                     trailing_comments(dangling).fmt(f)?;
 
@@ -344,78 +423,9 @@ impl FormatNodeRule<ExprLambda> for FormatExprLambda {
             }
         });
 
-            match self.layout {
+        match layout {
             ExprLambdaLayout::Assignment => fits_expanded(&fmt_body).fmt(f),
             ExprLambdaLayout::Default => fmt_body.fmt(f),
         }
-        } else {
-            body.format().fmt(f)
-        }
-    }
-}
-
-#[derive(Debug, Default, Copy, Clone)]
-pub enum ExprLambdaLayout {
-    #[default]
-    Default,
-
-    /// The [`ExprLambda`] is the direct child of an assignment expression, so it needs to use
-    /// `fits_expanded` to prefer parenthesizing its own body before the assignment tries to
-    /// parenthesize the whole lambda. For example, we want this formatting:
-    ///
-    /// ```py
-    /// long_assignment_target = lambda x, y, z: (
-    ///     x + y + z
-    /// )
-    /// ```
-    ///
-    /// instead of either of these:
-    ///
-    /// ```py
-    /// long_assignment_target = (
-    ///     lambda x, y, z: (
-    ///         x + y + z
-    ///     )
-    /// )
-    ///
-    /// long_assignment_target = (
-    ///     lambda x, y, z: x + y + z
-    /// )
-    /// ```
-    Assignment,
-}
-
-impl FormatRuleWithOptions<ExprLambda, PyFormatContext<'_>> for FormatExprLambda {
-    type Options = ExprLambdaLayout;
-
-    fn with_options(mut self, options: Self::Options) -> Self {
-        self.layout = options;
-        self
-    }
-}
-
-impl NeedsParentheses for ExprLambda {
-    fn needs_parentheses(
-        &self,
-        parent: AnyNodeRef,
-        _context: &PyFormatContext,
-    ) -> OptionalParentheses {
-        if parent.is_expr_await() {
-            OptionalParentheses::Always
-        } else {
-            OptionalParentheses::Multiline
-        }
-    }
-}
-
-struct FormatBody<'a> {
-    body: &'a Expr,
-    dangling: &'a [SourceComment],
-}
-
-impl Format<PyFormatContext<'_>> for FormatBody<'_> {
-    fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
-        let FormatBody { dangling, body } = self;
-        todo!()
     }
 }