Update mutable-argument-default (B006) to use extend-immutable-calls when determining if annotations are immutable (#6781)

Part of https://github.com/astral-sh/ruff/issues/3762

crates/ruff/resources/test/fixtures/flake8_bugbear/B006_extended.py

@@ -0,0 +1,18 @@
+import custom
+from custom import ImmutableTypeB
+
+
+def okay(foo: ImmutableTypeB = []):
+    ...
+
+
+def okay(foo: custom.ImmutableTypeA = []):
+    ...
+
+
+def okay(foo: custom.ImmutableTypeB = []):
+    ...
+
+
+def error_due_to_missing_import(foo: ImmutableTypeA = []):
+    ...

crates/ruff/src/rules/flake8_bugbear/mod.rs

@@ -71,8 +71,27 @@ mod tests {
     }
 
     #[test]
-    fn extend_immutable_calls() -> Result<()> {
+    fn extend_immutable_calls_arg_annotation() -> Result<()> {
-        let snapshot = "extend_immutable_calls".to_string();
+        let snapshot = "extend_immutable_calls_arg_annotation".to_string();
+        let diagnostics = test_path(
+            Path::new("flake8_bugbear/B006_extended.py"),
+            &Settings {
+                flake8_bugbear: super::settings::Settings {
+                    extend_immutable_calls: vec![
+                        "custom.ImmutableTypeA".to_string(),
+                        "custom.ImmutableTypeB".to_string(),
+                    ],
+                },
+                ..Settings::for_rule(Rule::MutableArgumentDefault)
+            },
+        )?;
+        assert_messages!(snapshot, diagnostics);
+        Ok(())
+    }
+
+    #[test]
+    fn extend_immutable_calls_arg_default() -> Result<()> {
+        let snapshot = "extend_immutable_calls_arg_default".to_string();
         let diagnostics = test_path(
             Path::new("flake8_bugbear/B008_extended.py"),
             &Settings {

crates/ruff/src/rules/flake8_bugbear/rules/mutable_argument_default.rs

@@ -1,3 +1,4 @@
+use ast::call_path::{from_qualified_name, CallPath};
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::is_docstring_stmt;
@@ -25,6 +26,10 @@ use crate::registry::AsRule;
 /// default, and initialize a new mutable object inside the function body
 /// for each call.
 ///
+/// Arguments with immutable type annotations will be ignored by this rule.
+/// Types outside of the standard library can be marked as immutable with the
+/// [`flake8-bugbear.extend-immutable-calls`] configuration option.
+///
 /// ## Example
 /// ```python
 /// def add_to_list(item, some_list=[]):
@@ -49,6 +54,9 @@ use crate::registry::AsRule;
 /// l2 = add_to_list(1)  # [1]
 /// ```
 ///
+/// ## Options
+/// - `flake8-bugbear.extend-immutable-calls`
+///
 /// ## References
 /// - [Python documentation: Default Argument Values](https://docs.python.org/3/tutorial/controlflow.html#default-argument-values)
 #[violation]
@@ -84,11 +92,18 @@ pub(crate) fn mutable_argument_default(checker: &mut Checker, function_def: &ast
             continue;
         };
 
+        let extend_immutable_calls: Vec<CallPath> = checker
+            .settings
+            .flake8_bugbear
+            .extend_immutable_calls
+            .iter()
+            .map(|target| from_qualified_name(target))
+            .collect();
+
         if is_mutable_expr(default, checker.semantic())
-            && !parameter
+            && !parameter.annotation.as_ref().is_some_and(|expr| {
-                .annotation
+                is_immutable_annotation(expr, checker.semantic(), extend_immutable_calls.as_slice())
-                .as_ref()
+            })
-                .is_some_and(|expr| is_immutable_annotation(expr, checker.semantic()))
         {
             let mut diagnostic = Diagnostic::new(MutableArgumentDefault, default.range());
 

crates/ruff/src/rules/flake8_bugbear/snapshots/ruff__rules__flake8_bugbear__tests__extend_immutable_calls_arg_annotation.snap

@@ -0,0 +1,22 @@
+---
+source: crates/ruff/src/rules/flake8_bugbear/mod.rs
+---
+B006_extended.py:17:55: B006 [*] Do not use mutable data structures for argument defaults
+   |
+17 | def error_due_to_missing_import(foo: ImmutableTypeA = []):
+   |                                                       ^^ B006
+18 |     ...
+   |
+   = help: Replace with `None`; initialize within function
+
+ℹ Possible fix
+14 14 |     ...
+15 15 | 
+16 16 | 
+17    |-def error_due_to_missing_import(foo: ImmutableTypeA = []):
+   17 |+def error_due_to_missing_import(foo: ImmutableTypeA = None):
+   18 |+    if foo is None:
+   19 |+        foo = []
+18 20 |     ...
+
+

crates/ruff/src/rules/ruff/rules/mutable_class_default.rs

@@ -60,7 +60,7 @@ pub(crate) fn mutable_class_default(checker: &mut Checker, class_def: &ast::Stmt
                     && is_mutable_expr(value, checker.semantic())
                     && !is_class_var_annotation(annotation, checker.semantic())
                     && !is_final_annotation(annotation, checker.semantic())
-                    && !is_immutable_annotation(annotation, checker.semantic())
+                    && !is_immutable_annotation(annotation, checker.semantic(), &[])
                     && !is_dataclass(class_def, checker.semantic())
                 {
                     // Avoid Pydantic models, which end up copying defaults on instance creation.

crates/ruff/src/rules/ruff/rules/mutable_dataclass_default.rs

@@ -76,7 +76,7 @@ pub(crate) fn mutable_dataclass_default(checker: &mut Checker, class_def: &ast::
         {
             if is_mutable_expr(value, checker.semantic())
                 && !is_class_var_annotation(annotation, checker.semantic())
-                && !is_immutable_annotation(annotation, checker.semantic())
+                && !is_immutable_annotation(annotation, checker.semantic(), &[])
             {
                 checker
                     .diagnostics

crates/ruff_python_semantic/src/analyze/typing.rs

@@ -186,12 +186,19 @@ pub fn to_pep604_operator(
 
 /// Return `true` if `Expr` represents a reference to a type annotation that resolves to an
 /// immutable type.
-pub fn is_immutable_annotation(expr: &Expr, semantic: &SemanticModel) -> bool {
+pub fn is_immutable_annotation(
+    expr: &Expr,
+    semantic: &SemanticModel,
+    extend_immutable_calls: &[CallPath],
+) -> bool {
     match expr {
         Expr::Name(_) | Expr::Attribute(_) => {
             semantic.resolve_call_path(expr).is_some_and(|call_path| {
                 is_immutable_non_generic_type(call_path.as_slice())
                     || is_immutable_generic_type(call_path.as_slice())
+                    || extend_immutable_calls
+                        .iter()
+                        .any(|target| call_path == *target)
             })
         }
         Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
@@ -200,17 +207,19 @@ pub fn is_immutable_annotation(expr: &Expr, semantic: &SemanticModel) -> bool {
                     true
                 } else if matches!(call_path.as_slice(), ["typing", "Union"]) {
                     if let Expr::Tuple(ast::ExprTuple { elts, .. }) = slice.as_ref() {
-                        elts.iter()
+                        elts.iter().all(|elt| {
-                            .all(|elt| is_immutable_annotation(elt, semantic))
+                            is_immutable_annotation(elt, semantic, extend_immutable_calls)
+                        })
                     } else {
                         false
                     }
                 } else if matches!(call_path.as_slice(), ["typing", "Optional"]) {
-                    is_immutable_annotation(slice, semantic)
+                    is_immutable_annotation(slice, semantic, extend_immutable_calls)
                 } else if is_pep_593_generic_type(call_path.as_slice()) {
                     if let Expr::Tuple(ast::ExprTuple { elts, .. }) = slice.as_ref() {
-                        elts.first()
+                        elts.first().is_some_and(|elt| {
-                            .is_some_and(|elt| is_immutable_annotation(elt, semantic))
+                            is_immutable_annotation(elt, semantic, extend_immutable_calls)
+                        })
                     } else {
                         false
                     }
@@ -224,7 +233,10 @@ pub fn is_immutable_annotation(expr: &Expr, semantic: &SemanticModel) -> bool {
             op: Operator::BitOr,
             right,
             range: _,
-        }) => is_immutable_annotation(left, semantic) && is_immutable_annotation(right, semantic),
+        }) => {
+            is_immutable_annotation(left, semantic, extend_immutable_calls)
+                && is_immutable_annotation(right, semantic, extend_immutable_calls)
+        }
         Expr::Constant(ast::ExprConstant {
             value: Constant::None,
             ..