Ensure ast::whitespace::indentation extracts whitespace

Co-authored-by: Charlie Marsh <charlie.r.marsh@gmail.com>
Signed-off-by: Anders Kaseorg <andersk@mit.edu>
This commit is contained in:
Anders Kaseorg 2023-01-18 01:55:41 -05:00 committed by Charlie Marsh
parent 715ea2d374
commit 462d81beb7
4 changed files with 76 additions and 57 deletions

View File

@ -7,12 +7,17 @@ use crate::ast::types::Range;
use crate::source_code::Locator;
/// Extract the leading indentation from a line.
pub fn indentation<'a, T>(locator: &'a Locator, located: &'a Located<T>) -> Cow<'a, str> {
pub fn indentation<'a, T>(locator: &'a Locator, located: &'a Located<T>) -> Option<Cow<'a, str>> {
let range = Range::from_located(located);
locator.slice_source_code_range(&Range::new(
let indentation = locator.slice_source_code_range(&Range::new(
Location::new(range.location.row(), 0),
Location::new(range.location.row(), range.location.column()),
))
));
if indentation.chars().all(char::is_whitespace) {
Some(indentation)
} else {
None
}
}
/// Extract the leading words from a line of text.

View File

@ -107,14 +107,16 @@ fn implicit_return(checker: &mut Checker, last_stmt: &Stmt) {
let mut diagnostic =
Diagnostic::new(violations::ImplicitReturn, Range::from_located(last_stmt));
if checker.patch(&RuleCode::RET503) {
let mut content = String::new();
content.push_str(&indentation(&checker.locator, last_stmt));
content.push_str("return None");
content.push('\n');
diagnostic.amend(Fix::insertion(
content,
Location::new(last_stmt.end_location.unwrap().row() + 1, 0),
));
if let Some(indent) = indentation(checker.locator, last_stmt) {
let mut content = String::new();
content.push_str(&indent);
content.push_str("return None");
content.push('\n');
diagnostic.amend(Fix::insertion(
content,
Location::new(last_stmt.end_location.unwrap().row() + 1, 0),
));
}
}
checker.diagnostics.push(diagnostic);
}

View File

@ -6,6 +6,7 @@ use crate::ast::whitespace::indentation;
use crate::checkers::ast::Checker;
use crate::fix::Fix;
use crate::registry::Diagnostic;
use crate::source_code::Locator;
use crate::violations;
#[derive(Debug)]
@ -41,6 +42,42 @@ fn extract_middle(contents: &str) -> Option<MiddleContent> {
})
}
/// Generate a [`Fix`] for a `stdout` and `stderr` [`Keyword`] pair.
fn generate_fix(locator: &Locator, stdout: &Keyword, stderr: &Keyword) -> Option<Fix> {
let first = if stdout.location < stderr.location {
stdout
} else {
stderr
};
let last = if stdout.location > stderr.location {
stdout
} else {
stderr
};
let mut contents = String::from("capture_output=True");
if let Some(middle) = extract_middle(
&locator.slice_source_code_range(&Range::new(first.end_location.unwrap(), last.location)),
) {
if middle.multi_line {
let Some(indent) = indentation(locator, first) else {
return None;
};
contents.push(',');
contents.push('\n');
contents.push_str(&indent);
} else {
contents.push(',');
contents.push(' ');
}
contents.push_str(middle.contents);
}
Some(Fix::replacement(
contents,
first.location,
last.end_location.unwrap(),
))
}
/// UP022
pub fn replace_stdout_stderr(checker: &mut Checker, expr: &Expr, kwargs: &[Keyword]) {
if checker
@ -69,38 +106,9 @@ pub fn replace_stdout_stderr(checker: &mut Checker, expr: &Expr, kwargs: &[Keywo
let mut diagnostic =
Diagnostic::new(violations::ReplaceStdoutStderr, Range::from_located(expr));
if checker.patch(diagnostic.kind.code()) {
let first = if stdout.location < stderr.location {
stdout
} else {
stderr
if let Some(fix) = generate_fix(checker.locator, stdout, stderr) {
diagnostic.amend(fix);
};
let last = if stdout.location > stderr.location {
stdout
} else {
stderr
};
let mut contents = String::from("capture_output=True");
if let Some(middle) =
extract_middle(&checker.locator.slice_source_code_range(&Range::new(
first.end_location.unwrap(),
last.location,
)))
{
if middle.multi_line {
contents.push(',');
contents.push('\n');
contents.push_str(&indentation(&checker.locator, first));
} else {
contents.push(',');
contents.push(' ');
}
contents.push_str(middle.contents);
}
diagnostic.amend(Fix::replacement(
contents,
first.location,
last.end_location.unwrap(),
));
}
checker.diagnostics.push(diagnostic);
}

View File

@ -227,13 +227,16 @@ pub fn rewrite_mock_import(checker: &mut Checker, stmt: &Stmt) {
{
// Generate the fix, if needed, which is shared between all `mock` imports.
let content = if checker.patch(&RuleCode::UP026) {
let indent = indentation(&checker.locator, stmt);
match format_import(stmt, &indent, checker.locator, checker.stylist) {
Ok(content) => Some(content),
Err(e) => {
error!("Failed to rewrite `mock` import: {e}");
None
if let Some(indent) = indentation(checker.locator, stmt) {
match format_import(stmt, &indent, checker.locator, checker.stylist) {
Ok(content) => Some(content),
Err(e) => {
error!("Failed to rewrite `mock` import: {e}");
None
}
}
} else {
None
}
} else {
None
@ -273,16 +276,17 @@ pub fn rewrite_mock_import(checker: &mut Checker, stmt: &Stmt) {
Range::from_located(stmt),
);
if checker.patch(&RuleCode::UP026) {
let indent = indentation(&checker.locator, stmt);
match format_import_from(stmt, &indent, checker.locator, checker.stylist) {
Ok(content) => {
diagnostic.amend(Fix::replacement(
content,
stmt.location,
stmt.end_location.unwrap(),
));
if let Some(indent) = indentation(checker.locator, stmt) {
match format_import_from(stmt, &indent, checker.locator, checker.stylist) {
Ok(content) => {
diagnostic.amend(Fix::replacement(
content,
stmt.location,
stmt.end_location.unwrap(),
));
}
Err(e) => error!("Failed to rewrite `mock` import: {e}"),
}
Err(e) => error!("Failed to rewrite `mock` import: {e}"),
}
}
checker.diagnostics.push(diagnostic);