mod generator; mod indexer; mod line_index; mod locator; mod stylist; pub use crate::source_code::line_index::{LineIndex, OneIndexed}; use crate::types::Range; pub use generator::Generator; pub use indexer::Indexer; pub use locator::Locator; use ruff_text_size::{TextRange, TextSize}; use rustpython_parser as parser; use rustpython_parser::ast::Location; use rustpython_parser::{lexer, Mode, ParseError}; use std::fmt::{Debug, Formatter}; use std::sync::Arc; pub use stylist::{LineEnding, Stylist}; /// Run round-trip source code generation on a given Python code. pub fn round_trip(code: &str, source_path: &str) -> Result { let locator = Locator::new(code); let python_ast = parser::parse_program(code, source_path)?; let tokens: Vec<_> = lexer::lex(code, Mode::Module).collect(); let stylist = Stylist::from_tokens(&tokens, &locator); let mut generator: Generator = (&stylist).into(); generator.unparse_suite(&python_ast); Ok(generator.generate()) } /// Gives access to the source code of a file and allows mapping between [`Location`] and byte offsets. #[derive(Debug)] pub struct SourceCode<'src, 'index> { text: &'src str, index: &'index LineIndex, } impl<'src, 'index> SourceCode<'src, 'index> { pub fn new(content: &'src str, index: &'index LineIndex) -> Self { Self { text: content, index, } } /// Take the source code up to the given [`Location`]. pub fn up_to(&self, location: Location) -> &'src str { let offset = self.index.location_offset(location, self.text); &self.text[TextRange::up_to(offset)] } /// Take the source code after the given [`Location`]. pub fn after(&self, location: Location) -> &'src str { let offset = self.index.location_offset(location, self.text); &self.text[usize::from(offset)..] } /// Take the source code between the given [`Range`]. pub fn slice>(&self, range: R) -> &'src str { let range = self.text_range(range); &self.text[range] } /// Converts a [`Location`] range to a byte offset range pub fn text_range>(&self, range: R) -> TextRange { let range = range.into(); let start = self.index.location_offset(range.location, self.text); let end = self.index.location_offset(range.end_location, self.text); TextRange::new(start, end) } /// Return the byte offset of the given [`Location`]. pub fn offset(&self, location: Location) -> TextSize { self.index.location_offset(location, self.text) } pub fn line_start(&self, line: OneIndexed) -> TextSize { self.index.line_start(line, self.text) } pub fn line_end(&self, line: OneIndexed) -> TextSize { self.index.line_end(line, self.text) } pub fn line_range(&self, line: OneIndexed) -> TextRange { self.index.line_range(line, self.text) } /// Returns a string with the lines spawning between location and end location. pub fn lines(&self, range: Range) -> &'src str { let start_line = self .index .line_range(OneIndexed::new(range.location.row()).unwrap(), self.text); let end_line = self.index.line_range( OneIndexed::new(range.end_location.row()).unwrap(), self.text, ); &self.text[TextRange::new(start_line.start(), end_line.end())] } /// Returns the source text of the line with the given index #[inline] pub fn line_text(&self, index: OneIndexed) -> &'src str { let range = self.index.line_range(index, self.text); &self.text[range] } /// Returns the source text pub fn text(&self) -> &'src str { self.text } /// Returns the number of lines #[inline] pub fn line_count(&self) -> usize { self.index.line_count() } } impl PartialEq for SourceCode<'_, '_> { fn eq(&self, other: &Self) -> bool { self.text == other.text } } impl Eq for SourceCode<'_, '_> {} /// A Builder for constructing a [`SourceFile`] pub struct SourceFileBuilder { name: Box, code: Option, } impl SourceFileBuilder { /// Creates a new builder for a file named `name`. pub fn new(name: &str) -> Self { Self { name: Box::from(name), code: None, } } /// Creates a enw builder for a file named `name` pub fn from_string(name: String) -> Self { Self { name: Box::from(name), code: None, } } /// Consumes `self` and returns a builder for a file with the source text and the [`LineIndex`] copied /// from `source`. #[must_use] pub fn source_code(mut self, source: &SourceCode) -> Self { self.set_source_code(source); self } /// Copies the source text and [`LineIndex`] from `source`. pub fn set_source_code(&mut self, source: &SourceCode) { self.code = Some(FileSourceCode { text: Box::from(source.text()), index: source.index.clone(), }); } /// Consumes `self` and returns a builder for a file with the source text `text`. Builds the [`LineIndex`] from `text`. #[must_use] pub fn source_text(self, text: &str) -> Self { self.source_code(&SourceCode::new(text, &LineIndex::from_source_text(text))) } /// Consumes `self` and returns a builder for a file with the source text `text`. Builds the [`LineIndex`] from `text`. #[must_use] pub fn source_text_string(mut self, text: String) -> Self { self.set_source_text_string(text); self } /// Copies the source text `text` and builds the [`LineIndex`] from `text`. pub fn set_source_text_string(&mut self, text: String) { self.code = Some(FileSourceCode { index: LineIndex::from_source_text(&text), text: Box::from(text), }); } /// Consumes `self` and returns the [`SourceFile`]. pub fn finish(self) -> SourceFile { SourceFile { inner: Arc::new(SourceFileInner { name: self.name, code: self.code, }), } } } /// A source file that is identified by its name. Optionally stores the source code and [`LineIndex`]. /// /// Cloning a [`SourceFile`] is cheap, because it only requires bumping a reference count. #[derive(Clone, Eq, PartialEq)] pub struct SourceFile { inner: Arc, } impl Debug for SourceFile { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("SourceFile") .field("name", &self.name()) .field("code", &self.source_code()) .finish() } } impl SourceFile { /// Returns the name of the source file (filename). #[inline] pub fn name(&self) -> &str { &self.inner.name } /// Returns `Some` with the source code if set, or `None`. #[inline] pub fn source_code(&self) -> Option { self.inner.code.as_ref().map(|code| SourceCode { text: &code.text, index: &code.index, }) } /// Returns `Some` with the source text if set, or `None`. #[inline] pub fn source_text(&self) -> Option<&str> { self.inner.code.as_ref().map(|code| &*code.text) } } #[derive(Eq, PartialEq)] struct SourceFileInner { name: Box, code: Option, } struct FileSourceCode { text: Box, index: LineIndex, } impl PartialEq for FileSourceCode { fn eq(&self, other: &Self) -> bool { // It should be safe to assume that the index for two source files are identical self.text == other.text } } impl Eq for FileSourceCode {}