From ddd541b1987493a69ef5e8f65ef8030eda6198cd Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 17 May 2023 17:11:41 -0400 Subject: [PATCH] Move `Insertion` into its own module (#4478) --- crates/ruff/src/importer.rs | 365 -------------------------- crates/ruff/src/importer/insertion.rs | 261 ++++++++++++++++++ crates/ruff/src/importer/mod.rs | 111 ++++++++ 3 files changed, 372 insertions(+), 365 deletions(-) delete mode 100644 crates/ruff/src/importer.rs create mode 100644 crates/ruff/src/importer/insertion.rs create mode 100644 crates/ruff/src/importer/mod.rs diff --git a/crates/ruff/src/importer.rs b/crates/ruff/src/importer.rs deleted file mode 100644 index 82d03f71b8..0000000000 --- a/crates/ruff/src/importer.rs +++ /dev/null @@ -1,365 +0,0 @@ -//! 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 ruff_text_size::TextSize; -use rustpython_parser::ast::{self, Ranged, Stmt, 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>, - ordered_imports: Vec<&'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, - ordered_imports: Vec::default(), - } - } - - /// Visit a top-level import statement. - pub fn visit_import(&mut self, import: &'a Stmt) { - self.ordered_imports.push(import); - } - - /// Return the import statement that precedes the given position, if any. - fn preceding_import(&self, at: TextSize) -> Option<&Stmt> { - self.ordered_imports - .partition_point(|stmt| stmt.start() < at) - .checked_sub(1) - .map(|idx| self.ordered_imports[idx]) - } - - /// 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, at: TextSize) -> Edit { - let required_import = import.to_string(); - if let Some(stmt) = self.preceding_import(at) { - // 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 `Stmt::ImportFrom` - /// preceding the given position, if any. - pub fn find_import_from(&self, module: &str, at: TextSize) -> Option<&Stmt> { - let mut import_from = None; - for stmt in &self.ordered_imports { - if stmt.start() >= at { - break; - } - if let Stmt::ImportFrom(ast::StmtImportFrom { - module: name, - level, - .. - }) = stmt - { - if level.map_or(true, |level| level.to_u32() == 0) - && name.as_ref().map_or(false, |name| name == module) - { - import_from = Some(*stmt); - } - } - } - import_from - } - - /// Add the given member to an existing `Stmt::ImportFrom` statement. - pub fn add_member(&self, stmt: &Stmt, member: &str) -> Result { - let mut tree = match_module(self.locator.slice(stmt.range()))?; - 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::range_replacement(state.to_string(), stmt.range())) - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct Insertion { - /// The content to add before the insertion. - prefix: &'static str, - /// The location at which to insert. - location: TextSize, - /// The content to add after the insertion. - suffix: &'static str, -} - -impl Insertion { - fn new(prefix: &'static str, location: TextSize, 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()) -} - -/// Find the location at which an "end-of-statement" 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 -/// import math -/// -/// -/// def foo(): -/// pass -/// ``` -/// -/// The location returned will be the start of new line after the last -/// import statement, which in this case is the line after `import math`, -/// along with a trailing newline suffix. -fn end_of_statement_insertion(stmt: &Stmt, locator: &Locator, stylist: &Stylist) -> Insertion { - let location = stmt.end(); - let mut tokens = - lexer::lex_starts_at(locator.after(location), Mode::Module, location).flatten(); - if let Some((Tok::Semi, range)) = tokens.next() { - // If the first token after the docstring is a semicolon, insert after the semicolon as an - // inline statement; - Insertion::new(" ", range.end(), ";") - } else { - // Otherwise, insert on the next line. - Insertion::new( - "", - locator.full_line_end(location), - 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_starts_at(locator.after(location), Mode::Module, location) - .flatten() - .next(); - if let Some((Tok::Semi, range)) = first_token { - return Insertion::new(" ", range.end(), ";"); - } - - // Otherwise, advance to the next row. - locator.full_line_end(location) - } else { - TextSize::default() - }; - - // Skip over any comments and empty lines. - for (tok, range) in - lexer::lex_starts_at(locator.after(location), Mode::Module, location).flatten() - { - if matches!(tok, Tok::Comment(..) | Tok::Newline) { - location = locator.full_line_end(range.end()); - } else { - break; - } - } - - return Insertion::new("", location, stylist.line_ending().as_str()); -} - -#[cfg(test)] -mod tests { - use anyhow::Result; - use ruff_text_size::TextSize; - use rustpython_parser as parser; - use rustpython_parser::lexer::LexResult; - - use ruff_python_ast::newlines::LineEnding; - use ruff_python_ast::source_code::{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("", TextSize::from(0), LineEnding::default().as_str()) - ); - - let contents = r#" -"""Hello, world!""""# - .trim_start(); - assert_eq!( - insert(contents)?, - Insertion::new("", TextSize::from(19), LineEnding::default().as_str()) - ); - - let contents = r#" -"""Hello, world!""" -"# - .trim_start(); - assert_eq!( - insert(contents)?, - Insertion::new("", TextSize::from(20), "\n") - ); - - let contents = r#" -"""Hello, world!""" -"""Hello, world!""" -"# - .trim_start(); - assert_eq!( - insert(contents)?, - Insertion::new("", TextSize::from(40), "\n") - ); - - let contents = r#" -x = 1 -"# - .trim_start(); - assert_eq!( - insert(contents)?, - Insertion::new("", TextSize::from(0), "\n") - ); - - let contents = r#" -#!/usr/bin/env python3 -"# - .trim_start(); - assert_eq!( - insert(contents)?, - Insertion::new("", TextSize::from(23), "\n") - ); - - let contents = r#" -#!/usr/bin/env python3 -"""Hello, world!""" -"# - .trim_start(); - assert_eq!( - insert(contents)?, - Insertion::new("", TextSize::from(43), "\n") - ); - - let contents = r#" -"""Hello, world!""" -#!/usr/bin/env python3 -"# - .trim_start(); - assert_eq!( - insert(contents)?, - Insertion::new("", TextSize::from(43), "\n") - ); - - let contents = r#" -"""%s""" % "Hello, world!" -"# - .trim_start(); - assert_eq!( - insert(contents)?, - Insertion::new("", TextSize::from(0), "\n") - ); - - let contents = r#" -"""Hello, world!"""; x = 1 -"# - .trim_start(); - assert_eq!( - insert(contents)?, - Insertion::new(" ", TextSize::from(20), ";") - ); - - let contents = r#" -"""Hello, world!"""; x = 1; y = \ - 2 -"# - .trim_start(); - assert_eq!( - insert(contents)?, - Insertion::new(" ", TextSize::from(20), ";") - ); - - Ok(()) - } -} diff --git a/crates/ruff/src/importer/insertion.rs b/crates/ruff/src/importer/insertion.rs new file mode 100644 index 0000000000..7b3258d6bf --- /dev/null +++ b/crates/ruff/src/importer/insertion.rs @@ -0,0 +1,261 @@ +use ruff_diagnostics::Edit; +use ruff_text_size::TextSize; +use rustpython_parser::ast::{Ranged, Stmt}; +use rustpython_parser::{lexer, Mode, Tok}; + +use ruff_python_ast::helpers::is_docstring_stmt; +use ruff_python_ast::source_code::{Locator, Stylist}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct Insertion { + /// The content to add before the insertion. + prefix: &'static str, + /// The location at which to insert. + location: TextSize, + /// The content to add after the insertion. + suffix: &'static str, +} + +impl Insertion { + /// Create an [`Insertion`] to insert (e.g.) an import after the end of the given [`Stmt`], + /// along with a prefix and suffix to use for the insertion. + /// + /// For example, given the following code: + /// + /// ```python + /// """Hello, world!""" + /// + /// import os + /// import math + /// + /// + /// def foo(): + /// pass + /// ``` + /// + /// The insertion returned will begin after the newline after the last import statement, which + /// in this case is the line after `import math`, and will include a trailing newline suffix. + pub(super) fn end_of_statement(stmt: &Stmt, locator: &Locator, stylist: &Stylist) -> Insertion { + let location = stmt.end(); + let mut tokens = + lexer::lex_starts_at(locator.after(location), Mode::Module, location).flatten(); + if let Some((Tok::Semi, range)) = tokens.next() { + // If the first token after the docstring is a semicolon, insert after the semicolon as an + // inline statement; + Insertion::new(" ", range.end(), ";") + } else { + // Otherwise, insert on the next line. + Insertion::new( + "", + locator.full_line_end(location), + stylist.line_ending().as_str(), + ) + } + } + + /// Create an [`Insertion`] to insert (e.g.) an import statement at the "top" of a given file, + /// along with a prefix and suffix to use for the insertion. + /// + /// For example, given the following code: + /// + /// ```python + /// """Hello, world!""" + /// + /// import os + /// ``` + /// + /// The insertion returned will begin at the start of the `import os` statement, and will + /// include a trailing newline suffix. + pub(super) fn top_of_file(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_starts_at(locator.after(location), Mode::Module, location) + .flatten() + .next(); + if let Some((Tok::Semi, range)) = first_token { + return Insertion::new(" ", range.end(), ";"); + } + + // Otherwise, advance to the next row. + locator.full_line_end(location) + } else { + TextSize::default() + }; + + // Skip over any comments and empty lines. + for (tok, range) in + lexer::lex_starts_at(locator.after(location), Mode::Module, location).flatten() + { + if matches!(tok, Tok::Comment(..) | Tok::Newline) { + location = locator.full_line_end(range.end()); + } else { + break; + } + } + + Insertion::new("", location, stylist.line_ending().as_str()) + } + + fn new(prefix: &'static str, location: TextSize, suffix: &'static str) -> Self { + Self { + prefix, + location, + suffix, + } + } + + /// Convert this [`Insertion`] into an [`Edit`] that inserts the given content. + pub(super) fn into_edit(self, content: &str) -> Edit { + let Insertion { + prefix, + location, + suffix, + } = self; + Edit::insertion(format!("{prefix}{content}{suffix}"), location) + } +} + +/// 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()) +} + +#[cfg(test)] +mod tests { + use anyhow::Result; + use ruff_text_size::TextSize; + use rustpython_parser as parser; + use rustpython_parser::lexer::LexResult; + + use ruff_python_ast::newlines::LineEnding; + use ruff_python_ast::source_code::{Locator, Stylist}; + + use super::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(Insertion::top_of_file(&program, &locator, &stylist)) + } + + #[test] + fn top_of_file() -> Result<()> { + let contents = ""; + assert_eq!( + insert(contents)?, + Insertion::new("", TextSize::from(0), LineEnding::default().as_str()) + ); + + let contents = r#" +"""Hello, world!""""# + .trim_start(); + assert_eq!( + insert(contents)?, + Insertion::new("", TextSize::from(19), LineEnding::default().as_str()) + ); + + let contents = r#" +"""Hello, world!""" +"# + .trim_start(); + assert_eq!( + insert(contents)?, + Insertion::new("", TextSize::from(20), "\n") + ); + + let contents = r#" +"""Hello, world!""" +"""Hello, world!""" +"# + .trim_start(); + assert_eq!( + insert(contents)?, + Insertion::new("", TextSize::from(40), "\n") + ); + + let contents = r#" +x = 1 +"# + .trim_start(); + assert_eq!( + insert(contents)?, + Insertion::new("", TextSize::from(0), "\n") + ); + + let contents = r#" +#!/usr/bin/env python3 +"# + .trim_start(); + assert_eq!( + insert(contents)?, + Insertion::new("", TextSize::from(23), "\n") + ); + + let contents = r#" +#!/usr/bin/env python3 +"""Hello, world!""" +"# + .trim_start(); + assert_eq!( + insert(contents)?, + Insertion::new("", TextSize::from(43), "\n") + ); + + let contents = r#" +"""Hello, world!""" +#!/usr/bin/env python3 +"# + .trim_start(); + assert_eq!( + insert(contents)?, + Insertion::new("", TextSize::from(43), "\n") + ); + + let contents = r#" +"""%s""" % "Hello, world!" +"# + .trim_start(); + assert_eq!( + insert(contents)?, + Insertion::new("", TextSize::from(0), "\n") + ); + + let contents = r#" +"""Hello, world!"""; x = 1 +"# + .trim_start(); + assert_eq!( + insert(contents)?, + Insertion::new(" ", TextSize::from(20), ";") + ); + + let contents = r#" +"""Hello, world!"""; x = 1; y = \ + 2 +"# + .trim_start(); + assert_eq!( + insert(contents)?, + Insertion::new(" ", TextSize::from(20), ";") + ); + + Ok(()) + } +} diff --git a/crates/ruff/src/importer/mod.rs b/crates/ruff/src/importer/mod.rs new file mode 100644 index 0000000000..4d0ac0d9fa --- /dev/null +++ b/crates/ruff/src/importer/mod.rs @@ -0,0 +1,111 @@ +//! 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 ruff_text_size::TextSize; +use rustpython_parser::ast::{self, Ranged, Stmt, Suite}; + +use ruff_diagnostics::Edit; +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}; +use crate::importer::insertion::Insertion; + +mod insertion; + +pub struct Importer<'a> { + python_ast: &'a Suite, + locator: &'a Locator<'a>, + stylist: &'a Stylist<'a>, + ordered_imports: Vec<&'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, + ordered_imports: Vec::default(), + } + } + + /// Visit a top-level import statement. + pub fn visit_import(&mut self, import: &'a Stmt) { + self.ordered_imports.push(import); + } + + /// Return the import statement that precedes the given position, if any. + fn preceding_import(&self, at: TextSize) -> Option<&Stmt> { + self.ordered_imports + .partition_point(|stmt| stmt.start() < at) + .checked_sub(1) + .map(|idx| self.ordered_imports[idx]) + } + + /// 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, at: TextSize) -> Edit { + let required_import = import.to_string(); + if let Some(stmt) = self.preceding_import(at) { + // Insert after the last top-level import. + Insertion::end_of_statement(stmt, self.locator, self.stylist) + .into_edit(&required_import) + } else { + // Insert at the top of the file. + Insertion::top_of_file(self.python_ast, self.locator, self.stylist) + .into_edit(&required_import) + } + } + + /// Return the top-level [`Stmt`] that imports the given module using `Stmt::ImportFrom` + /// preceding the given position, if any. + pub fn find_import_from(&self, module: &str, at: TextSize) -> Option<&Stmt> { + let mut import_from = None; + for stmt in &self.ordered_imports { + if stmt.start() >= at { + break; + } + if let Stmt::ImportFrom(ast::StmtImportFrom { + module: name, + level, + .. + }) = stmt + { + if level.map_or(true, |level| level.to_u32() == 0) + && name.as_ref().map_or(false, |name| name == module) + { + import_from = Some(*stmt); + } + } + } + import_from + } + + /// Add the given member to an existing `Stmt::ImportFrom` statement. + pub fn add_member(&self, stmt: &Stmt, member: &str) -> Result { + let mut tree = match_module(self.locator.slice(stmt.range()))?; + 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::range_replacement(state.to_string(), stmt.range())) + } +}