[pyupgrade] Implement import-replacement rule (UP035)

README.md

@@ -732,6 +732,8 @@ 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 | Replace old formatting imports with their new versions | 🛠 |
+| UP036 | import-replacements-six | Replace old formatting imports with their new versions | 🛠 |
 
 ### pep8-naming (N)
 

resources/test/fixtures/pyupgrade/UP035.py

@@ -0,0 +1,47 @@
+# These SHOULD change
+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 six.moves import getcwd, StringIO
+
+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,
+    )
+
+# These SHOULD NOT change
+from a import b

resources/test/fixtures/pyupgrade/UP036.py

@@ -0,0 +1,31 @@
+# These SHOULD change
+from six.moves import tkinter, badlib, queue
+
+from six.moves import tkinter as hello, badlib, queue
+
+from six.moves import dbm_gnu as gnud
+
+if True:
+    from six.moves import queue, urllib_request
+
+from six.moves.queue import Queue
+
+from six.moves import queue
+
+from six.moves import queue, map, getcwd
+
+if True:
+    if True:
+        from six.moves import tkinter, badlib, queue
+
+if True:
+    if False:
+        pass
+    from six.moves import tkinter, badlib, queue
+
+# These should NOT change
+if True: from six.moves import get_cwd, StringIO
+
+if True:from six.moves import get_cwd, StringIO
+
+from six.moves import dbm_gnu

ruff.schema.json

@@ -1819,6 +1819,8 @@
         "UP032",
         "UP033",
         "UP034",
+        "UP035",
+        "UP036",
         "W",
         "W2",
         "W29",

src/checkers/ast.rs

@@ -1038,6 +1038,12 @@ 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);
+                }
+                if self.settings.rules.enabled(&Rule::ImportReplacementsSix) {
+                    pyupgrade::rules::import_replacements_six(self, stmt, module, names);
+                }
                 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

@@ -255,6 +255,8 @@ ruff_macros::define_rule_mapping!(
     UP032 => violations::FString,
     UP033 => violations::FunctoolsCache,
     UP034 => violations::ExtraneousParentheses,
+    UP035 => violations::ImportReplacements,
+    UP036 => violations::ImportReplacementsSix,
     // pydocstyle
     D100 => violations::PublicModule,
     D101 => violations::PublicClass,

src/rules/pyupgrade/helpers.rs

@@ -1,5 +1,10 @@
 use once_cell::sync::Lazy;
 use regex::{Captures, Regex};
+use rustpython_ast::{AliasData, Located};
+
+use crate::ast::types::Range;
+use crate::ast::whitespace::indentation;
+use crate::source_code::Locator;
 
 static CURLY_ESCAPE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(\\N\{[^}]+})|([{}])").unwrap());
 
@@ -19,3 +24,80 @@ pub fn curly_escape(text: &str) -> String {
         })
         .to_string()
 }
+
+/// Converts a list of names and a module into a from import string. I do not
+/// love this and would MUCH rather use libcst, so if you have any better ideas
+/// feel free to let me know.
+pub fn get_fromimport_str(
+    names: &[AliasData],
+    module: &str,
+    multi_line: bool,
+    indent: &str,
+    short_indent: &str,
+) -> String {
+    if names.is_empty() {
+        return String::new();
+    }
+    let after_comma = if multi_line { '\n' } else { ' ' };
+    let start_imps = if multi_line { "(\n" } else { "" };
+    let after_imps = if multi_line {
+        format!("\n{short_indent})")
+    } else {
+        String::new()
+    };
+    let mut full_names: Vec<String> = vec![];
+    for name in names {
+        let asname_str = match &name.asname {
+            Some(item) => format!(" as {item}"),
+            None => String::new(),
+        };
+        let final_string = format!("{indent}{}{asname_str}", name.name);
+        full_names.push(final_string);
+    }
+    format!(
+        "from {} import {}{}{}",
+        module,
+        start_imps,
+        full_names.join(format!(",{after_comma}").as_str()),
+        after_imps
+    )
+}
+
+pub fn clean_indent<T>(locator: &Locator, located: &Located<T>) -> String {
+    match indentation(locator, located) {
+        // This is an opninionated way of formatting import statements
+        None => "    ".to_string(),
+        Some(item) => item.to_string(),
+    }
+}
+
+pub struct ImportFormatting {
+    pub multi_line: bool,
+    pub indent: String,
+    pub short_indent: String,
+    pub start_indent: String,
+}
+
+impl ImportFormatting {
+    pub fn new<T>(locator: &Locator, stmt: &Located<T>, args: &[Located<AliasData>]) -> Self {
+        let module_text = locator.slice_source_code_range(&Range::from_located(stmt));
+        let multi_line = module_text.contains('\n');
+        let start_indent = clean_indent(locator, stmt);
+        let indent = match args.get(0) {
+            None => panic!("No args for import"),
+            Some(item) if multi_line => clean_indent(locator, item),
+            Some(_) => String::new(),
+        };
+        let short_indent = if indent.len() > 3 {
+            indent[3..].to_string()
+        } else {
+            indent.to_string()
+        };
+        Self {
+            multi_line,
+            indent,
+            short_indent,
+            start_indent,
+        }
+    }
+}

src/rules/pyupgrade/mod.rs

@@ -58,6 +58,8 @@ 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")]
+    #[test_case(Rule::ImportReplacementsSix, Path::new("UP036.py"); "UP036")]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.code(), path.to_string_lossy());
         let diagnostics = test_path(