Handle multi-import lines (#307)

This commit is contained in:
Charlie Marsh 2022-10-03 15:22:46 -04:00 committed by GitHub
parent 64d8e25528
commit 0966bf2c66
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 113 additions and 69 deletions

View File

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

View File

@ -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<Fix> {
remove_stmt(stmt, parent, deleted)
pub fn remove_unused_imports(
locator: &mut SourceCodeLocator,
full_names: &[&str],
stmt: &Stmt,
parent: Option<&Stmt>,
deleted: &[&Stmt],
) -> Result<Fix> {
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);
}
}

View File

@ -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<usize>), Vec<String>> =
let mut unused: BTreeMap<(ImportKind, usize, Option<usize>), 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),
}
}

View File

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