From 9eb73cb7e02435c27809dbfd004daf2b36af8e29 Mon Sep 17 00:00:00 2001 From: InSync Date: Mon, 23 Dec 2024 16:29:46 +0700 Subject: [PATCH] [`pycodestyle`] Preserve original value format (`E731`) (#15097) Co-authored-by: Micha Reiser --- .../test/fixtures/pycodestyle/E731.py | 21 ++++ .../pycodestyle/rules/lambda_assignment.rs | 74 +++++++----- ...les__pycodestyle__tests__E731_E731.py.snap | 106 +++++++++++++++++- 3 files changed, 170 insertions(+), 31 deletions(-) diff --git a/crates/ruff_linter/resources/test/fixtures/pycodestyle/E731.py b/crates/ruff_linter/resources/test/fixtures/pycodestyle/E731.py index a8acec8eec562..e28d613857469 100644 --- a/crates/ruff_linter/resources/test/fixtures/pycodestyle/E731.py +++ b/crates/ruff_linter/resources/test/fixtures/pycodestyle/E731.py @@ -156,3 +156,24 @@ def scope(): 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 +) \ No newline at end of file diff --git a/crates/ruff_linter/src/rules/pycodestyle/rules/lambda_assignment.rs b/crates/ruff_linter/src/rules/pycodestyle/rules/lambda_assignment.rs index fbe1967c027f5..725d930dd4e30 100644 --- a/crates/ruff_linter/src/rules/pycodestyle/rules/lambda_assignment.rs +++ b/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 { - parameters, body, .. - }) = value - else { + let Expr::Lambda(lambda) = 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( - id, - parameters.as_deref(), - body, - annotation, - checker.semantic(), - checker.generator(), - ) - .universal_newlines() - .enumerate() + for (idx, line) in function(id, lambda, annotation, checker) + .universal_newlines() + .enumerate() { if idx == 0 { indented.push_str(&line); @@ -186,19 +177,21 @@ fn extract_types(annotation: &Expr, semantic: &SemanticModel) -> Option<(Vec, - body: &Expr, + lambda: &ExprLambda, annotation: Option<&Expr>, - semantic: &SemanticModel, - generator: Generator, + checker: &Checker, ) -> 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 } diff --git a/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E731_E731.py.snap b/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E731_E731.py.snap index c12a0033a0e5d..fdc897617e3dd 100644 --- a/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E731_E731.py.snap +++ b/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), -150 149 | -151 150 | -152 151 | from dataclasses import dataclass + 148 |+ return ( + 149 |+ i := 1, + 150 |+ ) +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 | )