Add import insertion support to autofix capabilities (#3787)

crates/ruff/resources/test/fixtures/pylint/sys_exit_alias_1.py

@@ -9,4 +9,8 @@ def main():
     quit(1)
 
 
-sys.exit(2)
+def main():
+    sys = 1
+
+    exit(1)
+    quit(1)

crates/ruff/src/autofix/helpers.rs

@@ -7,13 +7,16 @@ use rustpython_parser::ast::{ExcepthandlerKind, Expr, Keyword, Location, Stmt, S
 use rustpython_parser::{lexer, Mode, Tok};
 
 use ruff_diagnostics::Edit;
+use ruff_python_ast::context::Context;
 use ruff_python_ast::helpers;
 use ruff_python_ast::helpers::to_absolute;
+use ruff_python_ast::imports::{AnyImport, Import};
 use ruff_python_ast::newlines::NewlineWithTrailingNewline;
 use ruff_python_ast::source_code::{Indexer, Locator, Stylist};
 
 use crate::cst::helpers::compose_module_path;
 use crate::cst::matchers::match_module;
+use crate::importer::Importer;
 
 /// Determine if a body contains only a single statement, taking into account
 /// deleted.
@@ -444,6 +447,82 @@ pub fn remove_argument(
     }
 }
 
+/// Generate an [`Edit`] to reference the given symbol. Returns the [`Edit`] necessary to make the
+/// symbol available in the current scope along with the bound name of the symbol.
+///
+/// For example, assuming `module` is `"functools"` and `member` is `"lru_cache"`, this function
+/// could return an [`Edit`] to add `import functools` to the top of the file, alongside with the
+/// name on which the `lru_cache` symbol would be made available (`"functools.lru_cache"`).
+///
+/// Attempts to reuse existing imports when possible.
+pub fn get_or_import_symbol(
+    module: &str,
+    member: &str,
+    context: &Context,
+    importer: &Importer,
+    locator: &Locator,
+) -> Result<(Edit, String)> {
+    if let Some((source, binding)) = context.resolve_qualified_import_name(module, member) {
+        // If the symbol is already available in the current scope, use it.
+        //
+        // We also add a no-nop edit to force conflicts with any other fixes that might try to
+        // remove the import. Consider:
+        //
+        // ```py
+        // import sys
+        //
+        // quit()
+        // ```
+        //
+        // Assume you omit this no-op edit. If you run Ruff with `unused-imports` and
+        // `sys-exit-alias` over this snippet, it will generate two fixes: (1) remove the unused
+        // `sys` import; and (2) replace `quit()` with `sys.exit()`, under the assumption that `sys`
+        // is already imported and available.
+        //
+        // By adding this no-op edit, we force the `unused-imports` fix to conflict with the
+        // `sys-exit-alias` fix, and thus will avoid applying both fixes in the same pass.
+        let import_edit = Edit::replacement(
+            locator.slice(source).to_string(),
+            source.location,
+            source.end_location.unwrap(),
+        );
+        Ok((import_edit, binding))
+    } else {
+        if let Some(stmt) = importer.get_import_from(module) {
+            // Case 1: `from functools import lru_cache` is in scope, and we're trying to reference
+            // `functools.cache`; thus, we add `cache` to the import, and return `"cache"` as the
+            // bound name.
+            if context
+                .find_binding(member)
+                .map_or(true, |binding| binding.kind.is_builtin())
+            {
+                let import_edit = importer.add_member(stmt, member)?;
+                Ok((import_edit, member.to_string()))
+            } else {
+                bail!(
+                    "Unable to insert `{}` into scope due to name conflict",
+                    member
+                )
+            }
+        } else {
+            // Case 2: No `functools` import is in scope; thus, we add `import functools`, and
+            // return `"functools.cache"` as the bound name.
+            if context
+                .find_binding(module)
+                .map_or(true, |binding| binding.kind.is_builtin())
+            {
+                let import_edit = importer.add_import(&AnyImport::Import(Import::module(module)));
+                Ok((import_edit, format!("{module}.{member}")))
+            } else {
+                bail!(
+                    "Unable to insert `{}` into scope due to name conflict",
+                    module
+                )
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use anyhow::Result;

crates/ruff/src/checkers/ast/mod.rs

@@ -38,6 +38,7 @@ use crate::docstrings::definition::{
     transition_scope, Definition, DefinitionKind, Docstring, Documentable,
 };
 use crate::fs::relativize_path;
+use crate::importer::Importer;
 use crate::registry::{AsRule, Rule};
 use crate::rules::{
     flake8_2020, flake8_annotations, flake8_bandit, flake8_blind_except, flake8_boolean_trap,
@@ -70,6 +71,7 @@ pub struct Checker<'a> {
     pub locator: &'a Locator<'a>,
     pub stylist: &'a Stylist<'a>,
     pub indexer: &'a Indexer,
+    pub importer: Importer<'a>,
     // Stateful fields.
     pub ctx: Context<'a>,
     pub deferred: Deferred<'a>,
@@ -92,6 +94,7 @@ impl<'a> Checker<'a> {
         locator: &'a Locator,
         stylist: &'a Stylist,
         indexer: &'a Indexer,
+        importer: Importer<'a>,
     ) -> Checker<'a> {
         Checker {
             settings,
@@ -105,6 +108,7 @@ impl<'a> Checker<'a> {
             locator,
             stylist,
             indexer,
+            importer,
             ctx: Context::new(&settings.typing_modules, path, module_path),
             deferred: Deferred::default(),
             diagnostics: Vec::default(),
@@ -189,6 +193,18 @@ where
                 }
             }
         }
+
+        // Track each top-level import, to guide import insertions.
+        if matches!(
+            &stmt.node,
+            StmtKind::Import { .. } | StmtKind::ImportFrom { .. }
+        ) {
+            let scope_index = self.ctx.scope_id();
+            if scope_index.is_global() && self.ctx.current_stmt_parent().is_none() {
+                self.importer.visit_import(stmt);
+            }
+        }
+
         // Pre-visit.
         match &stmt.node {
             StmtKind::Global { names } => {
@@ -5438,6 +5454,7 @@ pub fn check_ast(
         locator,
         stylist,
         indexer,
+        Importer::new(python_ast, locator, stylist),
     );
     checker.bind_builtins();
 

crates/ruff/src/cst/matchers.rs

@@ -1,7 +1,7 @@
 use anyhow::{bail, Result};
 use libcst_native::{
-    Attribute, Call, Comparison, Dict, Expr, Expression, Import, ImportFrom, Module, SimpleString,
+    Attribute, Call, Comparison, Dict, Expr, Expression, Import, ImportAlias, ImportFrom,
-    SmallStatement, Statement,
+    ImportNames, Module, SimpleString, SmallStatement, Statement,
 };
 
 pub fn match_module(module_text: &str) -> Result<Module> {
@@ -54,6 +54,16 @@ pub fn match_import_from<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut I
     }
 }
 
+pub fn match_aliases<'a, 'b>(
+    import_from: &'a mut ImportFrom<'b>,
+) -> Result<&'a mut Vec<ImportAlias<'b>>> {
+    if let ImportNames::Aliases(aliases) = &mut import_from.names {
+        Ok(aliases)
+    } else {
+        bail!("Expected ImportNames::Aliases")
+    }
+}
+
 pub fn match_call<'a, 'b>(expression: &'a mut Expression<'b>) -> Result<&'a mut Call<'b>> {
     if let Expression::Call(call) = expression {
         Ok(call)

crates/ruff/src/importer.rs

@@ -0,0 +1,359 @@
+//! Add and modify import statements to make module members available during fix execution.
+
+use anyhow::Result;
+use libcst_native::{Codegen, CodegenState, ImportAlias, Name, NameOrAttribute};
+use rustc_hash::FxHashMap;
+use rustpython_parser::ast::{Location, Stmt, StmtKind, Suite};
+use rustpython_parser::{lexer, Mode, Tok};
+
+use ruff_diagnostics::Edit;
+use ruff_python_ast::helpers::is_docstring_stmt;
+use ruff_python_ast::imports::AnyImport;
+use ruff_python_ast::source_code::{Locator, Stylist};
+
+use crate::cst::matchers::{match_aliases, match_import_from, match_module};
+
+pub struct Importer<'a> {
+    python_ast: &'a Suite,
+    locator: &'a Locator<'a>,
+    stylist: &'a Stylist<'a>,
+    /// A map from module name to top-level `StmtKind::ImportFrom` statements.
+    import_from_map: FxHashMap<&'a str, &'a Stmt>,
+    /// The last top-level import statement.
+    trailing_import: Option<&'a Stmt>,
+}
+
+impl<'a> Importer<'a> {
+    pub fn new(python_ast: &'a Suite, locator: &'a Locator<'a>, stylist: &'a Stylist<'a>) -> Self {
+        Self {
+            python_ast,
+            locator,
+            stylist,
+            import_from_map: FxHashMap::default(),
+            trailing_import: None,
+        }
+    }
+
+    /// Visit a top-level import statement.
+    pub fn visit_import(&mut self, import: &'a Stmt) {
+        // Store a reference to the import statement in the appropriate map.
+        match &import.node {
+            StmtKind::Import { .. } => {
+                // Nothing to do here, we don't extend top-level `import` statements at all, so
+                // no need to track them.
+            }
+            StmtKind::ImportFrom { module, level, .. } => {
+                // Store a reverse-map from module name to `import ... from` statement.
+                if level.map_or(true, |level| level == 0) {
+                    if let Some(module) = module {
+                        self.import_from_map.insert(module.as_str(), import);
+                    }
+                }
+            }
+            _ => {
+                panic!("Expected StmtKind::Import | StmtKind::ImportFrom");
+            }
+        }
+
+        // Store a reference to the last top-level import statement.
+        self.trailing_import = Some(import);
+    }
+
+    /// Add an import statement to import the given module.
+    ///
+    /// If there are no existing imports, the new import will be added at the top
+    /// of the file. Otherwise, it will be added after the most recent top-level
+    /// import statement.
+    pub fn add_import(&self, import: &AnyImport) -> Edit {
+        let required_import = import.to_string();
+        if let Some(stmt) = self.trailing_import {
+            // Insert after the last top-level import.
+            let Insertion {
+                prefix,
+                location,
+                suffix,
+            } = end_of_statement_insertion(stmt, self.locator, self.stylist);
+            let content = format!("{prefix}{required_import}{suffix}");
+            Edit::insertion(content, location)
+        } else {
+            // Insert at the top of the file.
+            let Insertion {
+                prefix,
+                location,
+                suffix,
+            } = top_of_file_insertion(self.python_ast, self.locator, self.stylist);
+            let content = format!("{prefix}{required_import}{suffix}");
+            Edit::insertion(content, location)
+        }
+    }
+
+    /// Return the top-level [`Stmt`] that imports the given module using `StmtKind::ImportFrom`.
+    /// if it exists.
+    pub fn get_import_from(&self, module: &str) -> Option<&Stmt> {
+        self.import_from_map.get(module).copied()
+    }
+
+    /// Add the given member to an existing `StmtKind::ImportFrom` statement.
+    pub fn add_member(&self, stmt: &Stmt, member: &str) -> Result<Edit> {
+        let mut tree = match_module(self.locator.slice(stmt))?;
+        let import_from = match_import_from(&mut tree)?;
+        let aliases = match_aliases(import_from)?;
+        aliases.push(ImportAlias {
+            name: NameOrAttribute::N(Box::new(Name {
+                value: member,
+                lpar: vec![],
+                rpar: vec![],
+            })),
+            asname: None,
+            comma: aliases.last().and_then(|alias| alias.comma.clone()),
+        });
+        let mut state = CodegenState {
+            default_newline: &self.stylist.line_ending(),
+            default_indent: self.stylist.indentation(),
+            ..CodegenState::default()
+        };
+        tree.codegen(&mut state);
+        Ok(Edit::replacement(
+            state.to_string(),
+            stmt.location,
+            stmt.end_location.unwrap(),
+        ))
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct Insertion {
+    /// The content to add before the insertion.
+    prefix: &'static str,
+    /// The location at which to insert.
+    location: Location,
+    /// The content to add after the insertion.
+    suffix: &'static str,
+}
+
+impl Insertion {
+    fn new(prefix: &'static str, location: Location, suffix: &'static str) -> Self {
+        Self {
+            prefix,
+            location,
+            suffix,
+        }
+    }
+}
+
+/// Find the end of the last docstring.
+fn match_docstring_end(body: &[Stmt]) -> Option<Location> {
+    let mut iter = body.iter();
+    let Some(mut stmt) = iter.next() else {
+        return None;
+    };
+    if !is_docstring_stmt(stmt) {
+        return None;
+    }
+    for next in iter {
+        if !is_docstring_stmt(next) {
+            break;
+        }
+        stmt = next;
+    }
+    Some(stmt.end_location.unwrap())
+}
+
+/// Find the location at which a "top-of-file" import should be inserted,
+/// along with a prefix and suffix to use for the insertion.
+///
+/// For example, given the following code:
+///
+/// ```python
+/// """Hello, world!"""
+///
+/// import os
+/// ```
+///
+/// The location returned will be the start of the `import os` statement,
+/// along with a trailing newline suffix.
+fn end_of_statement_insertion(stmt: &Stmt, locator: &Locator, stylist: &Stylist) -> Insertion {
+    let location = stmt.end_location.unwrap();
+    let mut tokens = lexer::lex_located(locator.skip(location), Mode::Module, location).flatten();
+    if let Some((.., Tok::Semi, end)) = tokens.next() {
+        // If the first token after the docstring is a semicolon, insert after the semicolon as an
+        // inline statement;
+        Insertion::new(" ", end, ";")
+    } else {
+        // Otherwise, insert on the next line.
+        Insertion::new(
+            "",
+            Location::new(location.row() + 1, 0),
+            stylist.line_ending().as_str(),
+        )
+    }
+}
+
+/// Find the location at which a "top-of-file" import should be inserted,
+/// along with a prefix and suffix to use for the insertion.
+///
+/// For example, given the following code:
+///
+/// ```python
+/// """Hello, world!"""
+///
+/// import os
+/// ```
+///
+/// The location returned will be the start of the `import os` statement,
+/// along with a trailing newline suffix.
+fn top_of_file_insertion(body: &[Stmt], locator: &Locator, stylist: &Stylist) -> Insertion {
+    // Skip over any docstrings.
+    let mut location = if let Some(location) = match_docstring_end(body) {
+        // If the first token after the docstring is a semicolon, insert after the semicolon as an
+        // inline statement;
+        let first_token = lexer::lex_located(locator.skip(location), Mode::Module, location)
+            .flatten()
+            .next();
+        if let Some((.., Tok::Semi, end)) = first_token {
+            return Insertion::new(" ", end, ";");
+        }
+
+        // Otherwise, advance to the next row.
+        Location::new(location.row() + 1, 0)
+    } else {
+        Location::default()
+    };
+
+    // Skip over any comments and empty lines.
+    for (.., tok, end) in
+        lexer::lex_located(locator.skip(location), Mode::Module, location).flatten()
+    {
+        if matches!(tok, Tok::Comment(..) | Tok::Newline) {
+            location = Location::new(end.row() + 1, 0);
+        } else {
+            break;
+        }
+    }
+
+    return Insertion::new("", location, stylist.line_ending().as_str());
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use rustpython_parser as parser;
+    use rustpython_parser::ast::Location;
+    use rustpython_parser::lexer::LexResult;
+
+    use ruff_python_ast::source_code::{LineEnding, Locator, Stylist};
+
+    use crate::importer::{top_of_file_insertion, Insertion};
+
+    fn insert(contents: &str) -> Result<Insertion> {
+        let program = parser::parse_program(contents, "<filename>")?;
+        let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents);
+        let locator = Locator::new(contents);
+        let stylist = Stylist::from_tokens(&tokens, &locator);
+        Ok(top_of_file_insertion(&program, &locator, &stylist))
+    }
+
+    #[test]
+    fn top_of_file_insertions() -> Result<()> {
+        let contents = "";
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new("", Location::new(1, 0), LineEnding::default().as_str())
+        );
+
+        let contents = r#"
+"""Hello, world!""""#
+            .trim_start();
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new("", Location::new(2, 0), LineEnding::default().as_str())
+        );
+
+        let contents = r#"
+"""Hello, world!"""
+"#
+        .trim_start();
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new("", Location::new(2, 0), "\n")
+        );
+
+        let contents = r#"
+"""Hello, world!"""
+"""Hello, world!"""
+"#
+        .trim_start();
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new("", Location::new(3, 0), "\n")
+        );
+
+        let contents = r#"
+x = 1
+"#
+        .trim_start();
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new("", Location::new(1, 0), "\n")
+        );
+
+        let contents = r#"
+#!/usr/bin/env python3
+"#
+        .trim_start();
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new("", Location::new(2, 0), "\n")
+        );
+
+        let contents = r#"
+#!/usr/bin/env python3
+"""Hello, world!"""
+"#
+        .trim_start();
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new("", Location::new(3, 0), "\n")
+        );
+
+        let contents = r#"
+"""Hello, world!"""
+#!/usr/bin/env python3
+"#
+        .trim_start();
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new("", Location::new(3, 0), "\n")
+        );
+
+        let contents = r#"
+"""%s""" % "Hello, world!"
+"#
+        .trim_start();
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new("", Location::new(1, 0), "\n")
+        );
+
+        let contents = r#"
+"""Hello, world!"""; x = 1
+"#
+        .trim_start();
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new(" ", Location::new(1, 20), ";")
+        );
+
+        let contents = r#"
+"""Hello, world!"""; x = 1; y = \
+    2
+"#
+        .trim_start();
+        assert_eq!(
+            insert(contents)?,
+            Insertion::new(" ", Location::new(1, 20), ";")
+        );
+
+        Ok(())
+    }
+}

crates/ruff/src/lib.rs

@@ -19,6 +19,7 @@ mod doc_lines;
 mod docstrings;
 pub mod flake8_to_ruff;
 pub mod fs;
+mod importer;
 pub mod jupyter;
 mod lex;
 pub mod linter;

crates/ruff/src/rules/isort/helpers.rs

@@ -1,9 +1,8 @@
-use rustpython_parser::ast::{Location, Stmt};
+use rustpython_parser::ast::Stmt;
 use rustpython_parser::{lexer, Mode, Tok};
 
-use ruff_python_ast::helpers::is_docstring_stmt;
 use ruff_python_ast::newlines::StrExt;
-use ruff_python_ast::source_code::{Locator, Stylist};
+use ruff_python_ast::source_code::Locator;
 
 use crate::rules::isort::types::TrailingComma;
 
@@ -83,214 +82,3 @@ pub fn has_comment_break(stmt: &Stmt, locator: &Locator) -> bool {
     }
     false
 }
-
-/// Find the end of the last docstring.
-fn match_docstring_end(body: &[Stmt]) -> Option<Location> {
-    let mut iter = body.iter();
-    let Some(mut stmt) = iter.next() else {
-        return None;
-    };
-    if !is_docstring_stmt(stmt) {
-        return None;
-    }
-    for next in iter {
-        if !is_docstring_stmt(next) {
-            break;
-        }
-        stmt = next;
-    }
-    Some(stmt.end_location.unwrap())
-}
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub(super) struct Insertion {
-    /// The content to add before the insertion.
-    pub prefix: &'static str,
-    /// The location at which to insert.
-    pub location: Location,
-    /// The content to add after the insertion.
-    pub suffix: &'static str,
-}
-
-impl Insertion {
-    pub fn new(prefix: &'static str, location: Location, suffix: &'static str) -> Self {
-        Self {
-            prefix,
-            location,
-            suffix,
-        }
-    }
-}
-
-/// Find the location at which a "top-of-file" import should be inserted,
-/// along with a prefix and suffix to use for the insertion.
-///
-/// For example, given the following code:
-///
-/// ```python
-/// """Hello, world!"""
-///
-/// import os
-/// ```
-///
-/// The location returned will be the start of the `import os` statement,
-/// along with a trailing newline suffix.
-pub(super) fn top_of_file_insertion(
-    body: &[Stmt],
-    locator: &Locator,
-    stylist: &Stylist,
-) -> Insertion {
-    // Skip over any docstrings.
-    let mut location = if let Some(location) = match_docstring_end(body) {
-        // If the first token after the docstring is a semicolon, insert after the semicolon as an
-        // inline statement;
-        let first_token = lexer::lex_located(locator.skip(location), Mode::Module, location)
-            .flatten()
-            .next();
-        if let Some((.., Tok::Semi, end)) = first_token {
-            return Insertion::new(" ", end, ";");
-        }
-
-        // Otherwise, advance to the next row.
-        Location::new(location.row() + 1, 0)
-    } else {
-        Location::default()
-    };
-
-    // Skip over any comments and empty lines.
-    for (.., tok, end) in
-        lexer::lex_located(locator.skip(location), Mode::Module, location).flatten()
-    {
-        if matches!(tok, Tok::Comment(..) | Tok::Newline) {
-            location = Location::new(end.row() + 1, 0);
-        } else {
-            break;
-        }
-    }
-
-    return Insertion::new("", location, stylist.line_ending().as_str());
-}
-
-#[cfg(test)]
-mod tests {
-    use anyhow::Result;
-    use rustpython_parser as parser;
-    use rustpython_parser::ast::Location;
-    use rustpython_parser::lexer::LexResult;
-
-    use ruff_python_ast::source_code::{LineEnding, Locator, Stylist};
-
-    use crate::rules::isort::helpers::{top_of_file_insertion, Insertion};
-
-    fn insert(contents: &str) -> Result<Insertion> {
-        let program = parser::parse_program(contents, "<filename>")?;
-        let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents);
-        let locator = Locator::new(contents);
-        let stylist = Stylist::from_tokens(&tokens, &locator);
-        Ok(top_of_file_insertion(&program, &locator, &stylist))
-    }
-
-    #[test]
-    fn top_of_file_insertions() -> Result<()> {
-        let contents = "";
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new("", Location::new(1, 0), LineEnding::default().as_str())
-        );
-
-        let contents = r#"
-"""Hello, world!""""#
-            .trim_start();
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new("", Location::new(2, 0), LineEnding::default().as_str())
-        );
-
-        let contents = r#"
-"""Hello, world!"""
-"#
-        .trim_start();
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new("", Location::new(2, 0), "\n")
-        );
-
-        let contents = r#"
-"""Hello, world!"""
-"""Hello, world!"""
-"#
-        .trim_start();
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new("", Location::new(3, 0), "\n")
-        );
-
-        let contents = r#"
-x = 1
-"#
-        .trim_start();
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new("", Location::new(1, 0), "\n")
-        );
-
-        let contents = r#"
-#!/usr/bin/env python3
-"#
-        .trim_start();
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new("", Location::new(2, 0), "\n")
-        );
-
-        let contents = r#"
-#!/usr/bin/env python3
-"""Hello, world!"""
-"#
-        .trim_start();
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new("", Location::new(3, 0), "\n")
-        );
-
-        let contents = r#"
-"""Hello, world!"""
-#!/usr/bin/env python3
-"#
-        .trim_start();
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new("", Location::new(3, 0), "\n")
-        );
-
-        let contents = r#"
-"""%s""" % "Hello, world!"
-"#
-        .trim_start();
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new("", Location::new(1, 0), "\n")
-        );
-
-        let contents = r#"
-"""Hello, world!"""; x = 1
-"#
-        .trim_start();
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new(" ", Location::new(1, 20), ";")
-        );
-
-        let contents = r#"
-"""Hello, world!"""; x = 1; y = \
-    2
-"#
-        .trim_start();
-        assert_eq!(
-            insert(contents)?,
-            Insertion::new(" ", Location::new(1, 20), ";")
-        );
-
-        Ok(())
-    }
-}

crates/ruff/src/rules/isort/rules/add_required_imports.rs

@@ -2,19 +2,18 @@ use log::error;
 use rustpython_parser as parser;
 use rustpython_parser::ast::{Location, StmtKind, Suite};
 
-use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit};
+use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::is_docstring_stmt;
 use ruff_python_ast::imports::{Alias, AnyImport, Import, ImportFrom};
 use ruff_python_ast::source_code::{Locator, Stylist};
 use ruff_python_ast::types::Range;
 
+use crate::importer::Importer;
 use crate::registry::Rule;
-use crate::rules::isort::helpers::{top_of_file_insertion, Insertion};
+use crate::rules::isort::track::Block;
 use crate::settings::{flags, Settings};
 
-use super::super::track::Block;
-
 /// ## What it does
 /// Adds any required imports, as specified by the user, to the top of the
 /// file.
@@ -109,21 +108,12 @@ fn add_required_import(
     }
 
     // Always insert the diagnostic at top-of-file.
-    let required_import = required_import.to_string();
     let mut diagnostic = Diagnostic::new(
-        MissingRequiredImport(required_import.clone()),
+        MissingRequiredImport(required_import.to_string()),
         Range::new(Location::default(), Location::default()),
     );
     if autofix.into() && settings.rules.should_fix(Rule::MissingRequiredImport) {
-        let Insertion {
+        diagnostic.set_fix(Importer::new(python_ast, locator, stylist).add_import(required_import));
-            prefix,
-            location,
-            suffix,
-        } = top_of_file_insertion(python_ast, locator, stylist);
-        diagnostic.set_fix(Edit::insertion(
-            format!("{prefix}{required_import}{suffix}"),
-            location,
-        ));
     }
     Some(diagnostic)
 }

crates/ruff/src/rules/pylint/rules/sys_exit_alias.rs

@@ -1,9 +1,10 @@
 use rustpython_parser::ast::{Expr, ExprKind};
 
-use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Violation};
+use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::types::Range;
 
+use crate::autofix::helpers::get_or_import_symbol;
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
 
@@ -45,13 +46,18 @@ pub fn sys_exit_alias(checker: &mut Checker, func: &Expr) {
             Range::from(func),
         );
         if checker.patch(diagnostic.kind.rule()) {
-            if let Some(binding) = checker.ctx.resolve_qualified_import_name("sys", "exit") {
+            diagnostic.try_set_fix(|| {
-                diagnostic.set_fix(Edit::replacement(
+                let (import_edit, binding) = get_or_import_symbol(
-                    binding,
+                    "sys",
-                    func.location,
+                    "exit",
-                    func.end_location.unwrap(),
+                    &checker.ctx,
-                ));
+                    &checker.importer,
-            }
+                    checker.locator,
+                )?;
+                let reference_edit =
+                    Edit::replacement(binding, func.location, func.end_location.unwrap());
+                Ok(Fix::from_iter([import_edit, reference_edit]))
+            });
         }
         checker.diagnostics.push(diagnostic);
     }

crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_0.py.snap

@@ -14,7 +14,21 @@ expression: diagnostics
     row: 1
     column: 4
   fix:
-    edits: []
+    edits:
+      - content: "import sys\n"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 0
+      - content: sys.exit
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 4
   parent: ~
 - kind:
     name: SysExitAlias
@@ -28,7 +42,21 @@ expression: diagnostics
     row: 2
     column: 4
   fix:
-    edits: []
+    edits:
+      - content: "import sys\n"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 0
+      - content: sys.exit
+        location:
+          row: 2
+          column: 0
+        end_location:
+          row: 2
+          column: 4
   parent: ~
 - kind:
     name: SysExitAlias
@@ -42,7 +70,21 @@ expression: diagnostics
     row: 6
     column: 8
   fix:
-    edits: []
+    edits:
+      - content: "import sys\n"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 0
+      - content: sys.exit
+        location:
+          row: 6
+          column: 4
+        end_location:
+          row: 6
+          column: 8
   parent: ~
 - kind:
     name: SysExitAlias
@@ -56,6 +98,20 @@ expression: diagnostics
     row: 7
     column: 8
   fix:
-    edits: []
+    edits:
+      - content: "import sys\n"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 0
+      - content: sys.exit
+        location:
+          row: 7
+          column: 4
+        end_location:
+          row: 7
+          column: 8
   parent: ~
 

crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_1.py.snap

@@ -15,6 +15,13 @@ expression: diagnostics
     column: 4
   fix:
     edits:
+      - content: import sys
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 10
       - content: sys.exit
         location:
           row: 3
@@ -36,6 +43,13 @@ expression: diagnostics
     column: 4
   fix:
     edits:
+      - content: import sys
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 10
       - content: sys.exit
         location:
           row: 4
@@ -57,6 +71,13 @@ expression: diagnostics
     column: 8
   fix:
     edits:
+      - content: import sys
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 10
       - content: sys.exit
         location:
           row: 8
@@ -78,6 +99,13 @@ expression: diagnostics
     column: 8
   fix:
     edits:
+      - content: import sys
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 10
       - content: sys.exit
         location:
           row: 9
@@ -86,4 +114,32 @@ expression: diagnostics
           row: 9
           column: 8
   parent: ~
+- kind:
+    name: SysExitAlias
+    body: "Use `sys.exit()` instead of `exit`"
+    suggestion: "Replace `exit` with `sys.exit()`"
+    fixable: true
+  location:
+    row: 15
+    column: 4
+  end_location:
+    row: 15
+    column: 8
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: SysExitAlias
+    body: "Use `sys.exit()` instead of `quit`"
+    suggestion: "Replace `quit` with `sys.exit()`"
+    fixable: true
+  location:
+    row: 16
+    column: 4
+  end_location:
+    row: 16
+    column: 8
+  fix:
+    edits: []
+  parent: ~
 

crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_2.py.snap

@@ -15,6 +15,13 @@ expression: diagnostics
     column: 4
   fix:
     edits:
+      - content: import sys as sys2
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 18
       - content: sys2.exit
         location:
           row: 3
@@ -36,6 +43,13 @@ expression: diagnostics
     column: 4
   fix:
     edits:
+      - content: import sys as sys2
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 18
       - content: sys2.exit
         location:
           row: 4
@@ -57,6 +71,13 @@ expression: diagnostics
     column: 8
   fix:
     edits:
+      - content: import sys as sys2
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 18
       - content: sys2.exit
         location:
           row: 8
@@ -78,6 +99,13 @@ expression: diagnostics
     column: 8
   fix:
     edits:
+      - content: import sys as sys2
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 18
       - content: sys2.exit
         location:
           row: 9

crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_3.py.snap

@@ -15,6 +15,13 @@ expression: diagnostics
     column: 4
   fix:
     edits:
+      - content: from sys import exit
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 20
       - content: exit
         location:
           row: 4
@@ -36,6 +43,13 @@ expression: diagnostics
     column: 8
   fix:
     edits:
+      - content: from sys import exit
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 20
       - content: exit
         location:
           row: 9

crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_4.py.snap

@@ -15,6 +15,13 @@ expression: diagnostics
     column: 4
   fix:
     edits:
+      - content: from sys import exit as exit2
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 29
       - content: exit2
         location:
           row: 3
@@ -36,6 +43,13 @@ expression: diagnostics
     column: 4
   fix:
     edits:
+      - content: from sys import exit as exit2
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 29
       - content: exit2
         location:
           row: 4
@@ -57,6 +71,13 @@ expression: diagnostics
     column: 8
   fix:
     edits:
+      - content: from sys import exit as exit2
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 29
       - content: exit2
         location:
           row: 8
@@ -78,6 +99,13 @@ expression: diagnostics
     column: 8
   fix:
     edits:
+      - content: from sys import exit as exit2
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 29
       - content: exit2
         location:
           row: 9

crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_6.py.snap

@@ -14,7 +14,21 @@ expression: diagnostics
     row: 1
     column: 4
   fix:
-    edits: []
+    edits:
+      - content: "import sys\n"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 0
+      - content: sys.exit
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 4
   parent: ~
 - kind:
     name: SysExitAlias
@@ -28,6 +42,20 @@ expression: diagnostics
     row: 2
     column: 4
   fix:
-    edits: []
+    edits:
+      - content: "import sys\n"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 0
+      - content: sys.exit
+        location:
+          row: 2
+          column: 0
+        end_location:
+          row: 2
+          column: 4
   parent: ~
 

crates/ruff/src/rules/pyupgrade/rules/lru_cache_with_maxsize_none.rs

@@ -1,9 +1,10 @@
 use rustpython_parser::ast::{Constant, Expr, ExprKind, KeywordData};
 
-use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit};
+use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::types::Range;
 
+use crate::autofix::helpers::get_or_import_symbol;
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
 
@@ -57,16 +58,18 @@ pub fn lru_cache_with_maxsize_none(checker: &mut Checker, decorator_list: &[Expr
                     Range::new(func.end_location.unwrap(), expr.end_location.unwrap()),
                 );
                 if checker.patch(diagnostic.kind.rule()) {
-                    if let Some(binding) = checker
+                    diagnostic.try_set_fix(|| {
-                        .ctx
+                        let (import_edit, binding) = get_or_import_symbol(
-                        .resolve_qualified_import_name("functools", "cache")
+                            "functools",
-                    {
+                            "cache",
-                        diagnostic.set_fix(Edit::replacement(
+                            &checker.ctx,
-                            binding,
+                            &checker.importer,
-                            expr.location,
+                            checker.locator,
-                            expr.end_location.unwrap(),
+                        )?;
-                        ));
+                        let reference_edit =
-                    }
+                            Edit::replacement(binding, expr.location, expr.end_location.unwrap());
+                        Ok(Fix::from_iter([import_edit, reference_edit]))
+                    });
                 }
                 checker.diagnostics.push(diagnostic);
             }

crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_0.py.snap

@@ -15,6 +15,13 @@ expression: diagnostics
     column: 34
   fix:
     edits:
+      - content: import functools
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 16
       - content: functools.cache
         location:
           row: 4
@@ -36,6 +43,13 @@ expression: diagnostics
     column: 34
   fix:
     edits:
+      - content: import functools
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 16
       - content: functools.cache
         location:
           row: 10
@@ -57,6 +71,13 @@ expression: diagnostics
     column: 34
   fix:
     edits:
+      - content: import functools
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 16
       - content: functools.cache
         location:
           row: 15

crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_1.py.snap

@@ -14,7 +14,21 @@ expression: diagnostics
     row: 4
     column: 24
   fix:
-    edits: []
+    edits:
+      - content: "from functools import lru_cache, cache"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 31
+      - content: cache
+        location:
+          row: 4
+          column: 1
+        end_location:
+          row: 4
+          column: 24
   parent: ~
 - kind:
     name: LRUCacheWithMaxsizeNone
@@ -28,7 +42,21 @@ expression: diagnostics
     row: 10
     column: 24
   fix:
-    edits: []
+    edits:
+      - content: "from functools import lru_cache, cache"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 31
+      - content: cache
+        location:
+          row: 10
+          column: 1
+        end_location:
+          row: 10
+          column: 24
   parent: ~
 - kind:
     name: LRUCacheWithMaxsizeNone
@@ -42,6 +70,20 @@ expression: diagnostics
     row: 15
     column: 24
   fix:
-    edits: []
+    edits:
+      - content: "from functools import lru_cache, cache"
+        location:
+          row: 1
+          column: 0
+        end_location:
+          row: 1
+          column: 31
+      - content: cache
+        location:
+          row: 15
+          column: 1
+        end_location:
+          row: 15
+          column: 24
   parent: ~
 

crates/ruff_python_ast/src/context.rs

@@ -219,10 +219,15 @@ impl<'a> Context<'a> {
     ///
     /// ...then `resolve_qualified_import_name("sys", "version_info")` will return
     /// `Some("python_version")`.
-    pub fn resolve_qualified_import_name(&self, module: &str, member: &str) -> Option<String> {
+    pub fn resolve_qualified_import_name(
+        &self,
+        module: &str,
+        member: &str,
+    ) -> Option<(&Stmt, String)> {
         self.scopes().enumerate().find_map(|(scope_index, scope)| {
             scope.binding_ids().find_map(|binding_index| {
-                match &self.bindings[*binding_index].kind {
+                let binding = &self.bindings[*binding_index];
+                match &binding.kind {
                     // Ex) Given `module="sys"` and `object="exit"`:
                     // `import sys`         -> `sys.exit`
                     // `import sys as sys2` -> `sys2.exit`
@@ -234,7 +239,10 @@ impl<'a> Context<'a> {
                                 .take(scope_index)
                                 .all(|scope| scope.get(name).is_none())
                             {
-                                return Some(format!("{name}.{member}"));
+                                return Some((
+                                    binding.source.as_ref().unwrap().into(),
+                                    format!("{name}.{member}"),
+                                ));
                             }
                         }
                     }
@@ -250,7 +258,10 @@ impl<'a> Context<'a> {
                                     .take(scope_index)
                                     .all(|scope| scope.get(name).is_none())
                                 {
-                                    return Some((*name).to_string());
+                                    return Some((
+                                        binding.source.as_ref().unwrap().into(),
+                                        (*name).to_string(),
+                                    ));
                                 }
                             }
                         }
@@ -265,7 +276,10 @@ impl<'a> Context<'a> {
                                 .take(scope_index)
                                 .all(|scope| scope.get(name).is_none())
                             {
-                                return Some(format!("{name}.{member}"));
+                                return Some((
+                                    binding.source.as_ref().unwrap().into(),
+                                    format!("{name}.{member}"),
+                                ));
                             }
                         }
                     }

crates/ruff_python_ast/src/imports.rs

@@ -27,6 +27,17 @@ pub struct Alias<'a> {
     pub as_name: Option<&'a str>,
 }
 
+impl<'a> Import<'a> {
+    pub fn module(name: &'a str) -> Self {
+        Self {
+            name: Alias {
+                name,
+                as_name: None,
+            },
+        }
+    }
+}
+
 impl fmt::Display for AnyImport<'_> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {