diff --git a/resources/test/fixtures/flake8_simplify/SIM401.py b/resources/test/fixtures/flake8_simplify/SIM401.py index c18c6a3d43..b3ae3c2fce 100644 --- a/resources/test/fixtures/flake8_simplify/SIM401.py +++ b/resources/test/fixtures/flake8_simplify/SIM401.py @@ -79,3 +79,9 @@ if key in a_dict: else: var2 = value2 var = "default" + +# OK (complex default value) +if key in a_dict: + var = a_dict[key] +else: + var = foo() diff --git a/src/flake8_simplify/rules/ast_if.rs b/src/flake8_simplify/rules/ast_if.rs index a183b6fba2..08aa58e7a5 100644 --- a/src/flake8_simplify/rules/ast_if.rs +++ b/src/flake8_simplify/rules/ast_if.rs @@ -2,7 +2,8 @@ use rustpython_ast::{Cmpop, Constant, Expr, ExprContext, ExprKind, Stmt, StmtKin use crate::ast::comparable::ComparableExpr; use crate::ast::helpers::{ - contains_call_path, create_expr, create_stmt, has_comments, unparse_expr, unparse_stmt, + any_over_expr, contains_call_path, create_expr, create_stmt, has_comments, unparse_expr, + unparse_stmt, }; use crate::ast::types::Range; use crate::checkers::ast::Checker; @@ -257,7 +258,7 @@ pub fn use_dict_get_with_default( if orelse_var.len() != 1 { return; }; - let ExprKind::Compare { left: test_key, ops , comparators: test_dict } = &test.node else { + let ExprKind::Compare { left: test_key, ops , comparators: test_dict } = &test.node else { return; }; if test_dict.len() != 1 { @@ -284,6 +285,23 @@ pub fn use_dict_get_with_default( return; } + // Check that the default value is not "complex". + if any_over_expr(default_val, &|expr| { + matches!( + expr.node, + ExprKind::Call { .. } + | ExprKind::Await { .. } + | ExprKind::GeneratorExp { .. } + | ExprKind::ListComp { .. } + | ExprKind::SetComp { .. } + | ExprKind::DictComp { .. } + | ExprKind::Yield { .. } + | ExprKind::YieldFrom { .. } + ) + }) { + return; + } + let contents = unparse_stmt( &create_stmt(StmtKind::Assign { targets: vec![create_expr(expected_var.node.clone())],