DRY up unused import removal code (#1643)

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<Item = &'a str>,
+    stmt: &Stmt,
+    parent: Option<&Stmt>,
+    deleted: &[&Stmt],
+    locator: &SourceCodeLocator,
+) -> Result<Fix> {
+    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;

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 { .. })

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<Fix> {
-    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],

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<Fix> {
     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<Fix> {
-    // 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(),
-        ))
-    }
-}