diff --git a/crates/ruff_python_formatter/src/cli.rs b/crates/ruff_python_formatter/src/cli.rs index 87d7dc06aa..f605e66407 100644 --- a/crates/ruff_python_formatter/src/cli.rs +++ b/crates/ruff_python_formatter/src/cli.rs @@ -1,11 +1,75 @@ +#![allow(clippy::print_stdout)] + use std::path::PathBuf; -use clap::{command, Parser}; +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 { + /// Write back to the original files + Files, + /// Write to stdout + Stdout, +} #[derive(Parser)] #[command(author, version, about, long_about = None)] pub struct Cli { - /// Python file to round-trip. - #[arg(required = true)] - pub file: PathBuf, + /// Python files to format. If there are none, stdin will be used. `-` as stdin is not supported + pub files: Vec, + #[clap(long)] + pub emit: Option, + /// Run in 'check' mode. Exits with 0 if input is formatted correctly. Exits with 1 and prints + /// 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 19178aa6ac..63b789888c 100644 --- a/crates/ruff_python_formatter/src/main.rs +++ b/crates/ruff_python_formatter/src/main.rs @@ -1,15 +1,53 @@ -use std::fs; +use std::io::{stdout, Read, Write}; +use std::{fs, io}; -use anyhow::Result; +use anyhow::{bail, Context, Result}; use clap::Parser as ClapParser; -use ruff_python_formatter::cli::Cli; -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 { + let mut buffer = String::new(); + io::stdin().lock().read_to_string(&mut buffer)?; + Ok(buffer) +} #[allow(clippy::print_stdout)] fn main() -> Result<()> { - let cli = Cli::parse(); - let contents = fs::read_to_string(cli.file)?; - println!("{}", format_module(&contents)?.as_code()); + let cli: Cli = Cli::parse(); + + if cli.files.is_empty() { + if !matches!(cli.emit, None | Some(Emit::Stdout)) { + bail!( + "Can only write to stdout when formatting from stdin, but you asked for {:?}", + cli.emit + ); + } + let input = read_from_stdin()?; + let formatted = format_and_debug_print(&input, &cli)?; + if cli.check { + if formatted == input { + return Ok(()); + } + bail!("Content not correctly formatted") + } + stdout().lock().write_all(formatted.as_bytes())?; + } else { + for file in &cli.files { + let input = fs::read_to_string(file) + .with_context(|| format!("Could not read {}: ", file.display()))?; + let formatted = format_and_debug_print(&input, &cli)?; + match cli.emit { + Some(Emit::Stdout) => stdout().lock().write_all(formatted.as_bytes())?, + None | Some(Emit::Files) => { + fs::write(file, formatted.as_bytes()).with_context(|| { + format!("Could not write to {}, exiting", file.display()) + })?; + } + } + } + } + Ok(()) }