From 2e5c81b202360ae4b33f989a03f056d7963c202e Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 11 Aug 2023 01:03:56 -0400 Subject: [PATCH] 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` --- .../test/fixtures/flake8_bugbear/B006_B008.py | 29 ++++ .../src/checkers/ast/analyze/statement.rs | 24 +-- .../rules/mutable_argument_default.rs | 138 ++++++++++------- ...ke8_bugbear__tests__B006_B006_B008.py.snap | 139 +++++++++++++++--- 4 files changed, 244 insertions(+), 86 deletions(-) diff --git a/crates/ruff/resources/test/fixtures/flake8_bugbear/B006_B008.py b/crates/ruff/resources/test/fixtures/flake8_bugbear/B006_B008.py index 71146ea184..4b3492bb4c 100644 --- a/crates/ruff/resources/test/fixtures/flake8_bugbear/B006_B008.py +++ b/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""" diff --git a/crates/ruff/src/checkers/ast/analyze/statement.rs b/crates/ruff/src/checkers/ast/analyze/statement.rs index 25b96d7126..a5c56291ee 100644 --- a/crates/ruff/src/checkers/ast/analyze/statement.rs +++ b/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 { - is_async, - name, - decorator_list, - returns, - parameters, - body, - type_params, - range, - }) => { + Stmt::FunctionDef( + function_def @ ast::StmtFunctionDef { + is_async, + name, + decorator_list, + returns, + parameters, + body, + 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, diff --git a/crates/ruff/src/rules/flake8_bugbear/rules/mutable_argument_default.rs b/crates/ruff/src/rules/flake8_bugbear/rules/mutable_argument_default.rs index d6de9e965a..bee4481a28 100644 --- a/crates/ruff/src/rules/flake8_bugbear/rules/mutable_argument_default.rs +++ b/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::{ParameterWithDefault, Parameters, Ranged, Stmt}; -use ruff_python_parser::lexer::lex_starts_at; -use ruff_python_parser::{Mode, Tok}; +use ruff_python_ast::helpers::is_docstring_stmt; +use ruff_python_ast::{self as ast, Expr, Parameter, ParameterWithDefault, Ranged}; +use ruff_python_codegen::{Generator, Stylist}; +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 { - Some(format!( - "Replace mutable data structure with `None` in argument default and replace it with data structure inside the function if still `None`" - )) + Some(format!("Replace with `None`; initialize within function")) } } /// B006 -pub(crate) fn mutable_argument_default( - checker: &mut Checker, - parameters: &Parameters, - body: &[Stmt], - func_range: TextRange, -) { +pub(crate) fn mutable_argument_default(checker: &mut Checker, function_def: &ast::StmtFunctionDef) { // Scan in reverse order to right-align zip(). for ParameterWithDefault { parameter, default, range: _, - } in parameters + } in function_def + .parameters .posonlyargs .iter() - .chain(¶meters.args) - .chain(¶meters.kwonlyargs) + .chain(&function_def.parameters.args) + .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()) - && !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(body[0].start(), checker.locator()).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", 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])); + if checker.patch(diagnostic.kind.rule()) { + if let Some(fix) = move_initialization( + function_def, + parameter, + default, + checker.locator(), + checker.stylist(), + checker.indexer(), + checker.generator(), + ) { + diagnostic.set_fix(fix); + } } 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, - } +/// Generate a [`Fix`] to move a mutable argument default initialization +/// into the function body. +fn move_initialization( + function_def: &ast::StmtFunctionDef, + parameter: &Parameter, + default: &Expr, + locator: &Locator, + stylist: &Stylist, + indexer: &Indexer, + generator: Generator, +) -> Option { + 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])) } diff --git a/crates/ruff/src/rules/flake8_bugbear/snapshots/ruff__rules__flake8_bugbear__tests__B006_B006_B008.py.snap b/crates/ruff/src/rules/flake8_bugbear/snapshots/ruff__rules__flake8_bugbear__tests__B006_B006_B008.py.snap index dd5f00fbec..038b21f9a2 100644 --- a/crates/ruff/src/rules/flake8_bugbear/snapshots/ruff__rules__flake8_bugbear__tests__B006_B006_B008.py.snap +++ b/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