diff --git a/src/ast/whitespace.rs b/src/ast/whitespace.rs index 0a2c06ec78..19111f87f0 100644 --- a/src/ast/whitespace.rs +++ b/src/ast/whitespace.rs @@ -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) -> Cow<'a, str> { +pub fn indentation<'a, T>(locator: &'a Locator, located: &'a Located) -> Option> { 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. diff --git a/src/rules/flake8_return/rules.rs b/src/rules/flake8_return/rules.rs index f2d32666c0..27e09e610c 100644 --- a/src/rules/flake8_return/rules.rs +++ b/src/rules/flake8_return/rules.rs @@ -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); } diff --git a/src/rules/pyupgrade/rules/replace_stdout_stderr.rs b/src/rules/pyupgrade/rules/replace_stdout_stderr.rs index baec5c008b..c081b28dab 100644 --- a/src/rules/pyupgrade/rules/replace_stdout_stderr.rs +++ b/src/rules/pyupgrade/rules/replace_stdout_stderr.rs @@ -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 { }) } +/// Generate a [`Fix`] for a `stdout` and `stderr` [`Keyword`] pair. +fn generate_fix(locator: &Locator, stdout: &Keyword, stderr: &Keyword) -> Option { + 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); } diff --git a/src/rules/pyupgrade/rules/rewrite_mock_import.rs b/src/rules/pyupgrade/rules/rewrite_mock_import.rs index 1185dabc2b..105769e216 100644 --- a/src/rules/pyupgrade/rules/rewrite_mock_import.rs +++ b/src/rules/pyupgrade/rules/rewrite_mock_import.rs @@ -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);