[pycodestyle] Preserve original value format (E731) (#15097)

Co-authored-by: Micha Reiser <micha@reiser.io>

crates/ruff_linter/resources/test/fixtures/pycodestyle/E731.py

@@ -156,3 +156,24 @@ from typing import Callable
 class FilterDataclass:
     # OK
     filter: Callable[[str], bool] = lambda _: True
+
+
+# Regression tests for:
+# * https://github.com/astral-sh/ruff/issues/7720
+x = lambda: """
+        a
+    b
+"""
+
+# * https://github.com/astral-sh/ruff/issues/10277
+at_least_one_million = lambda _: _ >= 1_000_000
+
+x = lambda: (
+    # comment
+    5 + 10
+)
+
+x = lambda: (
+    # comment
+    y := 10
+)

crates/ruff_linter/src/rules/pycodestyle/rules/lambda_assignment.rs

@@ -1,9 +1,10 @@
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
+use ruff_python_ast::parenthesize::parenthesized_range;
 use ruff_python_ast::{
-    self as ast, Expr, Identifier, Parameter, ParameterWithDefault, Parameters, Stmt,
+    self as ast, AstNode, Expr, ExprEllipsisLiteral, ExprLambda, Identifier, Parameter,
+    ParameterWithDefault, Parameters, Stmt,
 };
-use ruff_python_codegen::Generator;
 use ruff_python_semantic::SemanticModel;
 use ruff_python_trivia::{has_leading_content, has_trailing_content, leading_indentation};
 use ruff_source_file::UniversalNewlines;
@@ -65,10 +66,7 @@ pub(crate) fn lambda_assignment(
         return;
     };
 
-    let Expr::Lambda(ast::ExprLambda {
+    let Expr::Lambda(lambda) = value else {
-        parameters, body, ..
-    }) = value
-    else {
         return;
     };
 
@@ -85,16 +83,9 @@ pub(crate) fn lambda_assignment(
         let first_line = checker.locator().line_str(stmt.start());
         let indentation = leading_indentation(first_line);
         let mut indented = String::new();
-        for (idx, line) in function(
+        for (idx, line) in function(id, lambda, annotation, checker)
-            id,
+            .universal_newlines()
-            parameters.as_deref(),
+            .enumerate()
-            body,
-            annotation,
-            checker.semantic(),
-            checker.generator(),
-        )
-        .universal_newlines()
-        .enumerate()
         {
             if idx == 0 {
                 indented.push_str(&line);
@@ -186,19 +177,21 @@ fn extract_types(annotation: &Expr, semantic: &SemanticModel) -> Option<(Vec<Exp
 /// Generate a function definition from a `lambda` expression.
 fn function(
     name: &str,
-    parameters: Option<&Parameters>,
+    lambda: &ExprLambda,
-    body: &Expr,
     annotation: Option<&Expr>,
-    semantic: &SemanticModel,
+    checker: &Checker,
-    generator: Generator,
 ) -> String {
+    // Use a dummy body. It gets replaced at the end with the actual body.
+    // This allows preserving the source formatting for the body.
     let body = Stmt::Return(ast::StmtReturn {
-        value: Some(Box::new(body.clone())),
+        value: Some(Box::new(Expr::EllipsisLiteral(
+            ExprEllipsisLiteral::default(),
+        ))),
         range: TextRange::default(),
     });
-    let parameters = parameters.cloned().unwrap_or_default();
+    let parameters = lambda.parameters.as_deref().cloned().unwrap_or_default();
     if let Some(annotation) = annotation {
-        if let Some((arg_types, return_type)) = extract_types(annotation, semantic) {
+        if let Some((arg_types, return_type)) = extract_types(annotation, checker.semantic()) {
             // A `lambda` expression can only have positional-only and positional-or-keyword
             // arguments. The order is always positional-only first, then positional-or-keyword.
             let new_posonlyargs = parameters
@@ -243,10 +236,12 @@ fn function(
                 type_params: None,
                 range: TextRange::default(),
             });
-            return generator.stmt(&func);
+            let generated = checker.generator().stmt(&func);
+
+            return replace_trailing_ellipsis_with_original_expr(generated, lambda, checker);
         }
     }
-    let func = Stmt::FunctionDef(ast::StmtFunctionDef {
+    let function = Stmt::FunctionDef(ast::StmtFunctionDef {
         is_async: false,
         name: Identifier::new(name.to_string(), TextRange::default()),
         parameters: Box::new(parameters),
@@ -256,5 +251,32 @@ fn function(
         type_params: None,
         range: TextRange::default(),
     });
-    generator.stmt(&func)
+    let generated = checker.generator().stmt(&function);
+
+    replace_trailing_ellipsis_with_original_expr(generated, lambda, checker)
+}
+
+fn replace_trailing_ellipsis_with_original_expr(
+    mut generated: String,
+    lambda: &ExprLambda,
+    checker: &Checker,
+) -> String {
+    let original_expr_range = parenthesized_range(
+        (&lambda.body).into(),
+        lambda.as_any_node_ref(),
+        checker.comment_ranges(),
+        checker.source(),
+    )
+    .unwrap_or(lambda.body.range());
+
+    let original_expr_in_source = checker.locator().slice(original_expr_range);
+
+    let placeholder_ellipsis_start = generated.rfind("...").unwrap();
+    let placeholder_ellipsis_end = placeholder_ellipsis_start + "...".len();
+
+    generated.replace_range(
+        placeholder_ellipsis_start..placeholder_ellipsis_end,
+        original_expr_in_source,
+    );
+    generated
 }

crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E731_E731.py.snap

@@ -279,7 +279,7 @@ E731.py:127:5: E731 [*] Do not assign a `lambda` expression, use a `def`
 126 126 | 
 127     |-    f: Callable[[str, int], tuple[str, int]] = lambda a, b: (a, b)
     127 |+    def f(a: str, b: int) -> tuple[str, int]:
-    128 |+        return a, b
+    128 |+        return (a, b)
 128 129 | 
 129 130 | 
 130 131 | def scope():
@@ -364,7 +364,103 @@ E731.py:147:5: E731 [*] Do not assign a `lambda` expression, use a `def`
 148     |-        i := 1,
 149     |-    )
     147 |+    def f():
-    148 |+        return (i := 1),
+    148 |+        return (
-150 149 | 
+    149 |+            i := 1,
-151 150 | 
+    150 |+        )
-152 151 | from dataclasses import dataclass
+150 151 | 
+151 152 | 
+152 153 | from dataclasses import dataclass
+
+E731.py:163:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+    |
+161 |   # Regression tests for:
+162 |   # * https://github.com/astral-sh/ruff/issues/7720
+163 | / x = lambda: """
+164 | |         a
+165 | |     b
+166 | | """
+    | |___^ E731
+167 |   
+168 |   # * https://github.com/astral-sh/ruff/issues/10277
+    |
+    = help: Rewrite `x` as a `def`
+
+ℹ Unsafe fix
+160 160 | 
+161 161 | # Regression tests for:
+162 162 | # * https://github.com/astral-sh/ruff/issues/7720
+163     |-x = lambda: """
+    163 |+def x():
+    164 |+    return """
+164 165 |         a
+165 166 |     b
+166 167 | """
+
+E731.py:169:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+    |
+168 | # * https://github.com/astral-sh/ruff/issues/10277
+169 | at_least_one_million = lambda _: _ >= 1_000_000
+    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E731
+170 | 
+171 | x = lambda: (
+    |
+    = help: Rewrite `at_least_one_million` as a `def`
+
+ℹ Unsafe fix
+166 166 | """
+167 167 | 
+168 168 | # * https://github.com/astral-sh/ruff/issues/10277
+169     |-at_least_one_million = lambda _: _ >= 1_000_000
+    169 |+def at_least_one_million(_):
+    170 |+    return _ >= 1_000_000
+170 171 | 
+171 172 | x = lambda: (
+172 173 |     # comment
+
+E731.py:171:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+    |
+169 |   at_least_one_million = lambda _: _ >= 1_000_000
+170 |   
+171 | / x = lambda: (
+172 | |     # comment
+173 | |     5 + 10
+174 | | )
+    | |_^ E731
+175 |   
+176 |   x = lambda: (
+    |
+    = help: Rewrite `x` as a `def`
+
+ℹ Unsafe fix
+168 168 | # * https://github.com/astral-sh/ruff/issues/10277
+169 169 | at_least_one_million = lambda _: _ >= 1_000_000
+170 170 | 
+171     |-x = lambda: (
+    171 |+def x():
+    172 |+    return (
+172 173 |     # comment
+173 174 |     5 + 10
+174 175 | )
+
+E731.py:176:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+    |
+174 |   )
+175 |   
+176 | / x = lambda: (
+177 | |     # comment
+178 | |     y := 10
+179 | | )
+    | |_^ E731
+    |
+    = help: Rewrite `x` as a `def`
+
+ℹ Unsafe fix
+173 173 |     5 + 10
+174 174 | )
+175 175 | 
+176     |-x = lambda: (
+    176 |+def x():
+    177 |+    return (
+177 178 |     # comment
+178 179 |     y := 10
+179 180 | )