Avoid boolean-trap errors in __setitem__ (#2636)

crates/ruff/resources/test/fixtures/flake8_boolean_trap/FBT.py

@@ -57,3 +57,12 @@ dict.fromkeys(("world",), True)
 {}.deploy(True, False)
 getattr(someobj, attrname, False)
 mylist.index(True)
+
+
+class Registry:
+    def __init__(self) -> None:
+        self._switches = [False] * len(Switch)
+
+    # FBT001: Boolean positional arg in function definition
+    def __setitem__(self, switch: Switch, value: bool) -> None:
+        self._switches[switch.value] = value

crates/ruff/src/checkers/ast.rs

@@ -696,6 +696,24 @@ where
                     flake8_pytest_style::rules::marks(self, decorator_list);
                 }
 
+                if self
+                    .settings
+                    .rules
+                    .enabled(&Rule::BooleanPositionalArgInFunctionDefinition)
+                {
+                    flake8_boolean_trap::rules::check_positional_boolean_in_def(self, name, args);
+                }
+
+                if self
+                    .settings
+                    .rules
+                    .enabled(&Rule::BooleanDefaultValueInFunctionDefinition)
+                {
+                    flake8_boolean_trap::rules::check_boolean_default_value_in_function_definition(
+                        self, name, args,
+                    );
+                }
+
                 self.check_builtin_shadowing(name, stmt, true);
 
                 // Visit the decorators and arguments, but avoid the body, which will be
@@ -3655,24 +3673,6 @@ where
             flake8_bugbear::rules::function_call_argument_default(self, arguments);
         }
 
-        // flake8-boolean-trap
-        if self
-            .settings
-            .rules
-            .enabled(&Rule::BooleanPositionalArgInFunctionDefinition)
-        {
-            flake8_boolean_trap::rules::check_positional_boolean_in_def(self, arguments);
-        }
-        if self
-            .settings
-            .rules
-            .enabled(&Rule::BooleanDefaultValueInFunctionDefinition)
-        {
-            flake8_boolean_trap::rules::check_boolean_default_value_in_function_definition(
-                self, arguments,
-            );
-        }
-
         // Bind, but intentionally avoid walking default expressions, as we handle them
         // upstream.
         for arg in &arguments.posonlyargs {

crates/ruff/src/rules/flake8_boolean_trap/rules.rs

@@ -38,7 +38,7 @@ impl Violation for BooleanPositionalValueInFunctionCall {
     }
 }
 
-const FUNC_NAME_ALLOWLIST: &[&str] = &[
+const FUNC_CALL_NAME_ALLOWLIST: &[&str] = &[
     "assertEqual",
     "assertEquals",
     "assertNotEqual",
@@ -54,16 +54,18 @@ const FUNC_NAME_ALLOWLIST: &[&str] = &[
     "setdefault",
 ];
 
+const FUNC_DEF_NAME_ALLOWLIST: &[&str] = &["__setitem__"];
+
 /// Returns `true` if an argument is allowed to use a boolean trap. To return
 /// `true`, the function name must be explicitly allowed, and the argument must
 /// be either the first or second argument in the call.
 fn allow_boolean_trap(func: &Expr) -> bool {
     if let ExprKind::Attribute { attr, .. } = &func.node {
-        return FUNC_NAME_ALLOWLIST.contains(&attr.as_ref());
+        return FUNC_CALL_NAME_ALLOWLIST.contains(&attr.as_ref());
     }
 
     if let ExprKind::Name { id, .. } = &func.node {
-        return FUNC_NAME_ALLOWLIST.contains(&id.as_ref());
+        return FUNC_CALL_NAME_ALLOWLIST.contains(&id.as_ref());
     }
 
     false
@@ -87,7 +89,10 @@ fn add_if_boolean(checker: &mut Checker, arg: &Expr, kind: DiagnosticKind) {
     }
 }
 
-pub fn check_positional_boolean_in_def(checker: &mut Checker, arguments: &Arguments) {
+pub fn check_positional_boolean_in_def(checker: &mut Checker, name: &str, arguments: &Arguments) {
+    if FUNC_DEF_NAME_ALLOWLIST.contains(&name) {
+        return;
+    }
     for arg in arguments.posonlyargs.iter().chain(arguments.args.iter()) {
         if arg.node.annotation.is_none() {
             continue;
@@ -117,8 +122,12 @@ pub fn check_positional_boolean_in_def(checker: &mut Checker, arguments: &Argume
 
 pub fn check_boolean_default_value_in_function_definition(
     checker: &mut Checker,
+    name: &str,
     arguments: &Arguments,
 ) {
+    if FUNC_DEF_NAME_ALLOWLIST.contains(&name) {
+        return;
+    }
     for arg in &arguments.defaults {
         add_if_boolean(checker, arg, BooleanDefaultValueInFunctionDefinition.into());
     }
@@ -129,10 +138,10 @@ pub fn check_boolean_positional_value_in_function_call(
     args: &[Expr],
     func: &Expr,
 ) {
+    if allow_boolean_trap(func) {
+        return;
+    }
     for arg in args {
-        if allow_boolean_trap(func) {
-            continue;
-        }
         add_if_boolean(checker, arg, BooleanPositionalValueInFunctionCall.into());
     }
 }