add fixes for __future__ import removal (#682)

README.md

@@ -441,8 +441,8 @@ For more, see [pyupgrade](https://pypi.org/project/pyupgrade/3.2.0/) on PyPI.
 | U007 | UsePEP604Annotation | Use `X \| Y` for type annotations | 🛠 |
 | U008 | SuperCallWithParameters | Use `super()` instead of `super(__class__, self)` | 🛠 |
 | U009 | PEP3120UnnecessaryCodingComment | utf-8 encoding declaration is unnecessary | 🛠 |
-| U010 | UnnecessaryFutureImport | Unnessary __future__ import `...` for target Python version |  |
+| U010 | UnnecessaryFutureImport | Unnecessary `__future__` import `...` for target Python version | 🛠 |
-| U011 | UnnecessaryLRUCacheParams | Unnessary parameters to functools.lru_cache | 🛠 |
+| U011 | UnnecessaryLRUCacheParams | Unnecessary parameters to functools.lru_cache | 🛠 |
 
 ### pep8-naming
 

resources/test/fixtures/U010.py

@@ -1,5 +1,13 @@
-from __future__ import annotations, nested_scopes, generators
+from __future__ import nested_scopes, generators
-
+from __future__ import with_statement, unicode_literals
 from __future__ import absolute_import, division
-
 from __future__ import generator_stop
+from __future__ import print_function, generator_stop
+from __future__ import invalid_module, generators
+
+if True:
+    from __future__ import generator_stop
+
+if True:
+    from __future__ import generator_stop
+    from __future__ import invalid_module, generators

src/check_ast.rs

@@ -606,6 +606,12 @@ where
                     }
                 }
 
+                if let Some("__future__") = module.as_deref() {
+                    if self.settings.enabled.contains(&CheckCode::U010) {
+                        pyupgrade::plugins::unnecessary_future_import(self, stmt, names);
+                    }
+                }
+
                 for alias in names {
                     if let Some("__future__") = module.as_deref() {
                         let name = alias.node.asname.as_ref().unwrap_or(&alias.node.name);
@@ -638,14 +644,6 @@ where
                             }
                         }
 
-                        if self.settings.enabled.contains(&CheckCode::U010) {
-                            pyupgrade::plugins::unnecessary_future_import(
-                                self,
-                                stmt,
-                                &alias.node.name,
-                            );
-                        }
-
                         if self.settings.enabled.contains(&CheckCode::F404) && !self.futures_allowed
                         {
                             self.add_check(Check::new(

src/checks.rs

@@ -442,7 +442,7 @@ pub enum CheckKind {
     UsePEP604Annotation,
     SuperCallWithParameters,
     PEP3120UnnecessaryCodingComment,
-    UnnecessaryFutureImport(String),
+    UnnecessaryFutureImport(Vec<String>),
     UnnecessaryLRUCacheParams,
     // pydocstyle
     BlankLineAfterLastSection(String),
@@ -689,7 +689,7 @@ impl CheckCode {
             CheckCode::U007 => CheckKind::UsePEP604Annotation,
             CheckCode::U008 => CheckKind::SuperCallWithParameters,
             CheckCode::U009 => CheckKind::PEP3120UnnecessaryCodingComment,
-            CheckCode::U010 => CheckKind::UnnecessaryFutureImport("...".to_string()),
+            CheckCode::U010 => CheckKind::UnnecessaryFutureImport(vec!["...".to_string()]),
             CheckCode::U011 => CheckKind::UnnecessaryLRUCacheParams,
             // pydocstyle
             CheckCode::D100 => CheckKind::PublicModule,
@@ -1595,11 +1595,17 @@ impl CheckKind {
             CheckKind::SuperCallWithParameters => {
                 "Use `super()` instead of `super(__class__, self)`".to_string()
             }
-            CheckKind::UnnecessaryFutureImport(name) => {
+            CheckKind::UnnecessaryFutureImport(names) => {
-                format!("Unnessary __future__ import `{name}` for target Python version")
+                if names.len() == 1 {
+                    let import = &names[0];
+                    format!("Unnecessary `__future__` import `{import}` for target Python version")
+                } else {
+                    let imports = names.iter().map(|name| format!("`{name}`")).join(", ");
+                    format!("Unnecessary `__future__` imports {imports} for target Python version")
+                }
             }
             CheckKind::UnnecessaryLRUCacheParams => {
-                "Unnessary parameters to functools.lru_cache".to_string()
+                "Unnecessary parameters to functools.lru_cache".to_string()
             }
             // pydocstyle
             CheckKind::FitsOnOneLine => "One-line docstring should fit on one line".to_string(),
@@ -1867,6 +1873,7 @@ impl CheckKind {
                 | CheckKind::UnnecessaryAbspath
                 | CheckKind::UnnecessaryCollectionCall(_)
                 | CheckKind::UnnecessaryComprehension(_)
+                | CheckKind::UnnecessaryFutureImport(_)
                 | CheckKind::UnnecessaryGeneratorDict
                 | CheckKind::UnnecessaryGeneratorList
                 | CheckKind::UnnecessaryGeneratorSet

src/pyupgrade/checks.rs

@@ -9,29 +9,6 @@ use crate::checks::{Check, CheckKind};
 use crate::pyupgrade::types::Primitive;
 use crate::settings::types::PythonVersion;
 
-pub const PY33_PLUS_REMOVE_FUTURES: &[&str] = &[
-    "nested_scopes",
-    "generators",
-    "with_statement",
-    "division",
-    "absolute_import",
-    "with_statement",
-    "print_function",
-    "unicode_literals",
-];
-
-pub const PY37_PLUS_REMOVE_FUTURES: &[&str] = &[
-    "nested_scopes",
-    "generators",
-    "with_statement",
-    "division",
-    "absolute_import",
-    "with_statement",
-    "print_function",
-    "unicode_literals",
-    "generator_stop",
-];
-
 /// U008
 pub fn super_args(
     scope: &Scope,
@@ -183,23 +160,6 @@ pub fn type_of_primitive(func: &Expr, args: &[Expr], location: Range) -> Option<
     None
 }
 
-/// U010
-pub fn unnecessary_future_import(
-    version: PythonVersion,
-    name: &str,
-    location: Range,
-) -> Option<Check> {
-    if (version >= PythonVersion::Py33 && PY33_PLUS_REMOVE_FUTURES.contains(&name))
-        || (version >= PythonVersion::Py37 && PY37_PLUS_REMOVE_FUTURES.contains(&name))
-    {
-        return Some(Check::new(
-            CheckKind::UnnecessaryFutureImport(name.to_string()),
-            location,
-        ));
-    }
-    None
-}
-
 /// U011
 pub fn unnecessary_lru_cache_params(
     decorator_list: &[Expr],

src/pyupgrade/fixes.rs

@@ -1,11 +1,13 @@
-use libcst_native::{Codegen, Expression, SmallStatement, Statement};
+use anyhow::Result;
-use rustpython_ast::{Expr, Keyword, Location};
+use libcst_native::{Codegen, Expression, ImportNames, SmallStatement, Statement};
+use rustpython_ast::{Expr, Keyword, Location, Stmt};
 use rustpython_parser::lexer;
 use rustpython_parser::lexer::Tok;
 
 use crate::ast::helpers;
 use crate::ast::types::Range;
-use crate::autofix::Fix;
+use crate::autofix::{self, Fix};
+use crate::cst::matchers::match_module;
 use crate::source_code_locator::SourceCodeLocator;
 
 /// Generate a fix to remove a base from a ClassDef statement.
@@ -41,7 +43,7 @@ pub fn remove_class_def_base(
         }
 
         return match (fix_start, fix_end) {
-            (Some(start), Some(end)) => Some(Fix::replacement("".to_string(), start, end)),
+            (Some(start), Some(end)) => Some(Fix::deletion(start, end)),
             _ => None,
         };
     }
@@ -133,6 +135,63 @@ pub fn remove_super_arguments(locator: &SourceCodeLocator, expr: &Expr) -> Optio
     None
 }
 
+/// U010
+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_froms.
+    let module_text = locator.slice_source_code_range(&Range::from_located(stmt));
+    let mut tree = match_module(&module_text)?;
+
+    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::ImportFrom(body)) = body.body.first_mut() {
+        body
+    } else {
+        return Err(anyhow::anyhow!(
+            "Expected node to be: SmallStatement::ImportFrom"
+        ));
+    };
+
+    let aliases = if let ImportNames::Aliases(aliases) = &mut body.names {
+        aliases
+    } else {
+        return Err(anyhow::anyhow!("Expected node to be: 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);
+    }
+
+    if let Some(alias) = aliases.last_mut() {
+        alias.comma = trailing_comma;
+    }
+
+    if aliases.is_empty() {
+        autofix::helpers::remove_stmt(stmt, parent, deleted)
+    } else {
+        let mut state = Default::default();
+        tree.codegen(&mut state);
+
+        Ok(Fix::replacement(
+            state.to_string(),
+            stmt.location,
+            stmt.end_location.unwrap(),
+        ))
+    }
+}
+
 /// U011
 pub fn remove_unnecessary_lru_cache_params(
     locator: &SourceCodeLocator,

src/pyupgrade/plugins/unnecessary_future_import.rs

@@ -1,15 +1,77 @@
+use std::collections::BTreeSet;
+
+use rustpython_ast::{AliasData, Located};
 use rustpython_parser::ast::Stmt;
 
 use crate::ast::types::Range;
 use crate::check_ast::Checker;
-use crate::pyupgrade::checks;
+use crate::checks::{Check, CheckKind};
+use crate::pyupgrade::fixes;
+use crate::settings::types::PythonVersion;
 
-pub fn unnecessary_future_import(checker: &mut Checker, stmt: &Stmt, name: &str) {
+const PY33_PLUS_REMOVE_FUTURES: &[&str] = &[
-    if let Some(check) = checks::unnecessary_future_import(
+    "nested_scopes",
-        checker.settings.target_version,
+    "generators",
-        name,
+    "with_statement",
+    "division",
+    "absolute_import",
+    "with_statement",
+    "print_function",
+    "unicode_literals",
+];
+
+const PY37_PLUS_REMOVE_FUTURES: &[&str] = &[
+    "nested_scopes",
+    "generators",
+    "with_statement",
+    "division",
+    "absolute_import",
+    "with_statement",
+    "print_function",
+    "unicode_literals",
+    "generator_stop",
+];
+
+/// U010
+pub fn unnecessary_future_import(checker: &mut Checker, stmt: &Stmt, names: &[Located<AliasData>]) {
+    let target_version = checker.settings.target_version;
+
+    let mut removable_index: Vec<usize> = vec![];
+    let mut removable_names: BTreeSet<&str> = BTreeSet::new();
+    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))
+        {
+            removable_index.push(index);
+            removable_names.insert(name);
+        }
+    }
+
+    if !removable_index.is_empty() {
+        let mut check = Check::new(
+            CheckKind::UnnecessaryFutureImport(
+                removable_names.into_iter().map(String::from).collect(),
+            ),
             Range::from_located(stmt),
+        );
+        if checker.patch() {
+            let context = checker.binding_context();
+            let deleted: Vec<&Stmt> = checker
+                .deletions
+                .iter()
+                .map(|index| checker.parents[*index])
+                .collect();
+            if let Ok(fix) = fixes::remove_unnecessary_future_import(
+                checker.locator,
+                &removable_index,
+                checker.parents[context.defined_by],
+                context.defined_in.map(|index| checker.parents[index]),
+                &deleted,
             ) {
+                check.amend(fix);
+            }
+        }
         checker.add_check(check);
     }
 }

src/snapshots/ruff__linter__tests__U010_U010.py.snap

@@ -3,48 +3,178 @@ source: src/linter.rs
 expression: checks
 ---
 - kind:
-    UnnecessaryFutureImport: nested_scopes
+    UnnecessaryFutureImport:
+      - generators
+      - nested_scopes
   location:
     row: 1
     column: 0
   end_location:
     row: 1
-    column: 61
+    column: 48
-  fix: ~
+  fix:
-- kind:
+    patch:
-    UnnecessaryFutureImport: generators
+      content: ""
       location:
         row: 1
         column: 0
       end_location:
-    row: 1
+        row: 2
-    column: 61
+        column: 0
-  fix: ~
+    applied: false
 - kind:
-    UnnecessaryFutureImport: absolute_import
+    UnnecessaryFutureImport:
+      - unicode_literals
+      - with_statement
+  location:
+    row: 2
+    column: 0
+  end_location:
+    row: 2
+    column: 55
+  fix:
+    patch:
+      content: ""
+      location:
+        row: 2
+        column: 0
+      end_location:
+        row: 3
+        column: 0
+    applied: false
+- kind:
+    UnnecessaryFutureImport:
+      - absolute_import
+      - division
   location:
     row: 3
     column: 0
   end_location:
     row: 3
     column: 48
-  fix: ~
+  fix:
-- kind:
+    patch:
-    UnnecessaryFutureImport: division
+      content: ""
       location:
         row: 3
         column: 0
       end_location:
-    row: 3
+        row: 4
-    column: 48
+        column: 0
-  fix: ~
+    applied: false
 - kind:
-    UnnecessaryFutureImport: generator_stop
+    UnnecessaryFutureImport:
+      - generator_stop
   location:
-    row: 5
+    row: 4
     column: 0
   end_location:
-    row: 5
+    row: 4
     column: 37
-  fix: ~
+  fix:
+    patch:
+      content: ""
+      location:
+        row: 4
+        column: 0
+      end_location:
+        row: 5
+        column: 0
+    applied: false
+- kind:
+    UnnecessaryFutureImport:
+      - generator_stop
+      - print_function
+  location:
+    row: 5
+    column: 0
+  end_location:
+    row: 5
+    column: 53
+  fix:
+    patch:
+      content: ""
+      location:
+        row: 5
+        column: 0
+      end_location:
+        row: 6
+        column: 0
+    applied: false
+- kind:
+    UnnecessaryFutureImport:
+      - generators
+  location:
+    row: 6
+    column: 0
+  end_location:
+    row: 6
+    column: 49
+  fix:
+    patch:
+      content: from __future__ import invalid_module
+      location:
+        row: 6
+        column: 0
+      end_location:
+        row: 6
+        column: 49
+    applied: false
+- kind:
+    UnnecessaryFutureImport:
+      - generator_stop
+  location:
+    row: 9
+    column: 4
+  end_location:
+    row: 9
+    column: 41
+  fix:
+    patch:
+      content: pass
+      location:
+        row: 9
+        column: 4
+      end_location:
+        row: 9
+        column: 41
+    applied: false
+- kind:
+    UnnecessaryFutureImport:
+      - generator_stop
+  location:
+    row: 12
+    column: 4
+  end_location:
+    row: 12
+    column: 41
+  fix:
+    patch:
+      content: ""
+      location:
+        row: 12
+        column: 0
+      end_location:
+        row: 13
+        column: 0
+    applied: false
+- kind:
+    UnnecessaryFutureImport:
+      - generators
+  location:
+    row: 13
+    column: 4
+  end_location:
+    row: 13
+    column: 53
+  fix:
+    patch:
+      content: from __future__ import invalid_module
+      location:
+        row: 13
+        column: 4
+      end_location:
+        row: 13
+        column: 53
+    applied: false