From 90244a0e5433123dd5a0c7c31ae03b20ab86cdcc Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 6 Sep 2023 13:36:15 +0100 Subject: [PATCH] Parenthesize C417 targets if necessary --- .../fixtures/flake8_comprehensions/C417.py | 5 + .../src/rules/flake8_comprehensions/fixes.rs | 229 +++++++++--------- ...8_comprehensions__tests__C417_C417.py.snap | 53 ++++ 3 files changed, 173 insertions(+), 114 deletions(-) diff --git a/crates/ruff/resources/test/fixtures/flake8_comprehensions/C417.py b/crates/ruff/resources/test/fixtures/flake8_comprehensions/C417.py index 0d7565336b0ac..92926c7ca1bfd 100644 --- a/crates/ruff/resources/test/fixtures/flake8_comprehensions/C417.py +++ b/crates/ruff/resources/test/fixtures/flake8_comprehensions/C417.py @@ -41,3 +41,8 @@ def func(arg1: int, arg2: int = 4): # Ok because multiple arguments are allowed. dict(map(lambda k, v: (k, v), keys, values)) + +# Regression test for: https://github.com/astral-sh/ruff/issues/7121 +map(lambda x: x, y if y else z) +map(lambda x: x, (y if y else z)) +map(lambda x: x, (x, y, z)) diff --git a/crates/ruff/src/rules/flake8_comprehensions/fixes.rs b/crates/ruff/src/rules/flake8_comprehensions/fixes.rs index b0d46e81e513c..13ee009216519 100644 --- a/crates/ruff/src/rules/flake8_comprehensions/fixes.rs +++ b/crates/ruff/src/rules/flake8_comprehensions/fixes.rs @@ -3,8 +3,8 @@ use itertools::Itertools; use libcst_native::{ Arg, AssignEqual, AssignTargetExpression, Call, Comment, CompFor, Dict, DictComp, DictElement, Element, EmptyLine, Expression, GeneratorExp, LeftCurlyBrace, LeftParen, LeftSquareBracket, - List, ListComp, Name, ParenthesizableWhitespace, ParenthesizedWhitespace, RightCurlyBrace, - RightParen, RightSquareBracket, Set, SetComp, SimpleString, SimpleWhitespace, + List, ListComp, Name, ParenthesizableWhitespace, ParenthesizedNode, ParenthesizedWhitespace, + RightCurlyBrace, RightParen, RightSquareBracket, Set, SetComp, SimpleString, SimpleWhitespace, TrailingWhitespace, Tuple, }; use ruff_python_ast::Expr; @@ -938,134 +938,135 @@ pub(crate) fn fix_unnecessary_map( let module_text = locator.slice(expr); let mut tree = match_expression(module_text)?; let call = match_call_mut(&mut tree)?; - let arg = match_arg(call)?; - let (args, lambda_func) = match &arg.value { - Expression::Call(outer_call) => { - let inner_lambda = outer_call.args.first().unwrap().value.clone(); - match &inner_lambda { - Expression::Lambda(..) => (outer_call.args.clone(), inner_lambda), - _ => { - bail!("Expected a lambda function") - } - } + let (lambda, iter) = match call.args.as_slice() { + [call] => { + let call = match_call(&call.value)?; + let [lambda, iter] = call.args.as_slice() else { + bail!("Expected two arguments"); + }; + let lambda = match_lambda(&lambda.value)?; + let iter = &iter.value; + (lambda, iter) } - Expression::Lambda(..) => (call.args.clone(), arg.value.clone()), - _ => { - bail!("Expected a lambda or call") + [lambda, iter] => { + let lambda = match_lambda(&lambda.value)?; + let iter = &iter.value; + (lambda, iter) } + _ => bail!("Expected a call or lambda"), }; - let func_body = match_lambda(&lambda_func)?; + // Format the lambda arguments as a comma-separated list. + let mut args_str = lambda.params.params.iter().map(|f| f.name.value).join(", "); + if args_str.is_empty() { + args_str = "_".to_string(); + } - if args.len() == 2 { - if func_body.params.params.iter().any(|f| f.default.is_some()) { - bail!("Currently not supporting default values"); - } + // Parenthesize the iterator, if necessary, as in: + // ```python + // map(lambda x: x, y if y else z) + // ``` + let iter = iter.clone(); + let iter = if iter.lpar().is_empty() + && iter.rpar().is_empty() + && matches!(iter, Expression::IfExp(_) | Expression::Lambda(_)) + { + iter.with_parens(LeftParen::default(), RightParen::default()) + } else { + iter + }; - let mut args_str = func_body - .params - .params - .iter() - .map(|f| f.name.value) - .join(", "); - if args_str.is_empty() { - args_str = "_".to_string(); + let compfor = Box::new(CompFor { + target: AssignTargetExpression::Name(Box::new(Name { + value: args_str.as_str(), + lpar: vec![], + rpar: vec![], + })), + iter, + ifs: vec![], + inner_for_in: None, + asynchronous: None, + whitespace_before: space(), + whitespace_after_for: space(), + whitespace_before_in: space(), + whitespace_after_in: space(), + }); + + match object_type { + ObjectType::Generator => { + tree = Expression::GeneratorExp(Box::new(GeneratorExp { + elt: lambda.body.clone(), + for_in: compfor, + lpar: vec![LeftParen::default()], + rpar: vec![RightParen::default()], + })); } - - let compfor = Box::new(CompFor { - target: AssignTargetExpression::Name(Box::new(Name { - value: args_str.as_str(), + ObjectType::List => { + tree = Expression::ListComp(Box::new(ListComp { + elt: lambda.body.clone(), + for_in: compfor, + lbracket: LeftSquareBracket::default(), + rbracket: RightSquareBracket::default(), lpar: vec![], rpar: vec![], - })), - iter: args.last().unwrap().value.clone(), - ifs: vec![], - inner_for_in: None, - asynchronous: None, - whitespace_before: space(), - whitespace_after_for: space(), - whitespace_before_in: space(), - whitespace_after_in: space(), - }); - - match object_type { - ObjectType::Generator => { - tree = Expression::GeneratorExp(Box::new(GeneratorExp { - elt: func_body.body.clone(), - for_in: compfor, - lpar: vec![LeftParen::default()], - rpar: vec![RightParen::default()], - })); - } - ObjectType::List => { - tree = Expression::ListComp(Box::new(ListComp { - elt: func_body.body.clone(), - for_in: compfor, - lbracket: LeftSquareBracket::default(), - rbracket: RightSquareBracket::default(), - lpar: vec![], - rpar: vec![], - })); - } - ObjectType::Set => { - tree = Expression::SetComp(Box::new(SetComp { - elt: func_body.body.clone(), - for_in: compfor, - lpar: vec![], - rpar: vec![], - lbrace: LeftCurlyBrace::default(), - rbrace: RightCurlyBrace::default(), - })); - } - ObjectType::Dict => { - let elements = match func_body.body.as_ref() { - Expression::Tuple(tuple) => &tuple.elements, - Expression::List(list) => &list.elements, - _ => { - bail!("Expected tuple or list for dictionary comprehension") - } - }; - let [key, value] = elements.as_slice() else { - bail!("Expected container to include two elements"); - }; - let Element::Simple { value: key, .. } = key else { - bail!("Expected container to use a key as the first element"); - }; - let Element::Simple { value, .. } = value else { - bail!("Expected container to use a value as the second element"); - }; + })); + } + ObjectType::Set => { + tree = Expression::SetComp(Box::new(SetComp { + elt: lambda.body.clone(), + for_in: compfor, + lpar: vec![], + rpar: vec![], + lbrace: LeftCurlyBrace::default(), + rbrace: RightCurlyBrace::default(), + })); + } + ObjectType::Dict => { + let elements = match lambda.body.as_ref() { + Expression::Tuple(tuple) => &tuple.elements, + Expression::List(list) => &list.elements, + _ => { + bail!("Expected tuple or list for dictionary comprehension") + } + }; + let [key, value] = elements.as_slice() else { + bail!("Expected container to include two elements"); + }; + let Element::Simple { value: key, .. } = key else { + bail!("Expected container to use a key as the first element"); + }; + let Element::Simple { value, .. } = value else { + bail!("Expected container to use a value as the second element"); + }; - tree = Expression::DictComp(Box::new(DictComp { - for_in: compfor, - lpar: vec![], - rpar: vec![], - key: Box::new(key.clone()), - value: Box::new(value.clone()), - lbrace: LeftCurlyBrace::default(), - rbrace: RightCurlyBrace::default(), - whitespace_before_colon: ParenthesizableWhitespace::default(), - whitespace_after_colon: ParenthesizableWhitespace::SimpleWhitespace( - SimpleWhitespace(" "), - ), - })); - } + tree = Expression::DictComp(Box::new(DictComp { + for_in: compfor, + lpar: vec![], + rpar: vec![], + key: Box::new(key.clone()), + value: Box::new(value.clone()), + lbrace: LeftCurlyBrace::default(), + rbrace: RightCurlyBrace::default(), + whitespace_before_colon: ParenthesizableWhitespace::default(), + whitespace_after_colon: ParenthesizableWhitespace::SimpleWhitespace( + SimpleWhitespace(" "), + ), + })); } + } - let mut content = tree.codegen_stylist(stylist); + let mut content = tree.codegen_stylist(stylist); - // If the expression is embedded in an f-string, surround it with spaces to avoid - // syntax errors. - if matches!(object_type, ObjectType::Set | ObjectType::Dict) { - if parent.is_some_and(Expr::is_formatted_value_expr) { - content = format!(" {content} "); - } + // If the expression is embedded in an f-string, surround it with spaces to avoid + // syntax errors. + if matches!(object_type, ObjectType::Set | ObjectType::Dict) { + if parent.is_some_and(Expr::is_formatted_value_expr) { + content = format!(" {content} "); } - - Ok(Edit::range_replacement(content, expr.range())) - } else { - bail!("Should have two arguments"); } + + Ok(Edit::range_replacement(content, expr.range())) } /// (C418) Convert `dict({"a": 1})` to `{"a": 1}` diff --git a/crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C417_C417.py.snap b/crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C417_C417.py.snap index 5152e13434eae..da480658a0b34 100644 --- a/crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C417_C417.py.snap +++ b/crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C417_C417.py.snap @@ -290,4 +290,57 @@ C417.py:35:1: C417 [*] Unnecessary `map` usage (rewrite using a generator expres 37 37 | # Ok because of the default parameters, and variadic arguments. 38 38 | map(lambda x=1: x, nums) +C417.py:46:1: C417 [*] Unnecessary `map` usage (rewrite using a generator expression) + | +45 | # Regression test for: https://github.com/astral-sh/ruff/issues/7121 +46 | map(lambda x: x, y if y else z) + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417 +47 | map(lambda x: x, (y if y else z)) +48 | map(lambda x: x, (x, y, z)) + | + = help: Replace `map` with a generator expression + +ℹ Suggested fix +43 43 | dict(map(lambda k, v: (k, v), keys, values)) +44 44 | +45 45 | # Regression test for: https://github.com/astral-sh/ruff/issues/7121 +46 |-map(lambda x: x, y if y else z) + 46 |+(x for x in (y if y else z)) +47 47 | map(lambda x: x, (y if y else z)) +48 48 | map(lambda x: x, (x, y, z)) + +C417.py:47:1: C417 [*] Unnecessary `map` usage (rewrite using a generator expression) + | +45 | # Regression test for: https://github.com/astral-sh/ruff/issues/7121 +46 | map(lambda x: x, y if y else z) +47 | map(lambda x: x, (y if y else z)) + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417 +48 | map(lambda x: x, (x, y, z)) + | + = help: Replace `map` with a generator expression + +ℹ Suggested fix +44 44 | +45 45 | # Regression test for: https://github.com/astral-sh/ruff/issues/7121 +46 46 | map(lambda x: x, y if y else z) +47 |-map(lambda x: x, (y if y else z)) + 47 |+(x for x in (y if y else z)) +48 48 | map(lambda x: x, (x, y, z)) + +C417.py:48:1: C417 [*] Unnecessary `map` usage (rewrite using a generator expression) + | +46 | map(lambda x: x, y if y else z) +47 | map(lambda x: x, (y if y else z)) +48 | map(lambda x: x, (x, y, z)) + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^ C417 + | + = help: Replace `map` with a generator expression + +ℹ Suggested fix +45 45 | # Regression test for: https://github.com/astral-sh/ruff/issues/7121 +46 46 | map(lambda x: x, y if y else z) +47 47 | map(lambda x: x, (y if y else z)) +48 |-map(lambda x: x, (x, y, z)) + 48 |+(x for x in (x, y, z)) +