Small refactors; add one failing test

README.md

@@ -813,7 +813,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 | Replace old formatting imports with their new versions | 🛠 |
+| UP035 | import-replacements | Import `{existing}` from `{replacement}` | 🛠 |
 
 ### flake8-2020 (YTT)
 

resources/test/fixtures/pyupgrade/UP035.py

@@ -29,6 +29,9 @@
 
 if True: from collections import Mapping
 
+if True: from collections import (
+    Mapping, Counter)
+
 import os
 from collections import Counter, Mapping
 import sys

src/checkers/ast.rs

@@ -1074,6 +1074,7 @@ where
                         stmt,
                         names,
                         module.as_ref().map(String::as_str),
+                        level.as_ref(),
                     );
                 }
                 if self.settings.rules.enabled(&Rule::UnnecessaryBuiltinImport) {

src/rules/pyupgrade/helpers.rs

@@ -1,10 +1,10 @@
 use once_cell::sync::Lazy;
 use regex::{Captures, Regex};
-use rustpython_ast::{AliasData, Located};
+use rustpython_ast::{Alias, AliasData, Located};
 
 use crate::ast::types::Range;
 use crate::ast::whitespace::indentation;
-use crate::source_code::Locator;
+use crate::source_code::{Locator, Stylist};
 
 static CURLY_BRACES: Lazy<Regex> = Lazy::new(|| Regex::new(r"(\\N\{[^}]+})|([{}])").unwrap());
 
@@ -27,76 +27,75 @@ pub fn curly_escape(text: &str) -> String {
 }
 
 /// Converts a list of names and a module into an `import from`-style import.
-pub fn get_fromimport_str(
+pub fn format_import_from(
     names: &[AliasData],
     module: &str,
     multi_line: bool,
     indent: &str,
-    short_indent: &str,
+    start_indent: &str,
+    stylist: &Stylist,
 ) -> 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})")
+    // Construct the whitespace strings.
+    let line_ending = stylist.line_ending().as_str();
+    let after_comma = if multi_line { line_ending } else { " " };
+    let before_imports = if multi_line {
+        format!("({line_ending}")
     } else {
-        String::new()
+        "".to_string()
     };
-    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);
-    }
+    let after_imports = if multi_line {
+        format!("{line_ending}{start_indent})")
+    } else {
+        "".to_string()
+    };
+
+    // Generate the formatted names.
+    let full_names = {
+        let full_names: Vec<String> = names
+            .iter()
+            .map(|name| match &name.asname {
+                Some(asname) => format!("{indent}{} as {asname}", name.name),
+                None => format!("{indent}{}", name.name),
+            })
+            .collect();
+        let mut full_names = full_names.join(&format!(",{}", after_comma));
+        if multi_line {
+            full_names.push(',');
+        }
+        full_names
+    };
+
     format!(
         "from {} import {}{}{}",
-        module,
-        start_imps,
-        full_names.join(format!(",{after_comma}").as_str()),
-        after_imps
+        module, before_imports, full_names, after_imports
     )
 }
 
-pub fn clean_indent<T>(locator: &Locator, located: &Located<T>) -> String {
-    match indentation(locator, located) {
-        // This is an opinionated way of formatting import statements.
-        None => "    ".to_string(),
-        Some(item) => item.to_string(),
-    }
-}
-
-pub struct ImportFormatting {
+pub struct ImportFormatting<'a> {
     pub multi_line: bool,
-    pub indent: String,
-    pub short_indent: String,
-    pub start_indent: String,
+    pub member_indent: &'a str,
+    pub stmt_indent: &'a str,
 }
 
-impl ImportFormatting {
-    pub fn new<T>(locator: &Locator, stmt: &Located<T>, args: &[Located<AliasData>]) -> Self {
+impl<'a> ImportFormatting<'a> {
+    pub fn new<T>(
+        locator: &'a Locator<'a>,
+        stylist: &'a Stylist<'a>,
+        stmt: &'a Located<T>,
+        args: &'a [Alias],
+    ) -> 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()
+        let multi_line = module_text.contains(stylist.line_ending().as_str());
+        let stmt_indent = indentation(locator, stmt).unwrap_or(stylist.indentation().as_str());
+        let member_indent = if multi_line {
+            indentation(locator, &args[0]).unwrap_or(stylist.indentation().as_str())
         } else {
-            indent.to_string()
+            ""
         };
         Self {
             multi_line,
-            indent,
-            short_indent,
-            start_indent,
+            member_indent,
+            stmt_indent,
         }
     }
 }

src/rules/pyupgrade/rules/import_replacements.rs

@@ -7,8 +7,9 @@ use crate::checkers::ast::Checker;
 use crate::define_violation;
 use crate::fix::Fix;
 use crate::registry::{Diagnostic, Rule};
-use crate::rules::pyupgrade::helpers::{get_fromimport_str, ImportFormatting};
+use crate::rules::pyupgrade::helpers::{format_import_from, ImportFormatting};
 use crate::settings::types::PythonVersion;
+use crate::source_code::Stylist;
 use crate::violation::AlwaysAutofixableViolation;
 
 define_violation!(
@@ -29,7 +30,7 @@ impl AlwaysAutofixableViolation for ImportReplacements {
 
     fn autofix_title(&self) -> String {
         let ImportReplacements {
-            existing,
+            existing: _,
             replacement,
         } = self;
         format!("Replace with `{replacement}`")
@@ -211,37 +212,32 @@ struct FixImports<'a> {
     module: &'a str,
     multi_line: bool,
     names: &'a [AliasData],
-    // This is the indent level of the first named import
-    indent: &'a str,
-    // This is the indent for the parentheses at the end of a multi-line statement
-    short_indent: &'a str,
-    // This is the indent of the actual import statement
-    starting_indent: &'a str,
+    // The indent level of the first named import.
+    member_indent: &'a str,
+    // The indent of the import statement.
+    stmt_indent: &'a str,
     version: PythonVersion,
+    stylist: &'a Stylist<'a>,
 }
 
 impl<'a> FixImports<'a> {
     fn new(
         module: &'a str,
-        multi_line: bool,
         names: &'a [AliasData],
-        indent: &'a str,
+        multi_line: bool,
+        member_indent: &'a str,
+        stmt_indent: &'a str,
         version: PythonVersion,
-        starting_indent: &'a str,
+        stylist: &'a Stylist,
     ) -> Self {
-        let short_indent = if indent.len() > 3 {
-            &indent[3..]
-        } else {
-            indent
-        };
         Self {
             module,
             multi_line,
             names,
-            indent,
-            short_indent,
-            starting_indent,
+            member_indent,
+            stmt_indent,
             version,
+            stylist,
         }
     }
 
@@ -315,35 +311,38 @@ impl<'a> FixImports<'a> {
 
     /// Converts the string of imports into new one
     fn create_new_str(&self, matches: &[&str], replace: &str) -> Option<String> {
-        let (matching_names, unmatching_names) = self.get_import_lists(matches);
-        let unmatching = get_fromimport_str(
-            &unmatching_names,
-            self.module,
-            self.multi_line,
-            self.indent,
-            self.short_indent,
-        );
-        let matching = get_fromimport_str(
+        let (matching_names, unmatching_names) = self.split_imports(matches);
+        if matching_names.is_empty() {
+            return None;
+        }
+
+        let matching = format_import_from(
             &matching_names,
             replace,
             self.multi_line,
-            self.indent,
-            self.short_indent,
+            self.member_indent,
+            self.stmt_indent,
+            self.stylist,
         );
-        // We don't replace if there is unmatching, because then we don't need
-        // to refactor
-        if !unmatching.is_empty() && !matching.is_empty() {
-            Some(format!("{unmatching}\n{}{matching}", self.starting_indent))
-        } else if !matching.is_empty() {
-            Some(matching)
-        } else {
-            None
+
+        if unmatching_names.is_empty() {
+            return Some(matching);
         }
+
+        let unmatching = format_import_from(
+            &unmatching_names,
+            self.module,
+            self.multi_line,
+            self.member_indent,
+            self.stmt_indent,
+            self.stylist,
+        );
+        Some(format!("{unmatching}\n{}{matching}", self.stmt_indent))
     }
 
     /// Returns a list of imports that does and does not have a match in the
     /// given list of matches
-    fn get_import_lists(&self, matches: &[&str]) -> (Vec<AliasData>, Vec<AliasData>) {
+    fn split_imports(&self, matches: &[&str]) -> (Vec<AliasData>, Vec<AliasData>) {
         let mut unmatching_names: Vec<AliasData> = vec![];
         let mut matching_names: Vec<AliasData> = vec![];
 
@@ -364,49 +363,45 @@ pub fn import_replacements(
     stmt: &Stmt,
     names: &[Alias],
     module: Option<&str>,
+    level: Option<&usize>,
 ) {
-    // Pyupgrade only works with import_from statements, so this linter does that as
-    // well
+    // Avoid relative imports.
+    if level.map_or(false, |level| *level > 0) {
+        return;
+    }
     let Some(module) = module else {
         return;
     };
+
     if !RELEVANT_MODULES.contains(&module) {
         return;
     }
 
-    let mut clean_names: Vec<AliasData> = vec![];
-    for name in names {
-        clean_names.push(name.node.clone());
-    }
-    let module_text = checker
-        .locator
-        .slice_source_code_range(&Range::from_located(stmt));
-    let formatting = ImportFormatting::new(checker.locator, stmt, names);
+    let formatting = ImportFormatting::new(checker.locator, checker.stylist, stmt, names);
+    let names: Vec<AliasData> = names.iter().map(|alias| alias.node.clone()).collect();
     let fixer = FixImports::new(
         module,
+        &names,
         formatting.multi_line,
-        &clean_names,
-        &formatting.indent,
+        &formatting.member_indent,
+        &formatting.stmt_indent,
         checker.settings.target_version,
-        &formatting.start_indent,
+        checker.stylist,
     );
-    let Some(clean_result) = fixer.check_replacement() else {
+    let Some(content) = fixer.check_replacement() else {
         return;
     };
-    if clean_result == module_text {
-        return;
-    }
-    let range = Range::from_located(stmt);
+
     let mut diagnostic = Diagnostic::new(
         ImportReplacements {
             existing: module.to_string(),
-            replacement: clean_result.to_string(),
+            replacement: content.to_string(),
         },
-        range,
+        Range::from_located(stmt),
     );
     if checker.patch(&Rule::ImportReplacements) {
         diagnostic.amend(Fix::replacement(
-            clean_result,
+            content,
             stmt.location,
             stmt.end_location.unwrap(),
         ));