mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-08-04 10:48:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Add import insertion support to autofix capabilities (#3787)

Browse source
This commit is contained in:
Charlie Marsh 2023-03-30 11:33:46 -04:00 • committed by GitHub
parent d7113d3995
commit 01357f62e5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
20 changed files with 820 additions and 265 deletions
Download patch file Download diff file Expand all files Collapse all files

6
crates/ruff/resources/test/fixtures/pylint/sys_exit_alias_1.py vendored
View file

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

79
crates/ruff/src/autofix/helpers.rs
View file

@ -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;

17
crates/ruff/src/checkers/ast/mod.rs
View file

@ -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();

14
crates/ruff/src/cst/matchers.rs
View file

@ -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<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)

359
crates/ruff/src/importer.rs Normal file
View file

@ -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(())
}
}

1
crates/ruff/src/lib.rs
View file

@ -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;

216
crates/ruff/src/rules/isort/helpers.rs
View file

@ -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(())
}
}

20
crates/ruff/src/rules/isort/rules/add_required_imports.rs
View file

@ -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)
}

22
crates/ruff/src/rules/pylint/rules/sys_exit_alias.rs
View file

@ -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);
}

64
crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_0.py.snap
View file

@ -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: ~

56
crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_1.py.snap
View file

@ -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: ~

28
crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_2.py.snap
View file

@ -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

14
crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_3.py.snap
View file

@ -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

28
crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_4.py.snap
View file

@ -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

32
crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_6.py.snap
View file

@ -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: ~

25
crates/ruff/src/rules/pyupgrade/rules/lru_cache_with_maxsize_none.rs
View file

@ -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);
}

21
crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_0.py.snap
View file

@ -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

48
crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_1.py.snap
View file

@ -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: ~

24
crates/ruff_python_ast/src/context.rs
View file

@ -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}"),
));
}
}
}

11
crates/ruff_python_ast/src/imports.rs
View file

@ -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 {