diff --git a/crates/ruff_python_formatter/src/cli.rs b/crates/ruff_python_formatter/src/cli.rs index a4f5664f0a705..f605e6640768a 100644 --- a/crates/ruff_python_formatter/src/cli.rs +++ b/crates/ruff_python_formatter/src/cli.rs @@ -1,6 +1,16 @@ +#![allow(clippy::print_stdout)] + use std::path::PathBuf; +use anyhow::{bail, Context, Result}; use clap::{command, Parser, ValueEnum}; +use rustpython_parser::lexer::lex; +use rustpython_parser::{parse_tokens, Mode}; + +use ruff_formatter::SourceCode; +use ruff_python_ast::source_code::CommentRangesBuilder; + +use crate::format_node; #[derive(ValueEnum, Clone, Debug)] pub enum Emit { @@ -21,4 +31,45 @@ pub struct Cli { /// a diff if formatting is required. #[clap(long)] pub check: bool, + #[clap(long)] + pub print_ir: bool, + #[clap(long)] + pub print_comments: bool, +} + +pub fn format_and_debug_print(input: &str, cli: &Cli) -> Result { + let mut tokens = Vec::new(); + let mut comment_ranges = CommentRangesBuilder::default(); + + for result in lex(input, Mode::Module) { + let (token, range) = match result { + Ok((token, range)) => (token, range), + Err(err) => bail!("Source contains syntax errors {err:?}"), + }; + + comment_ranges.visit_token(&token, range); + tokens.push(Ok((token, range))); + } + + let comment_ranges = comment_ranges.finish(); + + // Parse the AST. + let python_ast = parse_tokens(tokens, Mode::Module, "") + .with_context(|| "Syntax error in input")?; + + let formatted = format_node(&python_ast, &comment_ranges, input)?; + if cli.print_ir { + println!("{}", formatted.document().display(SourceCode::new(input))); + } + if cli.print_comments { + println!( + "{:?}", + formatted.context().comments().debug(SourceCode::new(input)) + ); + } + Ok(formatted + .print() + .with_context(|| "Failed to print the formatter IR")? + .as_code() + .to_string()) } diff --git a/crates/ruff_python_formatter/src/main.rs b/crates/ruff_python_formatter/src/main.rs index f4ce6d255837c..63b789888cc99 100644 --- a/crates/ruff_python_formatter/src/main.rs +++ b/crates/ruff_python_formatter/src/main.rs @@ -4,8 +4,7 @@ use std::{fs, io}; use anyhow::{bail, Context, Result}; use clap::Parser as ClapParser; -use ruff_python_formatter::cli::{Cli, Emit}; -use ruff_python_formatter::format_module; +use ruff_python_formatter::cli::{format_and_debug_print, Cli, Emit}; /// Read a `String` from `stdin`. pub(crate) fn read_from_stdin() -> Result { @@ -26,23 +25,23 @@ fn main() -> Result<()> { ); } let input = read_from_stdin()?; - let formatted = format_module(&input)?; + let formatted = format_and_debug_print(&input, &cli)?; if cli.check { - if formatted.as_code() == input { + if formatted == input { return Ok(()); } bail!("Content not correctly formatted") } - stdout().lock().write_all(formatted.as_code().as_bytes())?; + stdout().lock().write_all(formatted.as_bytes())?; } else { - for file in cli.files { - let unformatted = fs::read_to_string(&file) + for file in &cli.files { + let input = fs::read_to_string(file) .with_context(|| format!("Could not read {}: ", file.display()))?; - let formatted = format_module(&unformatted)?; + let formatted = format_and_debug_print(&input, &cli)?; match cli.emit { - Some(Emit::Stdout) => stdout().lock().write_all(formatted.as_code().as_bytes())?, + Some(Emit::Stdout) => stdout().lock().write_all(formatted.as_bytes())?, None | Some(Emit::Files) => { - fs::write(&file, formatted.as_code().as_bytes()).with_context(|| { + fs::write(file, formatted.as_bytes()).with_context(|| { format!("Could not write to {}, exiting", file.display()) })?; }