Allow streaming formatting

src/bin/main.rs

@@ -11,7 +11,7 @@ use tracing_subscriber::EnvFilter;
 use std::collections::HashMap;
 use std::env;
 use std::fs::File;
-use std::io::{self, Read, Write, stdout};
+use std::io::{self, BufRead, Read, Write, stdout};
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 
@@ -69,6 +69,11 @@ enum Operation {
     ConfigOutputCurrent { path: Option<String> },
     /// No file specified, read from stdin
     Stdin { input: String },
+    /// No file specified, read from stdin streamingly
+    StdinStream {
+        start_marker: String,
+        end_marker: String,
+    },
 }
 
 /// Rustfmt operations errors.
@@ -167,6 +172,18 @@ fn make_opts() -> Options {
         "Set options from command line. These settings take priority over .rustfmt.toml",
         "[key1=val1,key2=val2...]",
     );
+    opts.optopt(
+        "",
+        "start-marker",
+        "Start marker for streaming formatting",
+        "[Marker string]",
+    );
+    opts.optopt(
+        "",
+        "end-marker",
+        "End marker for streaming formatting",
+        "[Marker string]",
+    );
     opts.optopt(
         "",
         "style-edition",
@@ -268,6 +285,10 @@ fn execute(opts: &Options) -> Result<i32> {
             Ok(0)
         }
         Operation::Stdin { input } => format_string(input, options),
+        Operation::StdinStream {
+            start_marker,
+            end_marker,
+        } => format_stdin_stream(start_marker, end_marker, options),
         Operation::Format {
             files,
             minimal_config_path,
@@ -328,6 +349,50 @@ fn format_string(input: String, options: GetOptsOptions) -> Result<i32> {
     Ok(exit_code)
 }
 
+fn format_stdin_stream(
+    start_marker: String,
+    end_marker: String,
+    options: GetOptsOptions,
+) -> Result<i32> {
+    let stdin = io::stdin();
+    let mut reader = stdin.lock();
+    let mut line = String::new();
+    let mut in_chunk = false;
+    let mut buffer = String::new();
+
+    loop {
+        line.clear();
+        let bytes_read = reader.read_line(&mut line)?;
+        if bytes_read == 0 {
+            break; // EOF
+        }
+
+        let trimmed = line.trim_end_matches(&['\r', '\n'][..]);
+        if !in_chunk {
+            if trimmed == start_marker {
+                in_chunk = true;
+                buffer.clear();
+            }
+        } else {
+            if trimmed == end_marker {
+                in_chunk = false;
+                match format_string(buffer.clone(), options.clone()) {
+                    Ok(_) => {}
+                    Err(e) => {
+                        eprintln!("Error formatting chunk: {e:#}");
+                    }
+                }
+                println!("{}", end_marker);
+                io::stdout().flush()?;
+            } else {
+                buffer.push_str(&line);
+            }
+        }
+    }
+
+    Ok(0)
+}
+
 fn format(
     files: Vec<PathBuf>,
     minimal_config_path: Option<String>,
@@ -516,6 +581,15 @@ fn determine_operation(matches: &Matches) -> Result<Operation, OperationError> {
         if minimal_config_path.is_some() {
             return Err(OperationError::MinimalPathWithStdin);
         }
+        if let (Some(start_marker), Some(end_marker)) = (
+            matches.opt_str("start-marker"),
+            matches.opt_str("end-marker"),
+        ) {
+            return Ok(Operation::StdinStream {
+                start_marker,
+                end_marker,
+            });
+        }
         let mut buffer = String::new();
         io::stdin().read_to_string(&mut buffer)?;
 
@@ -548,6 +622,8 @@ struct GetOptsOptions {
     unstable_features: bool,
     error_on_unformatted: Option<bool>,
     print_misformatted_file_names: bool,
+    start_marker: Option<String>,
+    end_marker: Option<String>,
 }
 
 impl GetOptsOptions {
@@ -646,6 +722,24 @@ impl GetOptsOptions {
             options.print_misformatted_file_names = true;
         }
 
+        options.start_marker = matches.opt_str("start-marker");
+        options.end_marker = matches.opt_str("end-marker");
+
+        if (options.start_marker.is_some() && options.end_marker.is_none())
+            || (options.start_marker.is_none() && options.end_marker.is_some())
+        {
+            return Err(format_err!(
+                "Both `--start-marker` and `--end-marker` must be \
+                 specified for streaming formatting"
+            ));
+        }
+
+        if options.start_marker.is_some() && !matches.free.is_empty() {
+            return Err(format_err!(
+                "Cannot specify files with streaming formatting"
+            ));
+        }
+
         if !rust_nightly {
             if let Some(ref emit_mode) = options.emit_mode {
                 if !STABLE_EMIT_MODES.contains(emit_mode) {

tests/rustfmt/main.rs

@@ -274,3 +274,39 @@ fn rustfmt_allow_not_a_dir_errors() {
     assert_eq!(stdout, "");
     assert_eq!(stderr, "");
 }
+
+#[test]
+fn rustfmt_streaming_formatting() {
+    let rustfmt_exe = env!("CARGO_BIN_EXE_rustfmt");
+    let temp_dir = tempfile::tempdir().unwrap();
+    let mut cmd = Command::new(rustfmt_exe);
+    cmd.current_dir(temp_dir.path());
+    cmd.args(&["--start-marker", "START-CHUNK", "--end-marker", "END-CHUNK"]);
+    cmd.stdin(std::process::Stdio::piped());
+    cmd.stdout(std::process::Stdio::piped());
+    cmd.stderr(std::process::Stdio::piped());
+
+    let mut child = cmd.spawn().expect("failed to spawn rustfmt");
+
+    // Write multiple chunks to stdin
+    {
+        use std::io::Write;
+        let mut stdin = child.stdin.take().expect("failed to open stdin");
+        stdin
+            .write_all(b"START-CHUNK\nfn foo(  x:i32,y : u32  ) {}\nEND-CHUNK\n")
+            .unwrap();
+        stdin
+            .write_all(b"START-CHUNK\nfn bar(  ) -> bool   {false}\nEND-CHUNK\n")
+            .unwrap();
+    }
+
+    let output = child.wait_with_output().expect("failed to wait on child");
+    let stdout = String::from_utf8(output.stdout).expect("utf-8");
+    let stderr = String::from_utf8(output.stderr).expect("utf-8");
+
+    assert!(stdout.contains("fn foo(x: i32, y: u32) {"));
+    assert!(stdout.contains("fn bar() -> bool {"));
+    assert!(stdout.contains("false"));
+    assert!(stdout.contains("END-CHUNK"));
+    assert!(stderr.is_empty(), "actual stderr:\n{}", stderr);
+}