Ensure that B006 autofixes are inserted after imports

crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B006_5.py

@@ -0,0 +1,74 @@
+# Move mutable arguments below imports and docstrings
+# https://github.com/astral-sh/ruff/issues/7616
+
+
+def import_module_wrong(value: dict[str, str] = {}):
+    import os
+
+
+def import_module_with_values_wrong(value: dict[str, str] = {}):
+    import os
+
+    return 2
+
+
+def import_modules_wrong(value: dict[str, str] = {}):
+    import os
+    import sys
+    import itertools
+
+
+def from_import_module_wrong(value: dict[str, str] = {}):
+    from os import path
+
+
+def from_imports_module_wrong(value: dict[str, str] = {}):
+    from os import path
+    from sys import version_info
+
+
+def import_and_from_imports_module_wrong(value: dict[str, str] = {}):
+    import os
+    from sys import version_info
+
+
+def import_docstring_module_wrong(value: dict[str, str] = {}):
+    """Docstring"""
+    import os
+
+
+def import_module_wrong(value: dict[str, str] = {}):
+    """Docstring"""
+    import os; import sys
+
+
+def import_module_wrong(value: dict[str, str] = {}):
+    """Docstring"""
+    import os; import sys; x = 1
+
+
+def import_module_wrong(value: dict[str, str] = {}):
+    """Docstring"""
+    import os; import sys
+
+
+def import_module_wrong(value: dict[str, str] = {}):
+    import os; import sys
+
+
+def import_module_wrong(value: dict[str, str] = {}):
+    import os; import sys; x = 1
+
+
+def import_module_wrong(value: dict[str, str] = {}):
+    import os; import sys
+
+
+def import_module_wrong(value: dict[str, str] = {}): import os
+
+
+def import_module_wrong(value: dict[str, str] = {}): import os; import sys
+
+
+def import_module_wrong(value: dict[str, str] = {}): \
+    import os

crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B006_6.py

@@ -0,0 +1,5 @@
+# Import followed by whitespace with no newline
+# Same as B006_2.py, but import instead of docstring
+
+def foobar(foor, bar={}):    
+    import os                    

crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B006_7.py

@@ -0,0 +1,5 @@
+# Import with no newline
+# Same as B006_3.py, but import instead of docstring
+
+def foobar(foor, bar={}):    
+    import os

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

@@ -37,6 +37,9 @@ mod tests {
     #[test_case(Rule::MutableArgumentDefault, Path::new("B006_2.py"))]
     #[test_case(Rule::MutableArgumentDefault, Path::new("B006_3.py"))]
     #[test_case(Rule::MutableArgumentDefault, Path::new("B006_4.py"))]
+    #[test_case(Rule::MutableArgumentDefault, Path::new("B006_5.py"))]
+    #[test_case(Rule::MutableArgumentDefault, Path::new("B006_6.py"))]
+    #[test_case(Rule::MutableArgumentDefault, Path::new("B006_7.py"))]
     #[test_case(Rule::NoExplicitStacklevel, Path::new("B028.py"))]
     #[test_case(Rule::RaiseLiteral, Path::new("B016.py"))]
     #[test_case(Rule::RaiseWithoutFromInsideExcept, Path::new("B904.py"))]

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

@@ -139,16 +139,14 @@ fn move_initialization(
     indexer: &Indexer,
     generator: Generator,
 ) -> Option<Fix> {
-    let mut body = function_def.body.iter();
+    let mut body = function_def.body.iter().peekable();
 
-    let statement = body.next()?;
-    if indexer.in_multi_statement_line(statement, locator) {
+    // Avoid attempting to fix single-line functions.
+    let statement = body.peek()?;
+    if indexer.preceded_by_multi_statement_line(statement, locator) {
         return None;
     }
 
-    // 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());
 
@@ -164,34 +162,43 @@ fn move_initialization(
     ));
     content.push_str(stylist.line_ending().as_str());
 
+    // Determine the indentation depth of the function body.
+    let indentation = indentation_at_offset(statement.start(), locator)?;
+
     // 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.in_multi_statement_line(statement, locator) {
-                return None;
+    let mut content = textwrap::indent(&content, indentation).to_string();
+
+    // Find the position to insert the initialization after docstring and imports
+    let mut pos = locator.line_start(statement.start());
+    while let Some(statement) = body.next() {
+        // If the statement is a docstring or an import, insert _after_ it.
+        if is_docstring_stmt(statement)
+            || statement.is_import_stmt()
+            || statement.is_import_from_stmt()
+        {
+            if let Some(next) = body.peek() {
+                // If there's a second statement, insert _before_ it, but ensure this isn't a
+                // multi-statement line.
+                if indexer.in_multi_statement_line(statement, locator) {
+                    continue;
+                }
+                pos = locator.line_start(next.start());
+            } else if locator.full_line_end(statement.end()) == locator.text_len() {
+                // If the statement is at the end of the file, without a trailing newline, insert
+                // _after_ it with an extra newline.
+                content = format!("{}{}", stylist.line_ending().as_str(), content);
+                pos = locator.full_line_end(statement.end());
+                break;
+            } else {
+                // If this is the only statement, insert _after_ it.
+                pos = locator.full_line_end(statement.end());
+                break;
             }
-            Edit::insertion(content, locator.line_start(statement.start()))
-        } else if locator.full_line_end(statement.end()) == locator.text_len() {
-            // If the statement is at the end of the file, without a trailing newline, insert
-            // _after_ it with an extra newline.
-            Edit::insertion(
-                format!("{}{}", stylist.line_ending().as_str(), content),
-                locator.full_line_end(statement.end()),
-            )
         } else {
-            // If the docstring is the only statement, insert _after_ 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)
-    };
+            break;
+        };
+    }
 
+    let initialization_edit = Edit::insertion(content, pos);
     Some(Fix::manual_edits(default_edit, [initialization_edit]))
 }