STARTED TRANSITION FROM LIBCST TO CUSTOM

src/checkers/ast.rs

@@ -1042,7 +1042,7 @@ where
                     pyupgrade::rules::import_replacements(self, stmt, names, module);
                 }
                 if self.settings.rules.enabled(&Rule::ImportReplacementsSix) {
-                    pyupgrade::rules::import_replacements_six(self, stmt, module);
+                    pyupgrade::rules::import_replacements_six(self, stmt, module, names);
                 }
                 if self.settings.rules.enabled(&Rule::UnnecessaryBuiltinImport) {
                     if let Some(module) = module.as_deref() {

src/rules/pyupgrade/helpers.rs

@@ -1,5 +1,8 @@
+use rustpython_ast::{AliasData, Located};
 use once_cell::sync::Lazy;
 use regex::{Captures, Regex};
+use crate::source_code::Locator;
+use crate::ast::whitespace::indentation;
 
 static CURLY_ESCAPE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(\\N\{[^}]+})|([{}])").unwrap());
 
@@ -19,3 +22,48 @@ 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, name.name, asname_str);
+        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(),
+    }
+}

src/rules/pyupgrade/rules/import_replacements.rs

@@ -1,12 +1,11 @@
 use rustpython_ast::{AliasData, Located, Stmt};
 
 use crate::ast::types::Range;
-use crate::ast::whitespace::indentation;
 use crate::checkers::ast::Checker;
 use crate::fix::Fix;
 use crate::registry::{Diagnostic, Rule};
 use crate::settings::types::PythonVersion;
-use crate::source_code::Locator;
+use crate::rules::pyupgrade::helpers::{clean_indent, get_fromimport_str};
 use crate::violations;
 
 const BAD_MODULES: &[&str] = &[
@@ -317,8 +316,20 @@ 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 = self.get_str(&unmatching_names, self.module);
-        let matching = self.get_str(&matching_names, replace);
+        let unmatching = get_fromimport_str(
+            &unmatching_names,
+            self.module,
+            self.multi_line,
+            self.indent,
+            self.short_indent,
+        );
+        let matching = get_fromimport_str(
+            &matching_names,
+            replace,
+            self.multi_line,
+            self.indent,
+            self.short_indent,
+        );
         // We don't replace if there is just an unmatching, because then we don't need
         // to refactor
         if !unmatching.is_empty() && !matching.is_empty() {
@@ -345,43 +356,6 @@ impl<'a> FixImports<'a> {
         }
         (matching_names, unmatching_names)
     }
-
-    fn get_str(&self, names: &[AliasData], module: &str) -> String {
-        if names.is_empty() {
-            return String::new();
-        }
-        let after_comma = if self.multi_line { '\n' } else { ' ' };
-        let start_imps = if self.multi_line { "(\n" } else { "" };
-        let after_imps = if self.multi_line {
-            format!("\n{})", self.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!("{}{}{}", self.indent, name.name, asname_str);
-            full_names.push(final_string);
-        }
-        format!(
-            "from {} import {}{}{}",
-            module,
-            start_imps,
-            full_names.join(format!(",{}", after_comma).as_str()),
-            after_imps
-        )
-    }
-}
-
-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(),
-    }
 }
 
 /// UP035

src/rules/pyupgrade/rules/import_replacements_six.rs

@@ -1,17 +1,12 @@
-use std::collections::HashMap;
-
-use libcst_native::{
-    AsName, AssignTargetExpression, Codegen, CodegenState, ImportAlias, ImportNames,
-    NameOrAttribute,
-};
 use once_cell::sync::Lazy;
-use rustpython_ast::Stmt;
+use rustpython_ast::{AliasData, Located, Stmt};
+use std::collections::HashMap;
 
 use crate::ast::types::Range;
 use crate::checkers::ast::Checker;
-use crate::cst::matchers::{match_import_from, match_module};
 use crate::fix::Fix;
 use crate::registry::{Diagnostic, Rule};
+use crate::rules::pyupgrade::helpers::{clean_indent, get_fromimport_str};
 use crate::source_code::Locator;
 use crate::violations;
 
@@ -75,51 +70,72 @@ static REPLACE_MODS_URLLIB: Lazy<HashMap<&str, &str>> = Lazy::new(|| {
     m
 });
 
-fn get_asname(asname: &AsName) -> Option<String> {
-    if let AssignTargetExpression::Name(item) = &asname.name {
-        return Some(item.value.to_string());
+struct Formatting {
+    multi_line: bool,
+    indent: String,
+    short_indent: String,
+    start_indent: String,
+}
+
+impl Formatting {
+    fn new<T>(locator: &Locator, stmt: &Stmt, arg1: &Located<T>) -> 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 = clean_indent(locator, arg1);
+        let short_indent = if indent.len() > 3 {
+            indent[3..].to_string()
+        } else {
+            indent.to_string()
+        };
+        Self {
+            multi_line,
+            indent,
+            short_indent,
+            start_indent,
+        }
     }
-    None
 }
 
 fn refactor_segment(
     locator: &Locator,
     stmt: &Stmt,
     replace: &Lazy<HashMap<&str, &str>>,
+    names: &[Located<AliasData>],
+    module: &str,
 ) -> Option<String> {
-    let module_text = locator.slice_source_code_range(&Range::from_located(stmt));
-    let mut tree = match_module(&module_text).unwrap();
-    let mut import = match_import_from(&mut tree).unwrap();
     let mut new_entries = String::new();
-    let mut keep_names: Vec<ImportAlias<'_>> = vec![];
-    if let ImportNames::Aliases(item_names) = &import.names {
-        for name in item_names {
-            if let NameOrAttribute::N(the_name) = &name.name {
-                match replace.get(the_name.value) {
-                    Some(raw_name) => {
-                        new_entries.push_str(&format!("import {}", raw_name));
-                        if let Some(asname) = &name.asname {
-                            if let Some(final_name) = get_asname(asname) {
-                                new_entries.push_str(&format!(" as {}", final_name));
-                            }
-                        }
-                        new_entries.push('\n');
-                    }
-                    None => keep_names.push(name.clone()),
+    let mut keep_names: Vec<AliasData> = vec![];
+    let mut clean_names: Vec<AliasData> = vec![];
+    for name in names {
+        clean_names.push(name.node.clone());
+    }
+
+    let formatting = Formatting::new(locator, stmt, &names.get(0).unwrap());
+    for name in names {
+        match replace.get(name.node.name.as_str()) {
+            None => keep_names.push(name.node.clone()),
+            Some(item) => {
+                // MAKE SURE TO ADD IF STATEMENTS HERE
+                new_entries.push_str(&format!("{}import {item}", formatting.start_indent));
+                if let Some(final_name) = &name.node.asname {
+                    new_entries.push_str(&format!(" as {}", final_name));
                 }
-            } else {
-                keep_names.push(name.clone());
+                new_entries.push('\n');
             }
         }
     }
     // If nothing was different, there is no need to change
     if new_entries.is_empty() {
         return None;
     }
-    import.names = ImportNames::Aliases(keep_names);
-    let mut state = CodegenState::default();
-    import.codegen(&mut state);
-    let mut final_str = state.to_string();
+    let mut final_str = get_fromimport_str(
+        &keep_names,
+        module,
+        formatting.multi_line,
+        &formatting.indent,
+        &formatting.short_indent,
+    );
     final_str.push_str(&format!("\n{}", new_entries));
     if final_str.ends_with('\n') {
         final_str.pop();
@@ -128,17 +144,28 @@ fn refactor_segment(
 }
 
 /// UP036
-pub fn import_replacements_six(checker: &mut Checker, stmt: &Stmt, module: &Option<String>) {
+pub fn import_replacements_six(
+    checker: &mut Checker,
+    stmt: &Stmt,
+    module: &Option<String>,
+    names: &Vec<Located<AliasData>>,
+) {
     // Pyupgrade only works with import_from statements, so this linter does that as
     // well
 
-    // This only applies to six.moves libraries
     let final_string: Option<String>;
     if let Some(module_text) = module {
         if module_text == "six.moves" {
-            final_string = refactor_segment(checker.locator, stmt, &REPLACE_MODS);
+            final_string =
+                refactor_segment(checker.locator, stmt, &REPLACE_MODS, &names, module_text);
         } else if module_text == "six.moves.urllib" {
-            final_string = refactor_segment(checker.locator, stmt, &REPLACE_MODS_URLLIB);
+            final_string = refactor_segment(
+                checker.locator,
+                stmt,
+                &REPLACE_MODS_URLLIB,
+                &names,
+                module_text,
+            );
         } else {
             return;
         }