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`
This commit is contained in:
Charlie Marsh 2023-08-11 01:03:56 -04:00 committed by GitHub
parent cc151c35a8
commit 2e5c81b202
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 244 additions and 86 deletions

View file

@ -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"""

View file

@ -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,

View file

@ -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<String> {
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(&parameters.args)
.chain(&parameters.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<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]))
}

View file

@ -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