Ensure that B006 autofix respects docstrings (#6493)

## Summary

Some follow-ups to https://github.com/astral-sh/ruff/pull/6131 to ensure
that fixes are inserted _after_ function docstrings, and that fixes are
robust to a bunch of edge cases.

## Test Plan

`cargo test`

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

@@ -275,3 +275,32 @@ def mutable_annotations(
     d: typing_extensions.Annotated[Union[Set[str], abc.Sized], "annotation"] = set(),
 ):
     pass
+
+
+def single_line_func_wrong(value: dict[str, str] = {}):
+    """Docstring"""
+
+
+def single_line_func_wrong(value: dict[str, str] = {}):
+    """Docstring"""
+    ...
+
+
+def single_line_func_wrong(value: dict[str, str] = {}):
+    """Docstring"""; ...
+
+
+def single_line_func_wrong(value: dict[str, str] = {}):
+    """Docstring"""; \
+        ...
+
+
+def single_line_func_wrong(value: dict[str, str] = {
+    # This is a comment
+}):
+    """Docstring"""
+
+
+def single_line_func_wrong(value: dict[str, str] = {}) \
+    : \
+    """Docstring"""

crates/ruff/src/checkers/ast/analyze/statement.rs

@@ -69,16 +69,18 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 }
             }
         }
-        Stmt::FunctionDef(ast::StmtFunctionDef {
+        Stmt::FunctionDef(
-            is_async,
+            function_def @ ast::StmtFunctionDef {
-            name,
+                is_async,
-            decorator_list,
+                name,
-            returns,
+                decorator_list,
-            parameters,
+                returns,
-            body,
+                parameters,
-            type_params,
+                body,
-            range,
+                type_params,
-        }) => {
+                range: _,
+            },
+        ) => {
             if checker.enabled(Rule::DjangoNonLeadingReceiverDecorator) {
                 flake8_django::rules::non_leading_receiver_decorator(checker, decorator_list);
             }
@@ -205,7 +207,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 flake8_bugbear::rules::cached_instance_method(checker, decorator_list);
             }
             if checker.enabled(Rule::MutableArgumentDefault) {
-                flake8_bugbear::rules::mutable_argument_default(checker, parameters, body, *range);
+                flake8_bugbear::rules::mutable_argument_default(checker, function_def);
             }
             if checker.any_enabled(&[
                 Rule::UnnecessaryReturnNone,

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

@@ -1,13 +1,12 @@
 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::helpers::is_docstring_stmt;
-use ruff_python_ast::{ParameterWithDefault, Parameters, Ranged, Stmt};
+use ruff_python_ast::{self as ast, Expr, Parameter, ParameterWithDefault, Ranged};
-use ruff_python_parser::lexer::lex_starts_at;
+use ruff_python_codegen::{Generator, Stylist};
-use ruff_python_parser::{Mode, Tok};
+use ruff_python_index::Indexer;
 use ruff_python_semantic::analyze::typing::{is_immutable_annotation, is_mutable_expr};
-use ruff_python_trivia::indentation_at_offset;
+use ruff_python_trivia::{indentation_at_offset, textwrap};
 use ruff_source_file::Locator;
-use ruff_text_size::TextRange;
 
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
@@ -63,29 +62,23 @@ impl Violation for MutableArgumentDefault {
         format!("Do not use mutable data structures for argument defaults")
     }
     fn autofix_title(&self) -> Option<String> {
-        Some(format!(
+        Some(format!("Replace with `None`; initialize within function"))
-            "Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`"
-        ))
     }
 }
 
 /// B006
-pub(crate) fn mutable_argument_default(
+pub(crate) fn mutable_argument_default(checker: &mut Checker, function_def: &ast::StmtFunctionDef) {
-    checker: &mut Checker,
-    parameters: &Parameters,
-    body: &[Stmt],
-    func_range: TextRange,
-) {
     // Scan in reverse order to right-align zip().
     for ParameterWithDefault {
         parameter,
         default,
         range: _,
-    } in parameters
+    } in function_def
+        .parameters
         .posonlyargs
         .iter()
-        .chain(&parameters.args)
+        .chain(&function_def.parameters.args)
-        .chain(&parameters.kwonlyargs)
+        .chain(&function_def.parameters.kwonlyargs)
     {
         let Some(default) = default else {
             continue;
@@ -100,50 +93,81 @@ pub(crate) fn mutable_argument_default(
             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())
+            if checker.patch(diagnostic.kind.rule()) {
-                && !is_single_line(checker.locator(), func_range, body)
+                if let Some(fix) = move_initialization(
-            {
+                    function_def,
-                // Set the default arg value to None
+                    parameter,
-                let arg_edit = Edit::range_replacement("None".to_string(), default.range());
+                    default,
-
+                    checker.locator(),
-                // Add conditional check to set the default arg to its original value if still None
+                    checker.stylist(),
-                let mut check_lines = String::new();
+                    checker.indexer(),
-                let indentation =
+                    checker.generator(),
-                    indentation_at_offset(body[0].start(), checker.locator()).unwrap_or_default();
+                ) {
-                let indentation = leading_space(indentation);
+                    diagnostic.set_fix(fix);
-                // 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", parameter.name.as_str()).as_str());
-                check_lines.push_str(indentation);
-                check_lines.push_str(checker.stylist().indentation());
-                check_lines.push_str(
-                    format!(
-                        "{} = {}",
-                        parameter.name.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 {
+/// Generate a [`Fix`] to move a mutable argument default initialization
-    let arg_string = locator.slice(func_range);
+/// into the function body.
-    for (tok, range) in lex_starts_at(arg_string, Mode::Module, func_range.start()).flatten() {
+fn move_initialization(
-        match tok {
+    function_def: &ast::StmtFunctionDef,
-            Tok::Colon => {
+    parameter: &Parameter,
-                return !locator.contains_line_break(TextRange::new(range.end(), body[0].start()));
+    default: &Expr,
-            }
+    locator: &Locator,
-            _ => continue,
+    stylist: &Stylist,
-        }
+    indexer: &Indexer,
+    generator: Generator,
+) -> Option<Fix> {
+    let mut body = function_def.body.iter();
+
+    let statement = body.next()?;
+    if indexer.preceded_by_multi_statement_line(statement, locator) {
+        return None;
     }
-    false
+
+    // Determine the indentation depth of the function body.
+    let indentation = indentation_at_offset(statement.start(), locator)?;
+
+    // Set the default argument value to `None`.
+    let default_edit = Edit::range_replacement("None".to_string(), default.range());
+
+    // Add an `if`, to set the argument to its original value if still `None`.
+    let mut content = String::new();
+    content.push_str(&format!("if {} is None:", parameter.name.as_str()));
+    content.push_str(stylist.line_ending().as_str());
+    content.push_str(stylist.indentation());
+    content.push_str(&format!(
+        "{} = {}",
+        parameter.name.as_str(),
+        generator.expr(default)
+    ));
+    content.push_str(stylist.line_ending().as_str());
+
+    // Indent the edit to match the body indentation.
+    let content = textwrap::indent(&content, indentation).to_string();
+
+    let initialization_edit = if is_docstring_stmt(statement) {
+        // If the first statement in the function is a docstring, insert _after_ it.
+        if let Some(statement) = body.next() {
+            // If there's a second statement, insert _before_ it, but ensure this isn't a
+            // multi-statement line.
+            if indexer.preceded_by_multi_statement_line(statement, locator) {
+                return None;
+            }
+            Edit::insertion(content, locator.line_start(statement.start()))
+        } else {
+            // If the docstring is the only statement, insert _before_ it.
+            Edit::insertion(content, locator.full_line_end(statement.end()))
+        }
+    } else {
+        // Otherwise, insert before the first statement.
+        let at = locator.line_start(statement.start());
+        Edit::insertion(content, at)
+    };
+
+    Some(Fix::manual_edits(default_edit, [initialization_edit]))
 }

crates/ruff/src/rules/flake8_bugbear/snapshots/ruff__rules__flake8_bugbear__tests__B006_B006_B008.py.snap

@@ -7,7 +7,7 @@ B006_B008.py:63:25: B006 [*] Do not use mutable data structures for argument def
    |                         ^^^^^^^^^ B006
 64 |     ...
    |
-   = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+   = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 60 60 | # Flag mutable literals/comprehensions
@@ -27,7 +27,7 @@ B006_B008.py:67:30: B006 [*] Do not use mutable data structures for argument def
    |                              ^^ B006
 68 |     ...
    |
-   = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+   = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 64 64 |     ...
@@ -49,7 +49,7 @@ B006_B008.py:73:52: B006 [*] Do not use mutable data structures for argument def
    |                                                    ^^ B006
 74 |         pass
    |
-   = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+   = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 70 70 | 
@@ -72,7 +72,7 @@ B006_B008.py:77:31: B006 [*] Do not use mutable data structures for argument def
    | |_^ B006
 80 |       ...
    |
-   = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+   = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 74 74 |         pass
@@ -95,7 +95,7 @@ B006_B008.py:82:36: B006 Do not use mutable data structures for argument default
 82 | def single_line_func_wrong(value = {}): ...
    |                                    ^^ B006
    |
-   = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+   = help: Replace with `None`; initialize within function
 
 B006_B008.py:85:20: B006 [*] Do not use mutable data structures for argument defaults
    |
@@ -103,7 +103,7 @@ B006_B008.py:85:20: B006 [*] Do not use mutable data structures for argument def
    |                    ^^^^^ B006
 86 |     ...
    |
-   = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+   = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 82 82 | def single_line_func_wrong(value = {}): ...
@@ -123,7 +123,7 @@ B006_B008.py:89:20: B006 [*] Do not use mutable data structures for argument def
    |                    ^^^^^^^^^^^^^^^^^^^^^^^^^ B006
 90 |     ...
    |
-   = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+   = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 86 86 |     ...
@@ -143,7 +143,7 @@ B006_B008.py:93:32: B006 [*] Do not use mutable data structures for argument def
    |                                ^^^^^^^^^^^^^^^^^^^^^^^^^ B006
 94 |     ...
    |
-   = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+   = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 90 90 |     ...
@@ -163,7 +163,7 @@ B006_B008.py:97:26: B006 [*] Do not use mutable data structures for argument def
    |                          ^^^^^^^^^^^^^^^^^^^ B006
 98 |     ...
    |
-   = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+   = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 94  94  |     ...
@@ -184,7 +184,7 @@ B006_B008.py:102:46: B006 [*] Do not use mutable data structures for argument de
     |                                              ^^^^^^^^^^^^^^^^^^^^^^^^ B006
 103 |     pass
     |
-    = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+    = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 99  99  | 
@@ -204,7 +204,7 @@ B006_B008.py:106:46: B006 [*] Do not use mutable data structures for argument de
     |                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^ B006
 107 |     pass
     |
-    = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+    = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 103 103 |     pass
@@ -224,7 +224,7 @@ B006_B008.py:110:45: B006 [*] Do not use mutable data structures for argument de
     |                                             ^^^^^^^^^^^^^^^^^^^^^^^^ B006
 111 |     pass
     |
-    = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+    = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 107 107 |     pass
@@ -244,7 +244,7 @@ B006_B008.py:114:33: B006 [*] Do not use mutable data structures for argument de
     |                                 ^^ B006
 115 |     ...
     |
-    = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+    = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 111 111 |     pass
@@ -266,7 +266,7 @@ B006_B008.py:235:20: B006 [*] Do not use mutable data structures for argument de
     |                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ B006
 236 |     pass
     |
-    = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+    = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 232 232 | 
@@ -288,7 +288,7 @@ B006_B008.py:272:27: B006 [*] Do not use mutable data structures for argument de
 273 |     b: Optional[Dict[int, int]] = {},
 274 |     c: Annotated[Union[Set[str], abc.Sized], "annotation"] = set(),
     |
-    = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+    = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 269 269 | 
@@ -303,6 +303,8 @@ B006_B008.py:272:27: B006 [*] Do not use mutable data structures for argument de
     277 |+    if a is None:
     278 |+        a = []
 277 279 |     pass
+278 280 | 
+279 281 | 
 
 B006_B008.py:273:35: B006 [*] Do not use mutable data structures for argument defaults
     |
@@ -313,7 +315,7 @@ B006_B008.py:273:35: B006 [*] Do not use mutable data structures for argument de
 274 |     c: Annotated[Union[Set[str], abc.Sized], "annotation"] = set(),
 275 |     d: typing_extensions.Annotated[Union[Set[str], abc.Sized], "annotation"] = set(),
     |
-    = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+    = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 270 270 | 
@@ -327,6 +329,8 @@ B006_B008.py:273:35: B006 [*] Do not use mutable data structures for argument de
     277 |+    if b is None:
     278 |+        b = {}
 277 279 |     pass
+278 280 | 
+279 281 | 
 
 B006_B008.py:274:62: B006 [*] Do not use mutable data structures for argument defaults
     |
@@ -337,7 +341,7 @@ B006_B008.py:274:62: B006 [*] Do not use mutable data structures for argument de
 275 |     d: typing_extensions.Annotated[Union[Set[str], abc.Sized], "annotation"] = set(),
 276 | ):
     |
-    = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+    = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 271 271 | def mutable_annotations(
@@ -350,6 +354,8 @@ B006_B008.py:274:62: B006 [*] Do not use mutable data structures for argument de
     277 |+    if c is None:
     278 |+        c = set()
 277 279 |     pass
+278 280 | 
+279 281 | 
 
 B006_B008.py:275:80: B006 [*] Do not use mutable data structures for argument defaults
     |
@@ -360,7 +366,7 @@ B006_B008.py:275:80: B006 [*] Do not use mutable data structures for argument de
 276 | ):
 277 |     pass
     |
-    = help: Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`
+    = help: Replace with `None`; initialize within function
 
 ℹ Possible fix
 272 272 |     a: list[int] | None = [],
@@ -372,5 +378,102 @@ B006_B008.py:275:80: B006 [*] Do not use mutable data structures for argument de
     277 |+    if d is None:
     278 |+        d = set()
 277 279 |     pass
+278 280 | 
+279 281 | 
+
+B006_B008.py:280:52: B006 [*] Do not use mutable data structures for argument defaults
+    |
+280 | def single_line_func_wrong(value: dict[str, str] = {}):
+    |                                                    ^^ B006
+281 |     """Docstring"""
+    |
+    = help: Replace with `None`; initialize within function
+
+ℹ Possible fix
+277 277 |     pass
+278 278 | 
+279 279 | 
+280     |-def single_line_func_wrong(value: dict[str, str] = {}):
+    280 |+def single_line_func_wrong(value: dict[str, str] = None):
+281 281 |     """Docstring"""
+    282 |+    if value is None:
+    283 |+        value = {}
+282 284 | 
+283 285 | 
+284 286 | def single_line_func_wrong(value: dict[str, str] = {}):
+
+B006_B008.py:284:52: B006 [*] Do not use mutable data structures for argument defaults
+    |
+284 | def single_line_func_wrong(value: dict[str, str] = {}):
+    |                                                    ^^ B006
+285 |     """Docstring"""
+286 |     ...
+    |
+    = help: Replace with `None`; initialize within function
+
+ℹ Possible fix
+281 281 |     """Docstring"""
+282 282 | 
+283 283 | 
+284     |-def single_line_func_wrong(value: dict[str, str] = {}):
+    284 |+def single_line_func_wrong(value: dict[str, str] = None):
+285 285 |     """Docstring"""
+    286 |+    if value is None:
+    287 |+        value = {}
+286 288 |     ...
+287 289 | 
+288 290 | 
+
+B006_B008.py:289:52: B006 Do not use mutable data structures for argument defaults
+    |
+289 | def single_line_func_wrong(value: dict[str, str] = {}):
+    |                                                    ^^ B006
+290 |     """Docstring"""; ...
+    |
+    = help: Replace with `None`; initialize within function
+
+B006_B008.py:293:52: B006 Do not use mutable data structures for argument defaults
+    |
+293 | def single_line_func_wrong(value: dict[str, str] = {}):
+    |                                                    ^^ B006
+294 |     """Docstring"""; \
+295 |         ...
+    |
+    = help: Replace with `None`; initialize within function
+
+B006_B008.py:298:52: B006 [*] Do not use mutable data structures for argument defaults
+    |
+298 |   def single_line_func_wrong(value: dict[str, str] = {
+    |  ____________________________________________________^
+299 | |     # This is a comment
+300 | | }):
+    | |_^ B006
+301 |       """Docstring"""
+    |
+    = help: Replace with `None`; initialize within function
+
+ℹ Possible fix
+295 295 |         ...
+296 296 | 
+297 297 | 
+298     |-def single_line_func_wrong(value: dict[str, str] = {
+299     |-    # This is a comment
+300     |-}):
+    298 |+def single_line_func_wrong(value: dict[str, str] = None):
+301 299 |     """Docstring"""
+    300 |+    if value is None:
+    301 |+        value = {}
+302 302 | 
+303 303 | 
+304 304 | def single_line_func_wrong(value: dict[str, str] = {}) \
+
+B006_B008.py:304:52: B006 Do not use mutable data structures for argument defaults
+    |
+304 | def single_line_func_wrong(value: dict[str, str] = {}) \
+    |                                                    ^^ B006
+305 |     : \
+306 |     """Docstring"""
+    |
+    = help: Replace with `None`; initialize within function