-
Notifications
You must be signed in to change notification settings - Fork 919
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
aed248b
commit 808cb2d
Showing
18 changed files
with
713 additions
and
33 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -9,4 +9,8 @@ def main(): | |
| quit(1) | ||
|
|
||
|
|
||
| sys.exit(2) | ||
| def main(): | ||
| sys = 1 | ||
|
|
||
| exit(1) | ||
| quit(1) | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,217 @@ | ||
| //! Add and modify import statements to make symbols available during code generation. | ||
|
|
||
| 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::source_code::{Locator, Stylist}; | ||
|
|
||
| use crate::cst::matchers::{match_aliases, match_import_from, match_module}; | ||
| use crate::imports::AnyImport; | ||
|
|
||
| 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. | ||
| } | ||
| 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); | ||
| } | ||
| } | ||
| } | ||
| _ => { | ||
| unreachable!("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 (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 (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(), | ||
| )) | ||
| } | ||
| } | ||
|
|
||
| /// 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. | ||
| pub fn end_of_statement_insertion( | ||
| stmt: &Stmt, | ||
| locator: &Locator, | ||
| stylist: &Stylist, | ||
| ) -> (&'static str, Location, &'static str) { | ||
| 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; | ||
| (" ", end, ";") | ||
| } else { | ||
| // Otherwise, insert on the next line. | ||
| ( | ||
| "", | ||
| 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. | ||
| pub fn top_of_file_insertion( | ||
| body: &[Stmt], | ||
| locator: &Locator, | ||
| stylist: &Stylist, | ||
| ) -> (&'static str, Location, &'static str) { | ||
| // Skip over any docstrings. | ||
| let mut location = if let Some(location) = match_docstring_end(body) { | ||
| let mut tokens = lexer::lex_located(locator.skip(location), Mode::Module, location) | ||
| .flatten() | ||
| .peekable(); | ||
|
|
||
| // If the first token after the docstring is a semicolon, insert after the semicolon as an | ||
| // inline statement; | ||
| if let Some((.., Tok::Semi, end)) = tokens.peek() { | ||
| return (" ", *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 ("", location, stylist.line_ending().as_str()); | ||
| } |
Oops, something went wrong.