[refurb] Implement check-and-remove-from-set rule (FURB132)

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")

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,

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,
     })
 }

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::<DictChecker>(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::<SetChecker>(binding, name, semantic)
+}
+
 #[inline]
 fn find_parameter_by_name<'a>(
     parameters: &'a Parameters,

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(

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())
+}

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;