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 0000000000000..5e94524aebfbf --- /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 + + +# TODO: 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 92bb6ecfa25c0..fabc60e3eaf58 100644 --- a/crates/ruff/src/checkers/ast/analyze/statement.rs +++ b/crates/ruff/src/checkers/ast/analyze/statement.rs @@ -1056,6 +1056,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 54398a0017132..67eb3887db769 100644 --- a/crates/ruff/src/codes.rs +++ b/crates/ruff/src/codes.rs @@ -868,5 +868,8 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> { // refurb (Refurb, "113") => (RuleGroup::Nursery, rules::refurb::rules::RepeatedAppend), (Refurb, "131") => (RuleGroup::Nursery, rules::refurb::rules::DeleteFullSlice), + (Refurb, "132") => (RuleGroup::Nursery, 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 65e0d3acdd8c1..d1d92edd52670 100644 --- a/crates/ruff/src/rules/refurb/helpers.rs +++ b/crates/ruff/src/rules/refurb/helpers.rs @@ -149,6 +149,14 @@ impl BuiltinTypeChecker for DictChecker { const EXPR_TYPE: PythonType = PythonType::Dict; } +struct SetChecker; + +impl BuiltinTypeChecker for SetChecker { + const BUILTIN_TYPE_NAME: &'static str = "set"; + const TYPING_NAME: &'static str = "Set"; + const EXPR_TYPE: PythonType = PythonType::Set; +} + /// Test whether the given binding (and the given name) can be considered a list. /// For this, we check what value might be associated with it through it's initialization and /// what annotation it has (we consider `list` and `typing.List`). @@ -163,6 +171,13 @@ pub(super) fn is_dict<'a>(binding: &'a Binding, name: &str, semantic: &'a Semant check_type::(binding, name, semantic) } +/// Test whether the given binding (and the given name) can be considered a set. +/// For this, we check what value might be associated with it through it's initialization and +/// what annotation it has (we consider `set` and `typing.Set`). +pub(super) fn is_set<'a>(binding: &'a Binding, name: &str, semantic: &'a SemanticModel) -> bool { + check_type::(binding, name, semantic) +} + #[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 2537f47b139a4..1dddb615fbf2a 100644 --- a/crates/ruff/src/rules/refurb/mod.rs +++ b/crates/ruff/src/rules/refurb/mod.rs @@ -16,6 +16,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 0000000000000..e9597c428e365 --- /dev/null +++ b/crates/ruff/src/rules/refurb/rules/check_and_remove_from_set.rs @@ -0,0 +1,206 @@ +use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix}; +use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast::comparable::ComparableExpr; +use ruff_python_ast::helpers::contains_effect; +use ruff_python_ast::{self as ast, CmpOp, Expr, Stmt}; +use ruff_python_codegen::Generator; +use ruff_text_size::{Ranged, TextRange}; + +use crate::autofix::snippet::SourceCodeSnippet; +use crate::checkers::ast::Checker; +use crate::registry::AsRule; +use crate::rules::refurb::helpers::is_set; + +/// ## What it does +/// Checks for uses of `set#remove` that can be replaced with `set#discard`. +/// +/// ## Why is this bad? +/// If an element should be removed from a set if it is present, 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) +/// ``` +/// +/// ## References +/// - [Python documentation: `set.discard()`](https://docs.python.org/3/library/stdtypes.html?highlight=list#frozenset.discard) +#[violation] +pub struct CheckAndRemoveFromSet { + element: SourceCodeSnippet, + set: String, +} + +impl CheckAndRemoveFromSet { + fn suggestion(&self) -> String { + let set = &self.set; + let element = self.element.truncated_display(); + format!("{set}.discard({element})") + } +} + +impl AlwaysAutofixableViolation for CheckAndRemoveFromSet { + #[derive_message_formats] + fn message(&self) -> String { + let suggestion = self.suggestion(); + format!("Use `{suggestion}` instead of check and `remove`") + } + + fn autofix_title(&self) -> String { + let suggestion = self.suggestion(); + format!("Replace with `{suggestion}`") + } +} + +/// 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; + } + + // The `if` test should be `element in set`. + let Some((check_element, check_set)) = match_check(if_stmt) else { + return; + }; + + // The `if` body should be `set.remove(element)`. + let Some((remove_element, remove_set)) = match_remove(if_stmt) else { + return; + }; + + // ` + // `set` in the check should be the same as `set` in the body + if check_set.id != remove_set.id + // `element` in the check should be the same as `element` in the body + || !compare(&check_element.into(), &remove_element.into()) + // `element` shouldn't have a side effect, otherwise we might change the semantics of the program. + || contains_effect(check_element, |id| checker.semantic().is_builtin(id)) + { + return; + } + + // Check if what we assume is set is indeed a set. + if !checker + .semantic() + .resolve_name(check_set) + .map(|binding_id| checker.semantic().binding(binding_id)) + .map_or(false, |binding| { + is_set(binding, &check_set.id, checker.semantic()) + }) + { + return; + }; + + let mut diagnostic = Diagnostic::new( + CheckAndRemoveFromSet { + element: SourceCodeSnippet::from_str(checker.locator().slice(check_element)), + set: check_set.id.to_string(), + }, + if_stmt.range(), + ); + if checker.patch(diagnostic.kind.rule()) { + diagnostic.set_fix(Fix::suggested(Edit::replacement( + make_suggestion(check_set, check_element, checker.generator()), + if_stmt.start(), + if_stmt.end(), + ))); + } + checker.diagnostics.push(diagnostic); +} + +fn compare(lhs: &ComparableExpr, rhs: &ComparableExpr) -> bool { + lhs == rhs +} + +/// Match `if` condition to be `expr in name`, returns a tuple of (`expr`, `name`) on success. +fn match_check(if_stmt: &ast::StmtIf) -> Option<(&Expr, &ast::ExprName)> { + let ast::ExprCompare { + ops, + left, + comparators, + .. + } = if_stmt.test.as_compare_expr()?; + + 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)) +} + +/// Match `if` body to be `name.remove(expr)`, returns a tuple of (`expr`, `name`) on success. +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 ast::ExprCall { + func: attr, + arguments: ast::Arguments { args, keywords, .. }, + .. + } = expr.as_call_expr()?; + + let ast::ExprAttribute { + value: receiver, + attr: func_name, + .. + } = attr.as_attribute_expr()?; + + let Expr::Name(ref set @ ast::ExprName { .. }) = receiver.as_ref() else { + return None; + }; + + let [arg] = args.as_slice() else { + return None; + }; + + if func_name != "remove" || !keywords.is_empty() { + return None; + } + + Some((arg, set)) +} + +/// Construct the fix suggestion, ie `set.discard(element)`. +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()) +} diff --git a/crates/ruff/src/rules/refurb/rules/mod.rs b/crates/ruff/src/rules/refurb/rules/mod.rs index 8fef67abe0d02..1ea05511ee02a 100644 --- a/crates/ruff/src/rules/refurb/rules/mod.rs +++ b/crates/ruff/src/rules/refurb/rules/mod.rs @@ -1,6 +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 0000000000000..448c6ccecd614 --- /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 [*] Use `s.discard("x")` instead of check and `remove` + | +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 [*] Use `s3.discard("x")` instead of check and `remove` + | +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 [*] Use `s.discard(var)` instead of check and `remove` + | +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 [*] Use `s.discard(f"{var}:{var}")` instead of check and `remove` + | +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 c7036fc71c077..e4051de9f6a54 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -2046,6 +2046,7 @@ "FURB", "FURB113", "FURB131", + "FURB132", "G", "G0", "G00",