Format lambda expression (#5806)

crates/ruff_python_formatter/src/expression/expr_lambda.rs

@@ -1,22 +1,62 @@
+use crate::comments::dangling_comments;
 use crate::context::PyFormatContext;
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
-use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
+use crate::other::arguments::ArgumentsParentheses;
+use crate::AsFormat;
+use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::prelude::{space, text};
 use ruff_formatter::{write, Buffer, FormatResult};
-use ruff_python_ast::node::AnyNodeRef;
+use ruff_python_ast::node::{AnyNodeRef, AstNode};
 use rustpython_parser::ast::ExprLambda;
 
 #[derive(Default)]
 pub struct FormatExprLambda;
 
 impl FormatNodeRule<ExprLambda> for FormatExprLambda {
-    fn fmt_fields(&self, _item: &ExprLambda, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, item: &ExprLambda, f: &mut PyFormatter) -> FormatResult<()> {
+        let ExprLambda {
+            range: _,
+            args,
+            body,
+        } = item;
+
+        // It's possible for some `Arguments` of `lambda`s to be assigned dangling comments.
+        //
+        // a = (
+        //     lambda  # Dangling
+        //     : 1
+        // )
+        let comments = f.context().comments().clone();
+        let dangling = comments.dangling_comments(args.as_any_node_ref());
+
+        write!(f, [text("lambda")])?;
+
+        if !args.args.is_empty() {
+            write!(
+                f,
+                [
+                    space(),
+                    args.format()
+                        .with_options(ArgumentsParentheses::SkipInsideLambda),
+                ]
+            )?;
+        }
+
         write!(
             f,
-            [not_yet_implemented_custom_text(
-                "lambda NOT_YET_IMPLEMENTED_lambda: True"
-            )]
+            [
+                text(":"),
+                space(),
+                body.format(),
+                dangling_comments(dangling)
+            ]
         )
     }
+
+    fn fmt_dangling_comments(&self, _node: &ExprLambda, _f: &mut PyFormatter) -> FormatResult<()> {
+        // Override. Dangling comments are handled in `fmt_fields`.
+        Ok(())
+    }
 }
 
 impl NeedsParentheses for ExprLambda {

crates/ruff_python_formatter/src/other/arguments.rs

@@ -3,7 +3,7 @@ use std::usize;
 use ruff_text_size::{TextRange, TextSize};
 use rustpython_parser::ast::{Arguments, Ranged};
 
-use ruff_formatter::{format_args, write};
+use ruff_formatter::{format_args, write, FormatRuleWithOptions};
 use ruff_python_ast::node::{AnyNodeRef, AstNode};
 use ruff_python_whitespace::{first_non_trivia_token, SimpleTokenizer, Token, TokenKind};
 
@@ -16,8 +16,28 @@ use crate::expression::parentheses::parenthesized;
 use crate::prelude::*;
 use crate::FormatNodeRule;
 
+#[derive(Eq, PartialEq, Debug, Default)]
+pub enum ArgumentsParentheses {
+    #[default]
+    Default,
+
+    /// Arguments should never be inside parentheses for lambda expressions.
+    SkipInsideLambda,
+}
+
 #[derive(Default)]
-pub struct FormatArguments;
+pub struct FormatArguments {
+    parentheses: ArgumentsParentheses,
+}
+
+impl FormatRuleWithOptions<Arguments, PyFormatContext<'_>> for FormatArguments {
+    type Options = ArgumentsParentheses;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.parentheses = options;
+        self
+    }
+}
 
 impl FormatNodeRule<Arguments> for FormatArguments {
     fn fmt_fields(&self, item: &Arguments, f: &mut PyFormatter) -> FormatResult<()> {
@@ -170,7 +190,9 @@ impl FormatNodeRule<Arguments> for FormatArguments {
             + kwonlyargs.len()
             + usize::from(kwarg.is_some());
 
-        if num_arguments == 0 {
+        if self.parentheses == ArgumentsParentheses::SkipInsideLambda {
+            group(&format_inner).fmt(f)?;
+        } else if num_arguments == 0 {
             // No arguments, format any dangling comments between `()`
             write!(
                 f,