diff --git a/resources/test/fixtures/F401.py b/resources/test/fixtures/F401.py index 6f4e653197..15b15ed220 100644 --- a/resources/test/fixtures/F401.py +++ b/resources/test/fixtures/F401.py @@ -1,6 +1,5 @@ from __future__ import all_feature_names -import os -import functools +import functools, os from datetime import datetime from collections import ( Counter, diff --git a/src/autofix/fixes.rs b/src/autofix/fixes.rs index 9bb9f0426b..5c4b8ea617 100644 --- a/src/autofix/fixes.rs +++ b/src/autofix/fixes.rs @@ -251,14 +251,75 @@ fn remove_stmt(stmt: &Stmt, parent: Option<&Stmt>, deleted: &[&Stmt]) -> Result< } /// Generate a Fix to remove any unused imports from an `import` statement. -pub fn remove_unused_imports(stmt: &Stmt, parent: Option<&Stmt>, deleted: &[&Stmt]) -> Result { - remove_stmt(stmt, parent, deleted) +pub fn remove_unused_imports( + locator: &mut SourceCodeLocator, + full_names: &[&str], + stmt: &Stmt, + parent: Option<&Stmt>, + deleted: &[&Stmt], +) -> Result { + let mut tree = match libcst_native::parse_module( + locator.slice_source_code_range(&Range::from_located(stmt)), + None, + ) { + Ok(m) => m, + Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")), + }; + + let body = if let Some(Statement::Simple(body)) = tree.body.first_mut() { + body + } else { + return Err(anyhow::anyhow!("Expected node to be: Statement::Simple.")); + }; + let body = if let Some(SmallStatement::Import(body)) = body.body.first_mut() { + body + } else { + return Err(anyhow::anyhow!( + "Expected node to be: SmallStatement::ImportFrom." + )); + }; + let aliases = &mut body.names; + + // Preserve the trailing comma (or not) from the last entry. + let trailing_comma = aliases.last().and_then(|alias| alias.comma.clone()); + + // Identify unused imports from within the `import from`. + let mut removable = vec![]; + for (index, alias) in aliases.iter().enumerate() { + if let N(import_name) = &alias.name { + if full_names.contains(&import_name.value) { + removable.push(index); + } + } + } + // TODO(charlie): This is quadratic. + for index in removable.iter().rev() { + aliases.remove(*index); + } + + if let Some(alias) = aliases.last_mut() { + alias.comma = trailing_comma; + } + + if aliases.is_empty() { + remove_stmt(stmt, parent, deleted) + } else { + let mut state = Default::default(); + tree.codegen(&mut state); + + Ok(Fix { + content: state.to_string(), + location: stmt.location, + end_location: stmt.end_location, + applied: false, + }) + } } /// Generate a Fix to remove any unused imports from an `import from` statement. pub fn remove_unused_import_froms( locator: &mut SourceCodeLocator, - full_names: &[String], + full_names: &[&str], stmt: &Stmt, parent: Option<&Stmt>, deleted: &[&Stmt], @@ -301,7 +362,7 @@ pub fn remove_unused_import_froms( } else { name.value.to_string() }; - if full_names.contains(&import_name) { + if full_names.contains(&import_name.as_str()) { removable.push(index); } } diff --git a/src/check_ast.rs b/src/check_ast.rs index 1b2f5408b6..13f0e08d7a 100644 --- a/src/check_ast.rs +++ b/src/check_ast.rs @@ -1627,7 +1627,7 @@ impl<'a> Checker<'a> { if self.settings.select.contains(&CheckCode::F401) { // Collect all unused imports by location. (Multiple unused imports at the same // location indicates an `import from`.) - let mut unused: BTreeMap<(ImportKind, usize, Option), Vec> = + let mut unused: BTreeMap<(ImportKind, usize, Option), Vec<&str>> = BTreeMap::new(); for (name, binding) in scope.values.iter().rev() { @@ -1646,7 +1646,7 @@ impl<'a> Checker<'a> { context.defined_in, )) .or_insert(vec![]); - full_names.push(full_name.to_string()); + full_names.push(full_name); } BindingKind::Importation(full_name, context) | BindingKind::SubmoduleImportation(full_name, context) => { @@ -1657,7 +1657,7 @@ impl<'a> Checker<'a> { context.defined_in, )) .or_insert(vec![]); - full_names.push(full_name.to_string()); + full_names.push(full_name); } _ => {} } @@ -1680,35 +1680,19 @@ impl<'a> Checker<'a> { .map(|index| self.parents[*index]) .collect(); - match kind { - ImportKind::Import => { - match fixes::remove_unused_imports(child, parent, &deleted) { - Ok(fix) => { - if fix.content.is_empty() || fix.content == "pass" { - self.deletions.insert(defined_by); - } - check.amend(fix) - } - Err(e) => error!("Failed to fix unused imports: {}", e), - } - } - ImportKind::ImportFrom => { - match fixes::remove_unused_import_froms( - &mut self.locator, - &full_names, - child, - parent, - &deleted, - ) { - Ok(fix) => { - if fix.content.is_empty() || fix.content == "pass" { - self.deletions.insert(defined_by); - } - check.amend(fix) - } - Err(e) => error!("Failed to fix unused imports: {}", e), + let removal_fn = match kind { + ImportKind::Import => fixes::remove_unused_imports, + ImportKind::ImportFrom => fixes::remove_unused_import_froms, + }; + + match removal_fn(&mut self.locator, &full_names, child, parent, &deleted) { + Ok(fix) => { + if fix.content.is_empty() || fix.content == "pass" { + self.deletions.insert(defined_by); } + check.amend(fix) } + Err(e) => error!("Failed to fix unused imports: {}", e), } } diff --git a/src/snapshots/ruff__linter__tests__f401.snap b/src/snapshots/ruff__linter__tests__f401.snap index 44be9fa37a..47f1129106 100644 --- a/src/snapshots/ruff__linter__tests__f401.snap +++ b/src/snapshots/ruff__linter__tests__f401.snap @@ -5,120 +5,120 @@ expression: checks - kind: UnusedImport: functools location: - row: 3 + row: 2 column: 1 end_location: - row: 3 - column: 17 + row: 2 + column: 21 fix: - content: "" + content: import os location: - row: 3 + row: 2 column: 1 end_location: - row: 4 - column: 1 + row: 2 + column: 21 applied: false - kind: UnusedImport: collections.OrderedDict location: - row: 5 + row: 4 column: 1 end_location: - row: 9 + row: 8 column: 2 fix: content: "from collections import (\n Counter,\n namedtuple,\n)" location: - row: 5 + row: 4 column: 1 end_location: - row: 9 + row: 8 column: 2 applied: false - kind: UnusedImport: logging.handlers location: - row: 13 + row: 12 column: 1 end_location: - row: 13 + row: 12 column: 24 fix: - content: "" + content: import logging.handlers location: - row: 13 + row: 12 column: 1 end_location: - row: 14 - column: 1 + row: 12 + column: 24 applied: false - kind: UnusedImport: shelve location: - row: 34 + row: 33 column: 5 end_location: - row: 34 + row: 33 column: 18 fix: content: "" location: - row: 34 + row: 33 column: 1 end_location: - row: 35 + row: 34 column: 1 applied: false - kind: UnusedImport: importlib location: - row: 35 + row: 34 column: 5 end_location: - row: 35 + row: 34 column: 21 fix: content: pass location: - row: 35 + row: 34 column: 5 end_location: - row: 35 + row: 34 column: 21 applied: false - kind: UnusedImport: pathlib location: - row: 39 + row: 38 column: 5 end_location: - row: 39 + row: 38 column: 19 fix: content: "" location: - row: 39 + row: 38 column: 1 end_location: - row: 40 + row: 39 column: 1 applied: false - kind: UnusedImport: pickle location: - row: 54 + row: 53 column: 9 end_location: - row: 54 + row: 53 column: 22 fix: content: pass location: - row: 54 + row: 53 column: 9 end_location: - row: 54 + row: 53 column: 22 applied: false