[pyupgrade] Implement import-replacement rule (UP035) (#2049)

README.md

@@ -847,6 +847,7 @@ For more, see [pyupgrade](https://pypi.org/project/pyupgrade/) on PyPI.
 | UP032 | f-string | Use f-string instead of `format` call | 🛠 |
 | UP033 | functools-cache | Use `@functools.cache` instead of `@functools.lru_cache(maxsize=None)` | 🛠 |
 | UP034 | extraneous-parentheses | Avoid extraneous parentheses | 🛠 |
+| UP035 | import-replacements | Import from `{module}` instead: {names} | 🛠 |
 
 ### flake8-2020 (YTT)
 

resources/test/fixtures/pyupgrade/UP035.py

@@ -0,0 +1,50 @@
+# UP035
+from collections import Mapping
+
+from collections import Mapping as MAP
+
+from collections import Mapping, Sequence
+
+from collections import Counter, Mapping
+
+from collections import (Counter, Mapping)
+
+from collections import (Counter,
+                         Mapping)
+
+from collections import Counter, \
+                         Mapping
+
+from collections import Counter, Mapping, Sequence
+
+from collections import Mapping as mapping, Counter
+
+if True:
+    from collections import Mapping, Counter
+
+if True:
+    if True:
+        pass
+    from collections import Mapping, Counter
+
+if True: from collections import Mapping
+
+import os
+from collections import Counter, Mapping
+import sys
+
+if True:
+    from collections import (
+        Mapping,
+        Callable,
+        Bad,
+        Good,
+    )
+
+from typing import Callable, Match, Pattern, List
+
+if True: from collections import (
+    Mapping, Counter)
+
+# OK
+from a import b

ruff.schema.json

@@ -1892,6 +1892,7 @@
         "UP032",
         "UP033",
         "UP034",
+        "UP035",
         "W",
         "W2",
         "W29",

src/checkers/ast.rs

@@ -1073,6 +1073,15 @@ where
                 if self.settings.rules.enabled(&Rule::RewriteCElementTree) {
                     pyupgrade::rules::replace_c_element_tree(self, stmt);
                 }
+                if self.settings.rules.enabled(&Rule::ImportReplacements) {
+                    pyupgrade::rules::import_replacements(
+                        self,
+                        stmt,
+                        names,
+                        module.as_ref().map(String::as_str),
+                        level.as_ref(),
+                    );
+                }
                 if self.settings.rules.enabled(&Rule::UnnecessaryBuiltinImport) {
                     if let Some(module) = module.as_deref() {
                         pyupgrade::rules::unnecessary_builtin_import(self, stmt, module, names);

src/registry.rs

@@ -256,6 +256,7 @@ ruff_macros::define_rule_mapping!(
     UP032 => violations::FString,
     UP033 => violations::FunctoolsCache,
     UP034 => violations::ExtraneousParentheses,
+    UP035 => rules::pyupgrade::rules::ImportReplacements,
     // pydocstyle
     D100 => violations::PublicModule,
     D101 => violations::PublicClass,

src/rules/pyupgrade/fixes.rs

@@ -2,6 +2,8 @@ use libcst_native::{
     Codegen, CodegenState, Expression, ParenthesizableWhitespace, SmallStatement, Statement,
 };
 use rustpython_ast::{Expr, Keyword, Location};
+use rustpython_parser::lexer;
+use rustpython_parser::lexer::Tok;
 
 use crate::ast::types::Range;
 use crate::autofix::helpers::remove_argument;
@@ -58,3 +60,215 @@ pub fn remove_super_arguments(locator: &Locator, stylist: &Stylist, expr: &Expr)
         range.end_location,
     ))
 }
+
+/// Remove any imports matching `members` from an import-from statement.
+pub fn remove_import_members(contents: &str, members: &[&str]) -> String {
+    let mut names: Vec<Range> = vec![];
+    let mut commas: Vec<Range> = vec![];
+    let mut removal_indices: Vec<usize> = vec![];
+
+    // Find all Tok::Name tokens that are not preceded by Tok::As, and all Tok::Comma tokens.
+    let mut prev_tok = None;
+    for (start, tok, end) in lexer::make_tokenizer(contents)
+        .flatten()
+        .skip_while(|(_, tok, _)| !matches!(tok, Tok::Import))
+    {
+        if let Tok::Name { name } = &tok {
+            if matches!(prev_tok, Some(Tok::As)) {
+                // Adjust the location to take the alias into account.
+                names.last_mut().unwrap().end_location = end;
+            } else {
+                if members.contains(&name.as_str()) {
+                    removal_indices.push(names.len());
+                }
+                names.push(Range::new(start, end));
+            }
+        } else if matches!(tok, Tok::Comma) {
+            commas.push(Range::new(start, end));
+        }
+        prev_tok = Some(tok);
+    }
+
+    // Reconstruct the source code by skipping any names that are in `members`.
+    let locator = Locator::new(contents);
+    let mut output = String::with_capacity(contents.len());
+    let mut last_pos: Location = Location::new(1, 0);
+    let mut is_first = true;
+    for index in 0..names.len() {
+        if !removal_indices.contains(&index) {
+            is_first = false;
+            continue;
+        }
+
+        let (start_location, end_location) = if is_first {
+            (names[index].location, names[index + 1].location)
+        } else {
+            (commas[index - 1].location, names[index].end_location)
+        };
+
+        // Add all contents from `last_pos` to `fix.location`.
+        // It's possible that `last_pos` is after `fix.location`, if we're removing the first _two_
+        // members.
+        if start_location > last_pos {
+            let slice = locator.slice_source_code_range(&Range::new(last_pos, start_location));
+            output.push_str(slice);
+        }
+
+        last_pos = end_location;
+    }
+
+    // Add the remaining content.
+    let slice = locator.slice_source_code_at(last_pos);
+    output.push_str(slice);
+    output
+}
+
+#[cfg(test)]
+mod test {
+    use crate::rules::pyupgrade::fixes::remove_import_members;
+
+    #[test]
+    fn once() {
+        let source = r#"from foo import bar, baz, bop, qux as q"#;
+        let expected = r#"from foo import bar, baz, qux as q"#;
+        let actual = remove_import_members(source, &["bop"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn twice() {
+        let source = r#"from foo import bar, baz, bop, qux as q"#;
+        let expected = r#"from foo import bar, qux as q"#;
+        let actual = remove_import_members(source, &["baz", "bop"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn aliased() {
+        let source = r#"from foo import bar, baz, bop as boop, qux as q"#;
+        let expected = r#"from foo import bar, baz, qux as q"#;
+        let actual = remove_import_members(source, &["bop"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn parenthesized() {
+        let source = r#"from foo import (bar, baz, bop, qux as q)"#;
+        let expected = r#"from foo import (bar, baz, qux as q)"#;
+        let actual = remove_import_members(source, &["bop"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn last_import() {
+        let source = r#"from foo import bar, baz, bop, qux as q"#;
+        let expected = r#"from foo import bar, baz, bop"#;
+        let actual = remove_import_members(source, &["qux"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn first_import() {
+        let source = r#"from foo import bar, baz, bop, qux as q"#;
+        let expected = r#"from foo import baz, bop, qux as q"#;
+        let actual = remove_import_members(source, &["bar"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn first_two_imports() {
+        let source = r#"from foo import bar, baz, bop, qux as q"#;
+        let expected = r#"from foo import bop, qux as q"#;
+        let actual = remove_import_members(source, &["bar", "baz"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn first_two_imports_multiline() {
+        let source = r#"from foo import (
+    bar,
+    baz,
+    bop,
+    qux as q
+)"#;
+        let expected = r#"from foo import (
+    bop,
+    qux as q
+)"#;
+        let actual = remove_import_members(source, &["bar", "baz"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn multiline_once() {
+        let source = r#"from foo import (
+    bar,
+    baz,
+    bop,
+    qux as q,
+)"#;
+        let expected = r#"from foo import (
+    bar,
+    baz,
+    qux as q,
+)"#;
+        let actual = remove_import_members(source, &["bop"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn multiline_twice() {
+        let source = r#"from foo import (
+    bar,
+    baz,
+    bop,
+    qux as q,
+)"#;
+        let expected = r#"from foo import (
+    bar,
+    qux as q,
+)"#;
+        let actual = remove_import_members(source, &["baz", "bop"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn multiline_comment() {
+        let source = r#"from foo import (
+    bar,
+    baz,
+    # This comment should be removed.
+    bop,
+    # This comment should be retained.
+    qux as q,
+)"#;
+        let expected = r#"from foo import (
+    bar,
+    baz,
+    # This comment should be retained.
+    qux as q,
+)"#;
+        let actual = remove_import_members(source, &["bop"]);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn multi_comment_first_import() {
+        let source = r#"from foo import (
+    # This comment should be retained.
+    bar,
+    # This comment should be removed.
+    baz,
+    bop,
+    qux as q,
+)"#;
+        let expected = r#"from foo import (
+    # This comment should be retained.
+    baz,
+    bop,
+    qux as q,
+)"#;
+        let actual = remove_import_members(source, &["bar"]);
+        assert_eq!(expected, actual);
+    }
+}

src/rules/pyupgrade/mod.rs

@@ -57,6 +57,7 @@ mod tests {
     #[test_case(Rule::FString, Path::new("UP032.py"); "UP032")]
     #[test_case(Rule::FunctoolsCache, Path::new("UP033.py"); "UP033")]
     #[test_case(Rule::ExtraneousParentheses, Path::new("UP034.py"); "UP034")]
+    #[test_case(Rule::ImportReplacements, Path::new("UP035.py"); "UP035")]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.code(), path.to_string_lossy());
         let diagnostics = test_path(