diff --git a/crates/ruff/resources/test/fixtures/refurb/FURB132.py b/crates/ruff/resources/test/fixtures/refurb/FURB132.py new file mode 100644 index 00000000000000..65441fbe0264f5 --- /dev/null +++ b/crates/ruff/resources/test/fixtures/refurb/FURB132.py @@ -0,0 +1,80 @@ +from typing import Set +from some.package.name import foo, bar + + +s = set() +s2 = {} +s3: set[int] = foo() + +# these should match + +# FURB132 +if "x" in s: + s.remove("x") + + +# FURB132 +if "x" in s2: + s2.remove("x") + + +# FURB132 +if "x" in s3: + s3.remove("x") + + +var = "y" +# FURB132 +if var in s: + s.remove(var) + + +if f"{var}:{var}" in s: + s.remove(f"{var}:{var}") + + +def identity(x): + return x + + +# FURB132 +if identity("x") in s2: + s2.remove(identity("x")) + + +# these should not + +if "x" in s: + s.remove("y") + +s.discard("x") + +s2 = set() + +if "x" in s: + s2.remove("x") + +if "x" in s: + s.remove("x") + print("removed item") + +if bar() in s: + # bar() might have a side effect + s.remove(bar()) + +if "x" in s: + s.remove("x") +else: + print("not found") + +class Container: + def remove(self, item) -> None: + return + + def __contains__(self, other) -> bool: + return True + +c = Container() + +if "x" in c: + c.remove("x") diff --git a/crates/ruff/src/checkers/ast/analyze/statement.rs b/crates/ruff/src/checkers/ast/analyze/statement.rs index 00582b9ad15b54..90bb7a77af3fb5 100644 --- a/crates/ruff/src/checkers/ast/analyze/statement.rs +++ b/crates/ruff/src/checkers/ast/analyze/statement.rs @@ -1057,6 +1057,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) { checker.diagnostics.push(diagnostic); } } + if checker.enabled(Rule::CheckAndRemoveFromSet) { + refurb::rules::check_and_remove_from_set(checker, if_); + } if checker.source_type.is_stub() { if checker.any_enabled(&[ Rule::UnrecognizedVersionInfoCheck, diff --git a/crates/ruff/src/codes.rs b/crates/ruff/src/codes.rs index 6770ad96b5ae81..83d3854d406651 100644 --- a/crates/ruff/src/codes.rs +++ b/crates/ruff/src/codes.rs @@ -868,6 +868,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> { // refurb (Refurb, "113") => (RuleGroup::Unspecified, rules::refurb::rules::RepeatedAppend), (Refurb, "131") => (RuleGroup::Unspecified, rules::refurb::rules::DeleteFullSlice), + (Refurb, "132") => (RuleGroup::Unspecified, rules::refurb::rules::CheckAndRemoveFromSet), _ => return None, }) diff --git a/crates/ruff/src/rules/refurb/helpers.rs b/crates/ruff/src/rules/refurb/helpers.rs index 764b9dd17514af..24084161f04aad 100644 --- a/crates/ruff/src/rules/refurb/helpers.rs +++ b/crates/ruff/src/rules/refurb/helpers.rs @@ -6,8 +6,8 @@ use ruff_python_semantic::{Binding, BindingKind}; use crate::checkers::ast::Checker; trait TypeChecker { - const EXPR_TYPE: PythonType; fn match_annotation(checker: &Checker, annotation: &Expr) -> bool; + fn match_initializer(checker: &Checker, initializer: &Expr) -> bool; } fn check_type<'a, T: TypeChecker>(checker: &'a Checker, binding: &'a Binding, name: &str) -> bool { @@ -17,11 +17,7 @@ fn check_type<'a, T: TypeChecker>(checker: &'a Checker, binding: &'a Binding, na match binding.kind { BindingKind::Assignment => match stmt { Stmt::Assign(ast::StmtAssign { value, .. }) => { - let value_type: ResolvedPythonType = value.as_ref().into(); - let ResolvedPythonType::Atom(candidate) = value_type else { - return false; - }; - candidate == T::EXPR_TYPE + T::match_initializer(checker, value.as_ref()) } Stmt::AnnAssign(ast::StmtAnnAssign { annotation, .. }) => { T::match_annotation(checker, annotation.as_ref()) @@ -50,9 +46,10 @@ fn check_type<'a, T: TypeChecker>(checker: &'a Checker, binding: &'a Binding, na } } -trait AnnotationTypeChecker { +trait BuiltinTypeChecker { const TYPING_NAME: &'static str; const BUILTIN_TYPE_NAME: &'static str; + const EXPR_TYPE: PythonType; fn match_annotation(checker: &Checker, annotation: &Expr) -> bool { let Expr::Subscript(ast::ExprSubscript { value, .. }) = annotation else { @@ -64,6 +61,25 @@ trait AnnotationTypeChecker { .match_typing_expr(value, Self::TYPING_NAME) } + fn match_initializer(checker: &Checker, initializer: &Expr) -> bool { + Self::match_expr_type(initializer) || Self::match_builtin_constructor(checker, initializer) + } + + fn match_expr_type(initializer: &Expr) -> bool { + let init_type: ResolvedPythonType = initializer.into(); + match init_type { + ResolvedPythonType::Atom(atom) => atom == Self::EXPR_TYPE, + _ => false, + } + } + + fn match_builtin_constructor(checker: &Checker, initializer: &Expr) -> bool { + let Expr::Call(ast::ExprCall { func, .. }) = initializer else { + return false; + }; + Self::match_builtin_type(checker, func.as_ref()) + } + fn match_builtin_type(checker: &Checker, type_expr: &Expr) -> bool { let Expr::Name(ast::ExprName { id, .. }) = type_expr else { return false; @@ -72,32 +88,38 @@ trait AnnotationTypeChecker { } } +impl TypeChecker for T { + fn match_annotation(checker: &Checker, annotation: &Expr) -> bool { + ::match_annotation(checker, annotation) + } + + fn match_initializer(checker: &Checker, initializer: &Expr) -> bool { + ::match_initializer(checker, initializer) + } +} + struct ListChecker; -impl AnnotationTypeChecker for ListChecker { +impl BuiltinTypeChecker for ListChecker { const TYPING_NAME: &'static str = "List"; const BUILTIN_TYPE_NAME: &'static str = "list"; -} - -impl TypeChecker for ListChecker { const EXPR_TYPE: PythonType = PythonType::List; - fn match_annotation(checker: &Checker, annotation: &Expr) -> bool { - ::match_annotation(checker, annotation) - } } struct DictChecker; -impl AnnotationTypeChecker for DictChecker { +impl BuiltinTypeChecker for DictChecker { const TYPING_NAME: &'static str = "Dict"; const BUILTIN_TYPE_NAME: &'static str = "dict"; + const EXPR_TYPE: PythonType = PythonType::Dict; } -impl TypeChecker for DictChecker { - const EXPR_TYPE: PythonType = PythonType::Dict; - fn match_annotation(checker: &Checker, annotation: &Expr) -> bool { - ::match_annotation(checker, annotation) - } +struct SetChecker; + +impl BuiltinTypeChecker for SetChecker { + const TYPING_NAME: &'static str = "Set"; + const BUILTIN_TYPE_NAME: &'static str = "set"; + const EXPR_TYPE: PythonType = PythonType::Set; } pub(super) fn is_list<'a>(checker: &'a Checker, binding: &'a Binding, name: &str) -> bool { @@ -108,6 +130,10 @@ pub(super) fn is_dict<'a>(checker: &'a Checker, binding: &'a Binding, name: &str check_type::(checker, binding, name) } +pub(super) fn is_set<'a>(checker: &'a Checker, binding: &'a Binding, name: &str) -> bool { + check_type::(checker, binding, name) +} + #[inline] fn find_parameter_by_name<'a>( parameters: &'a Parameters, diff --git a/crates/ruff/src/rules/refurb/mod.rs b/crates/ruff/src/rules/refurb/mod.rs index c40524d1f3f8b5..095928a317460b 100644 --- a/crates/ruff/src/rules/refurb/mod.rs +++ b/crates/ruff/src/rules/refurb/mod.rs @@ -15,6 +15,7 @@ mod tests { #[test_case(Rule::RepeatedAppend, Path::new("FURB113.py"))] #[test_case(Rule::DeleteFullSlice, Path::new("FURB131.py"))] + #[test_case(Rule::CheckAndRemoveFromSet, Path::new("FURB132.py"))] fn rules(rule_code: Rule, path: &Path) -> Result<()> { let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy()); let diagnostics = test_path( diff --git a/crates/ruff/src/rules/refurb/rules/check_and_remove_from_set.rs b/crates/ruff/src/rules/refurb/rules/check_and_remove_from_set.rs new file mode 100644 index 00000000000000..913b58e5bbefba --- /dev/null +++ b/crates/ruff/src/rules/refurb/rules/check_and_remove_from_set.rs @@ -0,0 +1,207 @@ +use ast::{comparable::ComparableExpr, helpers::contains_effect, CmpOp, Ranged}; +use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix}; +use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast::{self as ast, Expr, Stmt}; +use ruff_python_codegen::Generator; +use ruff_python_semantic::Binding; +use ruff_text_size::TextRange; + +use crate::{checkers::ast::Checker, rules::refurb::helpers::is_set}; + +/// ## What it does +/// Checks for check and `remove` pattern that can be replaced via `discard`. +/// +/// ## Why is this bad? +/// It is more succinct and idiomatic to use `discard`. +/// +/// ## Example +/// ```python +/// nums = {123, 456} +/// +/// if 123 in nums: +/// nums.remove(123) +/// ``` +/// +/// Use instead: +/// ```python +/// nums = {123, 456} +/// +/// nums.discard(123) +/// ``` +#[violation] +pub struct CheckAndRemoveFromSet { + element: String, + set: String, +} + +impl AlwaysAutofixableViolation for CheckAndRemoveFromSet { + #[derive_message_formats] + fn message(&self) -> String { + let CheckAndRemoveFromSet { element, set } = self; + format!("Replace `if {element} in {set}: {set}.remove({element})` with `{set}.discard({element})`") + } + + fn autofix_title(&self) -> String { + let CheckAndRemoveFromSet { element, set } = self; + format!("Replace with `{set}.discard({element})`") + } +} + +// FURB132 +pub(crate) fn check_and_remove_from_set(checker: &mut Checker, if_stmt: &ast::StmtIf) { + // In order to fit the profile, we need if without else clauses and with only one statement in its body. + if if_stmt.body.len() != 1 || !if_stmt.elif_else_clauses.is_empty() { + return; + } + + // if test should be `element in set` + let Some((check_element, check_set)) = match_check(if_stmt) else { + return; + }; + + // if body should be `set.remove(element)` + let Some((remove_element, remove_set)) = match_remove(if_stmt) else { + return; + }; + + // `element` in the check should be the same as `element` in the body + if !compare(&check_element.into(), &remove_element.into()) + // `set` in the check should be the same as `set` in the body + || check_set.id != remove_set.id + // `element` shouldn't have a side effect, otherwise we might change the samntic of the program + || contains_effect(check_element, |id| checker.semantic().is_builtin(id)) + { + return; + } + + let Some(binding) = find_binding(checker, &check_set.id) else { + return; + }; + + if !is_set(checker, binding, &check_set.id) { + return; + }; + + let replacement = make_suggestion(check_set, check_element, checker.generator()); + let element_str = checker.generator().expr(check_element); + let mut diagnostic = Diagnostic::new( + CheckAndRemoveFromSet { + element: element_str, + set: check_set.id.to_string(), + }, + if_stmt.range(), + ); + diagnostic.set_fix(Fix::suggested(Edit::replacement( + replacement, + if_stmt.start(), + if_stmt.end(), + ))); + + checker.diagnostics.push(diagnostic); +} + +fn compare(lhs: &ComparableExpr, rhs: &ComparableExpr) -> bool { + lhs == rhs +} + +fn make_suggestion(set: &ast::ExprName, element: &Expr, generator: Generator) -> String { + // Here we construct `set.discard(element)` + // + // Let's make `set.discard`. + let attr = ast::ExprAttribute { + value: Box::new(set.clone().into()), + attr: ast::Identifier::new("discard".to_string(), TextRange::default()), + ctx: ast::ExprContext::Load, + range: TextRange::default(), + }; + // Make the actual call `set.discard(element)` + let call = ast::ExprCall { + func: Box::new(attr.into()), + arguments: ast::Arguments { + args: vec![element.clone()], + keywords: vec![], + range: TextRange::default(), + }, + range: TextRange::default(), + }; + // And finally, turn it into a statement. + let stmt = ast::StmtExpr { + value: Box::new(call.into()), + range: TextRange::default(), + }; + generator.stmt(&stmt.into()) +} + +fn find_binding<'a>(checker: &'a Checker, name: &str) -> Option<&'a Binding<'a>> { + // Let's find definition for var + let scope = checker.semantic().current_scope(); + let bindings: Vec<&Binding> = scope + .get_all(name) + .map(|binding_id| checker.semantic().binding(binding_id)) + .collect(); + + let [binding @ Binding { + source: Some(..), .. + }] = bindings.as_slice() + else { + return None; + }; + + Some(binding) +} + +fn match_check(if_stmt: &ast::StmtIf) -> Option<(&Expr, &ast::ExprName)> { + let Expr::Compare(ast::ExprCompare { + ops, + left, + comparators, + .. + }) = if_stmt.test.as_ref() + else { + return None; + }; + + if ops.as_slice() != [CmpOp::In] { + return None; + } + + let [Expr::Name(right @ ast::ExprName { .. })] = comparators.as_slice() else { + return None; + }; + + Some((left.as_ref(), right)) +} + +fn match_remove(if_stmt: &ast::StmtIf) -> Option<(&Expr, &ast::ExprName)> { + let [Stmt::Expr(ast::StmtExpr { value: expr, .. })] = if_stmt.body.as_slice() else { + return None; + }; + + let Expr::Call(ast::ExprCall { + func: attr, + arguments: ast::Arguments { args, keywords, .. }, + .. + }) = expr.as_ref() + else { + return None; + }; + + let Expr::Attribute(ast::ExprAttribute { + value: receiver, + attr: func_name, + .. + }) = attr.as_ref() + else { + return None; + }; + + let Expr::Name(ref set @ ast::ExprName { .. }) = receiver.as_ref() else { + return None; + }; + + if func_name != "remove" || args.len() != 1 || !keywords.is_empty() { + return None; + } + + Some((args.first().unwrap(), set)) +} diff --git a/crates/ruff/src/rules/refurb/rules/mod.rs b/crates/ruff/src/rules/refurb/rules/mod.rs index c897f03072c48b..1ea05511ee02a0 100644 --- a/crates/ruff/src/rules/refurb/rules/mod.rs +++ b/crates/ruff/src/rules/refurb/rules/mod.rs @@ -1,5 +1,7 @@ +pub(crate) use check_and_remove_from_set::*; pub(crate) use delete_full_slice::*; pub(crate) use repeated_append::*; +mod check_and_remove_from_set; mod delete_full_slice; mod repeated_append; diff --git a/crates/ruff/src/rules/refurb/snapshots/ruff__rules__refurb__tests__FURB132_FURB132.py.snap b/crates/ruff/src/rules/refurb/snapshots/ruff__rules__refurb__tests__FURB132_FURB132.py.snap new file mode 100644 index 00000000000000..2366c41cd78a18 --- /dev/null +++ b/crates/ruff/src/rules/refurb/snapshots/ruff__rules__refurb__tests__FURB132_FURB132.py.snap @@ -0,0 +1,84 @@ +--- +source: crates/ruff/src/rules/refurb/mod.rs +--- +FURB132.py:12:1: FURB132 [*] Replace `if "x" in s: s.remove("x")` with `s.discard("x")` + | +11 | # FURB132 +12 | / if "x" in s: +13 | | s.remove("x") + | |_________________^ FURB132 + | + = help: Replace with `s.discard("x")` + +ℹ Suggested fix +9 9 | # these should match +10 10 | +11 11 | # FURB132 +12 |-if "x" in s: +13 |- s.remove("x") + 12 |+s.discard("x") +14 13 | +15 14 | +16 15 | # FURB132 + +FURB132.py:22:1: FURB132 [*] Replace `if "x" in s3: s3.remove("x")` with `s3.discard("x")` + | +21 | # FURB132 +22 | / if "x" in s3: +23 | | s3.remove("x") + | |__________________^ FURB132 + | + = help: Replace with `s3.discard("x")` + +ℹ Suggested fix +19 19 | +20 20 | +21 21 | # FURB132 +22 |-if "x" in s3: +23 |- s3.remove("x") + 22 |+s3.discard("x") +24 23 | +25 24 | +26 25 | var = "y" + +FURB132.py:28:1: FURB132 [*] Replace `if var in s: s.remove(var)` with `s.discard(var)` + | +26 | var = "y" +27 | # FURB132 +28 | / if var in s: +29 | | s.remove(var) + | |_________________^ FURB132 + | + = help: Replace with `s.discard(var)` + +ℹ Suggested fix +25 25 | +26 26 | var = "y" +27 27 | # FURB132 +28 |-if var in s: +29 |- s.remove(var) + 28 |+s.discard(var) +30 29 | +31 30 | +32 31 | if f"{var}:{var}" in s: + +FURB132.py:32:1: FURB132 [*] Replace `if f"{var}:{var}" in s: s.remove(f"{var}:{var}")` with `s.discard(f"{var}:{var}")` + | +32 | / if f"{var}:{var}" in s: +33 | | s.remove(f"{var}:{var}") + | |____________________________^ FURB132 + | + = help: Replace with `s.discard(f"{var}:{var}")` + +ℹ Suggested fix +29 29 | s.remove(var) +30 30 | +31 31 | +32 |-if f"{var}:{var}" in s: +33 |- s.remove(f"{var}:{var}") + 32 |+s.discard(f"{var}:{var}") +34 33 | +35 34 | +36 35 | def identity(x): + + diff --git a/ruff.schema.json b/ruff.schema.json index 9095e32319d80d..00daa5dfae56c4 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -2049,6 +2049,7 @@ "FURB113", "FURB13", "FURB131", + "FURB132", "G", "G0", "G00",