[flake8-bugbear] add autofix for B006 mutable_argument_default

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

@@ -68,6 +68,20 @@ def this_is_also_wrong(value={}):
     ...
 
 
+class Foo:
+    @staticmethod
+    def this_is_also_wrong_and_more_indented(value={}):
+        pass
+
+
+def multiline_arg_wrong(value={
+
+}):
+    ...
+
+def single_line_func_wrong(value = {}): ...
+
+
 def and_this(value=set()):
     ...
 

crates/ruff/src/checkers/ast/mod.rs

@@ -343,6 +343,7 @@ where
                 returns,
                 args,
                 body,
+                range,
                 ..
             })
             | Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef {
@@ -351,8 +352,12 @@ where
                 returns,
                 args,
                 body,
+                range,
                 ..
             }) => {
+                if self.enabled(Rule::MutableArgumentDefault) {
+                    flake8_bugbear::rules::mutable_argument_default(self, args, body, *range);
+                }
                 if self.enabled(Rule::DjangoNonLeadingReceiverDecorator) {
                     flake8_django::rules::non_leading_receiver_decorator(self, decorator_list);
                 }
@@ -4010,9 +4015,6 @@ where
     }
 
     fn visit_arguments(&mut self, arguments: &'b Arguments) {
-        if self.enabled(Rule::MutableArgumentDefault) {
-            flake8_bugbear::rules::mutable_argument_default(self, arguments);
-        }
         if self.enabled(Rule::FunctionCallInDefaultArgument) {
             flake8_bugbear::rules::function_call_argument_default(self, arguments);
         }

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

@@ -1,8 +1,15 @@
-use rustpython_parser::ast::{ArgWithDefault, Arguments, Ranged};
+use rustpython_parser::ast::{ArgWithDefault, Arguments, Ranged, Stmt};
 
-use ruff_diagnostics::{Diagnostic, Violation};
+use crate::registry::AsRule;
+use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::docstrings::leading_space;
+use ruff_python_ast::source_code::Locator;
+use ruff_python_ast::whitespace::indentation_at_offset;
 use ruff_python_semantic::analyze::typing::{is_immutable_annotation, is_mutable_expr};
+use ruff_text_size::TextRange;
+use rustpython_parser::lexer::lex_starts_at;
+use rustpython_parser::{Mode, Tok};
 
 use crate::checkers::ast::Checker;
 
@@ -50,14 +57,26 @@ use crate::checkers::ast::Checker;
 pub struct MutableArgumentDefault;
 
 impl Violation for MutableArgumentDefault {
+    const AUTOFIX: AutofixKind = AutofixKind::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
-        format!("Do not use mutable data structures for argument defaults")
+        format!("Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`")
+    }
+    fn autofix_title(&self) -> Option<String> {
+        Some(format!(
+            "Do not use mutable data structures for argument defaults"
+        ))
     }
 }
 
 /// B006
-pub(crate) fn mutable_argument_default(checker: &mut Checker, arguments: &Arguments) {
+pub(crate) fn mutable_argument_default(
+    checker: &mut Checker,
+    arguments: &Arguments,
+    body: &[Stmt],
+    func_range: TextRange,
+) {
     // Scan in reverse order to right-align zip().
     for ArgWithDefault {
         def,
@@ -78,9 +97,53 @@ pub(crate) fn mutable_argument_default(checker: &mut Checker, arguments: &Argume
                 is_immutable_annotation(expr, checker.semantic())
             })
         {
-            checker
-                .diagnostics
-                .push(Diagnostic::new(MutableArgumentDefault, default.range()));
+            let mut diagnostic = Diagnostic::new(MutableArgumentDefault, default.range());
+
+            // If the function body is on the same line as the function def, do not fix
+            if checker.patch(diagnostic.kind.rule())
+                && !is_single_line(checker.locator, func_range, body)
+            {
+                // Set the default arg value to None
+                let arg_edit = Edit::range_replacement("None".to_string(), default.range());
+
+                // Add conditional check to set the default arg to its original value if still None
+                let mut check_lines = String::new();
+                let indentation =
+                    indentation_at_offset(checker.locator, body[0].start()).unwrap_or_default();
+                let indentation = leading_space(indentation);
+                // body[0].start() starts at correct indentation so we do need to add indentation
+                // before pushing the if statement
+                check_lines.push_str(format!("if {} is None:\n", def.arg.as_str()).as_str());
+                check_lines.push_str(indentation);
+                check_lines.push_str(checker.stylist.indentation());
+                check_lines.push_str(
+                    format!(
+                        "{} = {}",
+                        def.arg.as_str(),
+                        checker.generator().expr(default),
+                    )
+                    .as_str(),
+                );
+                check_lines.push_str(&checker.stylist.line_ending());
+                check_lines.push_str(indentation);
+                let check_edit = Edit::insertion(check_lines, body[0].start());
+
+                diagnostic.set_fix(Fix::manual_edits(arg_edit, [check_edit]));
+            }
+            checker.diagnostics.push(diagnostic);
+        }
+    }
+}
+
+fn is_single_line(locator: &Locator, func_range: TextRange, body: &[Stmt]) -> bool {
+    let arg_string = locator.slice(func_range);
+    for (tok, range) in lex_starts_at(arg_string, Mode::Module, func_range.start()).flatten() {
+        match tok {
+            Tok::Colon => {
+                return !locator.contains_line_break(TextRange::new(range.end(), body[0].start()))
+            }
+            _ => continue,
         }
     }
+    false
 }