diff --git a/crates/ruff/resources/test/fixtures/pylint/sys_exit_alias_1.py b/crates/ruff/resources/test/fixtures/pylint/sys_exit_alias_1.py index 006b7909f2..af245b8d47 100644 --- a/crates/ruff/resources/test/fixtures/pylint/sys_exit_alias_1.py +++ b/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) diff --git a/crates/ruff/src/autofix/helpers.rs b/crates/ruff/src/autofix/helpers.rs index f1831cc1d7..020b86893f 100644 --- a/crates/ruff/src/autofix/helpers.rs +++ b/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; diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs index 9db1769839..ad5ed6444a 100644 --- a/crates/ruff/src/checkers/ast/mod.rs +++ b/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(); diff --git a/crates/ruff/src/cst/matchers.rs b/crates/ruff/src/cst/matchers.rs index 9a5ea83a38..fbebdca0a8 100644 --- a/crates/ruff/src/cst/matchers.rs +++ b/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, - SmallStatement, Statement, + Attribute, Call, Comparison, Dict, Expr, Expression, Import, ImportAlias, ImportFrom, + ImportNames, Module, SimpleString, SmallStatement, Statement, }; pub fn match_module(module_text: &str) -> Result { @@ -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>> { + 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) diff --git a/crates/ruff/src/importer.rs b/crates/ruff/src/importer.rs new file mode 100644 index 0000000000..19db7cf48c --- /dev/null +++ b/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 { + 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 { + 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 { + let program = parser::parse_program(contents, "")?; + let tokens: Vec = 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(()) + } +} diff --git a/crates/ruff/src/lib.rs b/crates/ruff/src/lib.rs index 7736b7a7eb..df12559b1b 100644 --- a/crates/ruff/src/lib.rs +++ b/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; diff --git a/crates/ruff/src/rules/isort/helpers.rs b/crates/ruff/src/rules/isort/helpers.rs index cdec6c1ace..510c958983 100644 --- a/crates/ruff/src/rules/isort/helpers.rs +++ b/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 { - 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 { - let program = parser::parse_program(contents, "")?; - let tokens: Vec = 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(()) - } -} diff --git a/crates/ruff/src/rules/isort/rules/add_required_imports.rs b/crates/ruff/src/rules/isort/rules/add_required_imports.rs index 0ab96e53f3..b3c9975e79 100644 --- a/crates/ruff/src/rules/isort/rules/add_required_imports.rs +++ b/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 { - prefix, - location, - suffix, - } = top_of_file_insertion(python_ast, locator, stylist); - diagnostic.set_fix(Edit::insertion( - format!("{prefix}{required_import}{suffix}"), - location, - )); + diagnostic.set_fix(Importer::new(python_ast, locator, stylist).add_import(required_import)); } Some(diagnostic) } diff --git a/crates/ruff/src/rules/pylint/rules/sys_exit_alias.rs b/crates/ruff/src/rules/pylint/rules/sys_exit_alias.rs index b8d9ee9bc3..b0db6067a4 100644 --- a/crates/ruff/src/rules/pylint/rules/sys_exit_alias.rs +++ b/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.set_fix(Edit::replacement( - binding, - func.location, - func.end_location.unwrap(), - )); - } + diagnostic.try_set_fix(|| { + let (import_edit, binding) = get_or_import_symbol( + "sys", + "exit", + &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); } diff --git a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_0.py.snap b/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_0.py.snap index 1ac50fdd49..c0ef383624 100644 --- a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_0.py.snap +++ b/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: ~ diff --git a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_1.py.snap b/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_1.py.snap index 9194663f6c..6bc29cbca8 100644 --- a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_1.py.snap +++ b/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: ~ diff --git a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_2.py.snap b/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_2.py.snap index aaae64b983..9de290bf6f 100644 --- a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_2.py.snap +++ b/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 diff --git a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_3.py.snap b/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_3.py.snap index 52d2f6d5a4..801a625462 100644 --- a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_3.py.snap +++ b/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 diff --git a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_4.py.snap b/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_4.py.snap index 73ad483c42..3c39d1915e 100644 --- a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_4.py.snap +++ b/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 diff --git a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_6.py.snap b/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_6.py.snap index 1f55fd9ccf..4d1e19362f 100644 --- a/crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_6.py.snap +++ b/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: ~ diff --git a/crates/ruff/src/rules/pyupgrade/rules/lru_cache_with_maxsize_none.rs b/crates/ruff/src/rules/pyupgrade/rules/lru_cache_with_maxsize_none.rs index d2bcbe7ea2..db512c3c53 100644 --- a/crates/ruff/src/rules/pyupgrade/rules/lru_cache_with_maxsize_none.rs +++ b/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 - .ctx - .resolve_qualified_import_name("functools", "cache") - { - diagnostic.set_fix(Edit::replacement( - binding, - expr.location, - expr.end_location.unwrap(), - )); - } + diagnostic.try_set_fix(|| { + let (import_edit, binding) = get_or_import_symbol( + "functools", + "cache", + &checker.ctx, + &checker.importer, + checker.locator, + )?; + let reference_edit = + Edit::replacement(binding, expr.location, expr.end_location.unwrap()); + Ok(Fix::from_iter([import_edit, reference_edit])) + }); } checker.diagnostics.push(diagnostic); } diff --git a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_0.py.snap b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_0.py.snap index 6e0425c16e..97069eddb0 100644 --- a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_0.py.snap +++ b/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 diff --git a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_1.py.snap b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_1.py.snap index 41b265ae32..1eee466d68 100644 --- a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_1.py.snap +++ b/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: ~ diff --git a/crates/ruff_python_ast/src/context.rs b/crates/ruff_python_ast/src/context.rs index 5e8673efa8..e0e4f8e7f9 100644 --- a/crates/ruff_python_ast/src/context.rs +++ b/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 { + 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}"), + )); } } } diff --git a/crates/ruff_python_ast/src/imports.rs b/crates/ruff_python_ast/src/imports.rs index 0565efc89a..61f8c46a5e 100644 --- a/crates/ruff_python_ast/src/imports.rs +++ b/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 {