Remove lex_starts_at usage from UP035 fix logic

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -762,7 +762,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 pyupgrade::rules::deprecated_c_element_tree(checker, stmt);
             }
             if checker.enabled(Rule::DeprecatedImport) {
-                pyupgrade::rules::deprecated_import(checker, stmt, names, module, level);
+                pyupgrade::rules::deprecated_import(checker, import_from);
             }
             if checker.enabled(Rule::UnnecessaryBuiltinImport) {
                 if let Some(module) = module {

crates/ruff_linter/src/rules/pyupgrade/fixes.rs

@@ -1,52 +1,49 @@
-use ruff_python_parser::{lexer, Mode, Tok};
+use ruff_python_ast::{Alias, StmtImportFrom};
+use ruff_python_parser::{lexer, Mode, Tok, TokenKind, Tokens};
 use ruff_source_file::Locator;
-use ruff_text_size::{TextRange, TextSize};
+use ruff_text_size::{Ranged, TextRange, TextSize};
 
 /// Remove any imports matching `members` from an import-from statement.
-pub(crate) fn remove_import_members(contents: &str, members: &[&str]) -> String {
-    let mut names: Vec<TextRange> = vec![];
-    let mut commas: Vec<TextRange> = 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 (tok, range) in lexer::lex(contents, Mode::Module)
-        .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.
-                let last_range = names.last_mut().unwrap();
-                *last_range = TextRange::new(last_range.start(), range.end());
+pub(crate) fn remove_import_members(
+    locator: &Locator<'_>,
+    import_from_stmt: &StmtImportFrom,
+    tokens: &Tokens,
+    members_to_remove: &[&str],
+) -> String {
+    let commas: Vec<TextRange> = tokens
+        .tokens_in_range(import_from_stmt.range())
+        .iter()
+        .skip_while(|token| token.kind() != TokenKind::Import)
+        .filter_map(|token| {
+            if token.kind() == TokenKind::Comma {
+                Some(token.range())
             } else {
-                if members.contains(&&**name) {
-                    removal_indices.push(names.len());
-                }
-                names.push(range);
+                None
             }
-        } else if matches!(tok, Tok::Comma) {
-            commas.push(range);
-        }
-        prev_tok = Some(tok);
-    }
+        })
+        .collect();
 
     // 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 output = String::with_capacity(import_from_stmt.range().len().to_usize());
     let mut last_pos = TextSize::default();
     let mut is_first = true;
-    for index in 0..names.len() {
-        if !removal_indices.contains(&index) {
+
+    for (index, member) in import_from_stmt.names.iter().enumerate() {
+        if !members_to_remove.contains(&member.name.as_str()) {
             is_first = false;
             continue;
         }
 
         let range = if is_first {
-            TextRange::new(names[index].start(), names[index + 1].start())
+            TextRange::new(
+                import_from_stmt.names[index].start(),
+                import_from_stmt.names[index + 1].start(),
+            )
         } else {
-            TextRange::new(commas[index - 1].start(), names[index].end())
+            TextRange::new(
+                commas[index - 1].start(),
+                import_from_stmt.names[index].end(),
+            )
         };
 
         // Add all contents from `last_pos` to `fix.location`.
@@ -68,61 +65,80 @@ pub(crate) fn remove_import_members(contents: &str, members: &[&str]) -> String
 
 #[cfg(test)]
 mod tests {
-    use crate::rules::pyupgrade::fixes::remove_import_members;
+    use ruff_python_parser::parse_module;
+    use ruff_source_file::Locator;
+
+    use super::remove_import_members;
+
+    fn test_helper(source: &str, members_to_remove: &[&str]) -> String {
+        let program = parse_module(source).unwrap();
+        let import_from_stmt = program
+            .suite()
+            .first()
+            .expect("source should have one statement")
+            .as_import_from_stmt()
+            .expect("first statement should be an import from statement");
+        remove_import_members(
+            &Locator::new(source),
+            import_from_stmt,
+            program.tokens(),
+            members_to_remove,
+        )
+    }
 
     #[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"]);
+        let actual = test_helper(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"]);
+        let actual = test_helper(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"]);
+        let actual = test_helper(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"]);
+        let actual = test_helper(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"]);
+        let actual = test_helper(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"]);
+        let actual = test_helper(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"]);
+        let actual = test_helper(source, &["bar", "baz"]);
         assert_eq!(expected, actual);
     }
 
@@ -138,7 +154,7 @@ mod tests {
     bop,
     qux as q
 )";
-        let actual = remove_import_members(source, &["bar", "baz"]);
+        let actual = test_helper(source, &["bar", "baz"]);
         assert_eq!(expected, actual);
     }
 
@@ -155,7 +171,7 @@ mod tests {
     baz,
     qux as q,
 )";
-        let actual = remove_import_members(source, &["bop"]);
+        let actual = test_helper(source, &["bop"]);
         assert_eq!(expected, actual);
     }
 
@@ -171,7 +187,7 @@ mod tests {
     bar,
     qux as q,
 )";
-        let actual = remove_import_members(source, &["baz", "bop"]);
+        let actual = test_helper(source, &["baz", "bop"]);
         assert_eq!(expected, actual);
     }
 
@@ -191,7 +207,7 @@ mod tests {
     # This comment should be retained.
     qux as q,
 )";
-        let actual = remove_import_members(source, &["bop"]);
+        let actual = test_helper(source, &["bop"]);
         assert_eq!(expected, actual);
     }
 
@@ -211,7 +227,7 @@ mod tests {
     bop,
     qux as q,
 )";
-        let actual = remove_import_members(source, &["bar"]);
+        let actual = test_helper(source, &["bar"]);
         assert_eq!(expected, actual);
     }
 }

crates/ruff_linter/src/rules/pyupgrade/rules/deprecated_import.rs

@@ -1,10 +1,11 @@
 use itertools::Itertools;
-use ruff_python_ast::{Alias, Stmt};
+use ruff_python_ast::{Alias, Stmt, StmtImportFrom};
 
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::whitespace::indentation;
 use ruff_python_codegen::Stylist;
+use ruff_python_parser::Tokens;
 use ruff_source_file::Locator;
 use ruff_text_size::Ranged;
 
@@ -362,29 +363,29 @@ const TYPING_EXTENSIONS_TO_COLLECTIONS_ABC_312: &[&str] = &["Buffer"];
 const TYPING_EXTENSIONS_TO_TYPES_312: &[&str] = &["get_original_bases"];
 
 struct ImportReplacer<'a> {
-    stmt: &'a Stmt,
+    import_from_stmt: &'a StmtImportFrom,
     module: &'a str,
-    members: &'a [Alias],
     locator: &'a Locator<'a>,
     stylist: &'a Stylist<'a>,
+    tokens: &'a Tokens,
     version: PythonVersion,
 }
 
 impl<'a> ImportReplacer<'a> {
     const fn new(
-        stmt: &'a Stmt,
+        import_from_stmt: &'a StmtImportFrom,
         module: &'a str,
-        members: &'a [Alias],
         locator: &'a Locator<'a>,
         stylist: &'a Stylist<'a>,
+        tokens: &'a Tokens,
         version: PythonVersion,
     ) -> Self {
         Self {
-            stmt,
+            import_from_stmt,
             module,
-            members,
             locator,
             stylist,
+            tokens,
             version,
         }
     }
@@ -394,7 +395,7 @@ impl<'a> ImportReplacer<'a> {
         let mut operations = vec![];
         if self.module == "typing" {
             if self.version >= PythonVersion::Py39 {
-                for member in self.members {
+                for member in &self.import_from_stmt.names {
                     if let Some(target) = TYPING_TO_RENAME_PY39.iter().find_map(|(name, target)| {
                         if &member.name == *name {
                             Some(*target)
@@ -563,7 +564,7 @@ impl<'a> ImportReplacer<'a> {
             let fix = Some(matched);
             Some((operation, fix))
         } else {
-            let indentation = indentation(self.locator, self.stmt);
+            let indentation = indentation(self.locator, self.import_from_stmt);
 
             // If we have matched _and_ unmatched names, but the import is not on its own
             // line, we can't add a statement after it. For example, if we have
@@ -583,7 +584,9 @@ impl<'a> ImportReplacer<'a> {
 
             let matched = ImportReplacer::format_import_from(&matched_names, target);
             let unmatched = fixes::remove_import_members(
-                self.locator.slice(self.stmt.range()),
+                self.locator,
+                self.import_from_stmt,
+                self.tokens,
                 &matched_names
                     .iter()
                     .map(|name| name.name.as_str())
@@ -611,7 +614,7 @@ impl<'a> ImportReplacer<'a> {
     fn partition_imports(&self, candidates: &[&str]) -> (Vec<&Alias>, Vec<&Alias>) {
         let mut matched_names = vec![];
         let mut unmatched_names = vec![];
-        for name in self.members {
+        for name in &self.import_from_stmt.names {
             if candidates.contains(&name.name.as_str()) {
                 matched_names.push(name);
             } else {
@@ -638,35 +641,32 @@ impl<'a> ImportReplacer<'a> {
 }
 
 /// UP035
-pub(crate) fn deprecated_import(
-    checker: &mut Checker,
-    stmt: &Stmt,
-    names: &[Alias],
-    module: Option<&str>,
-    level: u32,
-) {
+pub(crate) fn deprecated_import(checker: &mut Checker, import_from_stmt: &StmtImportFrom) {
     // Avoid relative and star imports.
-    if level > 0 {
+    if import_from_stmt.level > 0 {
         return;
     }
-    if names.first().is_some_and(|name| &name.name == "*") {
+    if import_from_stmt
+        .names
+        .first()
+        .is_some_and(|name| &name.name == "*")
+    {
         return;
     }
-    let Some(module) = module else {
+    let Some(module) = import_from_stmt.module.as_deref() else {
         return;
     };
 
     if !is_relevant_module(module) {
         return;
     }
 
-    let members: Vec<Alias> = names.iter().map(Clone::clone).collect();
     let fixer = ImportReplacer::new(
-        stmt,
+        import_from_stmt,
         module,
-        &members,
         checker.locator(),
         checker.stylist(),
+        checker.program().tokens(),
         checker.settings.target_version,
     );
 
@@ -675,12 +675,12 @@ pub(crate) fn deprecated_import(
             DeprecatedImport {
                 deprecation: Deprecation::WithoutRename(operation),
             },
-            stmt.range(),
+            import_from_stmt.range(),
         );
         if let Some(content) = fix {
             diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
                 content,
-                stmt.range(),
+                import_from_stmt.range(),
             )));
         }
         checker.diagnostics.push(diagnostic);
@@ -691,7 +691,7 @@ pub(crate) fn deprecated_import(
             DeprecatedImport {
                 deprecation: Deprecation::WithRename(operation),
             },
-            stmt.range(),
+            import_from_stmt.range(),
         );
         checker.diagnostics.push(diagnostic);
     }