diff --git a/src/autofix/helpers.rs b/src/autofix/helpers.rs index b7dc0628558d2..d187d5527c892 100644 --- a/src/autofix/helpers.rs +++ b/src/autofix/helpers.rs @@ -1,11 +1,17 @@ use anyhow::{bail, Result}; use itertools::Itertools; +use libcst_native::{ + Codegen, CodegenState, ImportNames, ParenthesizableWhitespace, SmallStatement, Statement, +}; use rustpython_parser::ast::{ExcepthandlerKind, Location, Stmt, StmtKind}; use crate::ast::helpers; use crate::ast::helpers::to_absolute; +use crate::ast::types::Range; use crate::ast::whitespace::LinesWithTrailingNewline; use crate::autofix::Fix; +use crate::cst::helpers::compose_module_path; +use crate::cst::matchers::match_module; use crate::source_code_locator::SourceCodeLocator; /// Determine if a body contains only a single statement, taking into account @@ -185,6 +191,85 @@ pub fn delete_stmt( } } +/// Generate a `Fix` to remove any unused imports from an `import` statement. +pub fn remove_unused_imports<'a>( + unused_imports: impl Iterator, + stmt: &Stmt, + parent: Option<&Stmt>, + deleted: &[&Stmt], + locator: &SourceCodeLocator, +) -> Result { + let module_text = locator.slice_source_code_range(&Range::from_located(stmt)); + let mut tree = match_module(&module_text)?; + + let Some(Statement::Simple(body)) = tree.body.first_mut() else { + bail!("Expected Statement::Simple"); + }; + + let (aliases, import_module) = match body.body.first_mut() { + Some(SmallStatement::Import(import_body)) => (&mut import_body.names, None), + Some(SmallStatement::ImportFrom(import_body)) => { + if let ImportNames::Aliases(names) = &mut import_body.names { + (names, import_body.module.as_ref()) + } else { + bail!("Expected ImportNames::Aliases") + } + } + _ => bail!("Expected SmallStatement::ImportFrom or SmallStatement::Import"), + }; + + // Preserve the trailing comma (or not) from the last entry. + let trailing_comma = aliases.last().and_then(|alias| alias.comma.clone()); + + for unused_import in unused_imports { + let alias_index = aliases.iter().position(|alias| { + let full_name = match import_module { + Some(module_name) => format!( + "{}.{}", + compose_module_path(module_name), + compose_module_path(&alias.name) + ), + None => compose_module_path(&alias.name), + }; + full_name == unused_import + }); + + if let Some(index) = alias_index { + aliases.remove(index); + } + } + + // But avoid destroying any trailing comments. + if let Some(alias) = aliases.last_mut() { + let has_comment = if let Some(comma) = &alias.comma { + match &comma.whitespace_after { + ParenthesizableWhitespace::SimpleWhitespace(_) => false, + ParenthesizableWhitespace::ParenthesizedWhitespace(whitespace) => { + whitespace.first_line.comment.is_some() + } + } + } else { + false + }; + if !has_comment { + alias.comma = trailing_comma; + } + } + + if aliases.is_empty() { + delete_stmt(stmt, parent, deleted, locator) + } else { + let mut state = CodegenState::default(); + tree.codegen(&mut state); + + Ok(Fix::replacement( + state.to_string(), + stmt.location, + stmt.end_location.unwrap(), + )) + } +} + #[cfg(test)] mod tests { use anyhow::Result; diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs index 670f843a1c49e..563acf70c6f15 100644 --- a/src/checkers/ast.rs +++ b/src/checkers/ast.rs @@ -37,7 +37,7 @@ use crate::source_code_locator::SourceCodeLocator; use crate::source_code_style::SourceCodeStyleDetector; use crate::visibility::{module_visibility, transition_scope, Modifier, Visibility, VisibleScope}; use crate::{ - docstrings, flake8_2020, flake8_annotations, flake8_bandit, flake8_blind_except, + autofix, docstrings, flake8_2020, flake8_annotations, flake8_bandit, flake8_blind_except, flake8_boolean_trap, flake8_bugbear, flake8_builtins, flake8_comprehensions, flake8_datetimez, flake8_debugger, flake8_errmsg, flake8_implicit_str_concat, flake8_import_conventions, flake8_pie, flake8_print, flake8_pytest_style, flake8_return, flake8_simplify, @@ -3892,8 +3892,10 @@ impl<'a> Checker<'a> { let fix = if !ignore_init && self.patch(&CheckCode::F401) { let deleted: Vec<&Stmt> = self.deletions.iter().map(|node| node.0).collect(); - match pyflakes::fixes::remove_unused_imports( - &unused_imports, + match autofix::helpers::remove_unused_imports( + unused_imports + .iter() + .map(|(full_name, _)| full_name.as_str()), child, parent, &deleted, @@ -3917,7 +3919,7 @@ impl<'a> Checker<'a> { let multiple = unused_imports.len() > 1; for (full_name, range) in unused_imports { let mut check = Check::new( - CheckKind::UnusedImport(full_name.clone(), ignore_init, multiple), + CheckKind::UnusedImport(full_name.to_string(), ignore_init, multiple), *range, ); if matches!(child.node, StmtKind::ImportFrom { .. }) diff --git a/src/pyflakes/fixes.rs b/src/pyflakes/fixes.rs index edce435e79a5f..531ec9140a23c 100644 --- a/src/pyflakes/fixes.rs +++ b/src/pyflakes/fixes.rs @@ -1,96 +1,13 @@ use anyhow::{bail, Result}; -use libcst_native::{ - Call, Codegen, CodegenState, Dict, DictElement, Expression, ImportNames, - ParenthesizableWhitespace, SmallStatement, Statement, -}; -use rustpython_ast::{Expr, Stmt}; +use libcst_native::{Call, Codegen, CodegenState, Dict, DictElement, Expression}; +use rustpython_ast::Expr; use crate::ast::types::Range; -use crate::autofix::{helpers, Fix}; -use crate::cst::helpers::compose_module_path; +use crate::autofix::Fix; use crate::cst::matchers::{match_expr, match_module}; use crate::python::string::strip_quotes_and_prefixes; use crate::source_code_locator::SourceCodeLocator; -/// Generate a `Fix` to remove any unused imports from an `import` statement. -pub fn remove_unused_imports( - unused_imports: &Vec<(&String, &Range)>, - stmt: &Stmt, - parent: Option<&Stmt>, - deleted: &[&Stmt], - locator: &SourceCodeLocator, -) -> Result { - let module_text = locator.slice_source_code_range(&Range::from_located(stmt)); - let mut tree = match_module(&module_text)?; - - let Some(Statement::Simple(body)) = tree.body.first_mut() else { - bail!("Expected Statement::Simple"); - }; - - let (aliases, import_module) = match body.body.first_mut() { - Some(SmallStatement::Import(import_body)) => (&mut import_body.names, None), - Some(SmallStatement::ImportFrom(import_body)) => { - if let ImportNames::Aliases(names) = &mut import_body.names { - (names, import_body.module.as_ref()) - } else { - bail!("Expected ImportNames::Aliases") - } - } - _ => bail!("Expected SmallStatement::ImportFrom or SmallStatement::Import"), - }; - - // Preserve the trailing comma (or not) from the last entry. - let trailing_comma = aliases.last().and_then(|alias| alias.comma.clone()); - - for (name_to_remove, _) in unused_imports { - let alias_index = aliases.iter().position(|alias| { - let full_name = match import_module { - Some(module_name) => format!( - "{}.{}", - compose_module_path(module_name), - compose_module_path(&alias.name) - ), - None => compose_module_path(&alias.name), - }; - &full_name.as_str() == name_to_remove - }); - - if let Some(index) = alias_index { - aliases.remove(index); - } - } - - // But avoid destroying any trailing comments. - if let Some(alias) = aliases.last_mut() { - let has_comment = if let Some(comma) = &alias.comma { - match &comma.whitespace_after { - ParenthesizableWhitespace::SimpleWhitespace(_) => false, - ParenthesizableWhitespace::ParenthesizedWhitespace(whitespace) => { - whitespace.first_line.comment.is_some() - } - } - } else { - false - }; - if !has_comment { - alias.comma = trailing_comma; - } - } - - if aliases.is_empty() { - helpers::delete_stmt(stmt, parent, deleted, locator) - } else { - let mut state = CodegenState::default(); - tree.codegen(&mut state); - - Ok(Fix::replacement( - state.to_string(), - stmt.location, - stmt.end_location.unwrap(), - )) - } -} - /// Generate a `Fix` to remove unused keys from format dict. pub fn remove_unused_format_arguments_from_dict( unused_arguments: &[&str], diff --git a/src/pyupgrade/fixes.rs b/src/pyupgrade/fixes.rs index 90f429d7255de..7692aaf26d0ad 100644 --- a/src/pyupgrade/fixes.rs +++ b/src/pyupgrade/fixes.rs @@ -1,15 +1,12 @@ -use anyhow::Result; use libcst_native::{ - Codegen, CodegenState, Expression, ImportNames, ParenthesizableWhitespace, SmallStatement, - Statement, + Codegen, CodegenState, Expression, ParenthesizableWhitespace, SmallStatement, Statement, }; -use rustpython_ast::{Expr, Keyword, Location, Stmt}; +use rustpython_ast::{Expr, Keyword, Location}; use rustpython_parser::lexer; use rustpython_parser::lexer::Tok; use crate::ast::types::Range; -use crate::autofix::{self, Fix}; -use crate::cst::matchers::match_module; +use crate::autofix::Fix; use crate::source_code_locator::SourceCodeLocator; /// Generate a fix to remove a base from a `ClassDef` statement. @@ -103,6 +100,7 @@ pub fn remove_class_def_base( } } +/// Generate a fix to remove arguments from a `super` call. pub fn remove_super_arguments(locator: &SourceCodeLocator, expr: &Expr) -> Option { let range = Range::from_located(expr); let contents = locator.slice_source_code_range(&range); @@ -132,67 +130,3 @@ pub fn remove_super_arguments(locator: &SourceCodeLocator, expr: &Expr) -> Optio range.end_location, )) } - -/// UP010 -pub fn remove_unnecessary_future_import( - locator: &SourceCodeLocator, - removable: &[usize], - stmt: &Stmt, - parent: Option<&Stmt>, - deleted: &[&Stmt], -) -> Result { - // TODO(charlie): DRY up with pyflakes::fixes::remove_unused_import_from. - let module_text = locator.slice_source_code_range(&Range::from_located(stmt)); - let mut tree = match_module(&module_text)?; - - let Some(Statement::Simple(body)) = tree.body.first_mut() else { - return Err(anyhow::anyhow!("Expected Statement::Simple")); - }; - let Some(SmallStatement::ImportFrom(body)) = body.body.first_mut() else { - return Err(anyhow::anyhow!( - "Expected SmallStatement::ImportFrom" - )); - }; - - let ImportNames::Aliases(aliases) = &mut body.names else { - return Err(anyhow::anyhow!("Expected Aliases")); - }; - - // Preserve the trailing comma (or not) from the last entry. - let trailing_comma = aliases.last().and_then(|alias| alias.comma.clone()); - - // TODO(charlie): This is quadratic. - for index in removable.iter().rev() { - aliases.remove(*index); - } - - // But avoid destroying any trailing comments. - if let Some(alias) = aliases.last_mut() { - let has_comment = if let Some(comma) = &alias.comma { - match &comma.whitespace_after { - ParenthesizableWhitespace::SimpleWhitespace(_) => false, - ParenthesizableWhitespace::ParenthesizedWhitespace(whitespace) => { - whitespace.first_line.comment.is_some() - } - } - } else { - false - }; - if !has_comment { - alias.comma = trailing_comma; - } - } - - if aliases.is_empty() { - autofix::helpers::delete_stmt(stmt, parent, deleted, locator) - } else { - let mut state = CodegenState::default(); - tree.codegen(&mut state); - - Ok(Fix::replacement( - state.to_string(), - stmt.location, - stmt.end_location.unwrap(), - )) - } -} diff --git a/src/pyupgrade/plugins/unnecessary_future_import.rs b/src/pyupgrade/plugins/unnecessary_future_import.rs index 2ddc5976f8aa8..f1d2fa69f85c7 100644 --- a/src/pyupgrade/plugins/unnecessary_future_import.rs +++ b/src/pyupgrade/plugins/unnecessary_future_import.rs @@ -1,12 +1,11 @@ use itertools::Itertools; use log::error; -use rustc_hash::FxHashSet; -use rustpython_ast::{AliasData, Located}; +use rustpython_ast::{Alias, AliasData, Located}; use rustpython_parser::ast::Stmt; use crate::ast::types::Range; +use crate::autofix; use crate::checkers::ast::Checker; -use crate::pyupgrade::fixes; use crate::registry::{Check, CheckKind}; use crate::settings::types::PythonVersion; @@ -37,26 +36,25 @@ const PY37_PLUS_REMOVE_FUTURES: &[&str] = &[ pub fn unnecessary_future_import(checker: &mut Checker, stmt: &Stmt, names: &[Located]) { let target_version = checker.settings.target_version; - let mut removable_index: Vec = vec![]; - let mut removable_names: FxHashSet<&str> = FxHashSet::default(); - for (index, alias) in names.iter().enumerate() { - let name = alias.node.name.as_str(); - if (target_version >= PythonVersion::Py33 && PY33_PLUS_REMOVE_FUTURES.contains(&name)) - || (target_version >= PythonVersion::Py37 && PY37_PLUS_REMOVE_FUTURES.contains(&name)) + let mut unused_imports: Vec<&Alias> = vec![]; + for alias in names { + if (target_version >= PythonVersion::Py33 + && PY33_PLUS_REMOVE_FUTURES.contains(&alias.node.name.as_str())) + || (target_version >= PythonVersion::Py37 + && PY37_PLUS_REMOVE_FUTURES.contains(&alias.node.name.as_str())) { - removable_index.push(index); - removable_names.insert(name); + unused_imports.push(alias); } } - if removable_index.is_empty() { + if unused_imports.is_empty() { return; } let mut check = Check::new( CheckKind::UnnecessaryFutureImport( - removable_names - .into_iter() - .map(String::from) + unused_imports + .iter() + .map(|alias| alias.node.name.to_string()) .sorted() .collect(), ), @@ -67,12 +65,16 @@ pub fn unnecessary_future_import(checker: &mut Checker, stmt: &Stmt, names: &[Lo let deleted: Vec<&Stmt> = checker.deletions.iter().map(|node| node.0).collect(); let defined_by = checker.current_stmt(); let defined_in = checker.current_stmt_parent(); - match fixes::remove_unnecessary_future_import( - checker.locator, - &removable_index, + let unused_imports: Vec = unused_imports + .iter() + .map(|alias| format!("__future__.{}", alias.node.name)) + .collect(); + match autofix::helpers::remove_unused_imports( + unused_imports.iter().map(std::string::String::as_str), defined_by.0, defined_in.map(|node| node.0), &deleted, + checker.locator, ) { Ok(fix) => { if fix.content.is_empty() || fix.content == "pass" {