Skip to content

Commit

Permalink
Use atomic write for pip compile output (#6274)
Browse files Browse the repository at this point in the history
## Summary

This ensures that we don't stream output to the `--output-file`, since
other processes may rely on reading it.

Closes #6239.
  • Loading branch information
charliermarsh authored Aug 20, 2024
1 parent f10ccc4 commit 9892a4a
Showing 1 changed file with 28 additions and 18 deletions.
46 changes: 28 additions & 18 deletions crates/uv/src/commands/pip/compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use std::env;
use std::io::stdout;
use std::path::Path;

use anstream::{eprint, AutoStream, StripStream};
use anstream::{eprint, AutoStream};
use anyhow::{anyhow, Result};
use itertools::Itertools;
use owo_colors::OwoColorize;
Expand Down Expand Up @@ -377,7 +377,7 @@ pub(crate) async fn pip_compile(
};

// Write the resolved dependencies to the output channel.
let mut writer = OutputWriter::new(!quiet || output_file.is_none(), output_file)?;
let mut writer = OutputWriter::new(!quiet || output_file.is_none(), output_file);

if include_header {
writeln!(
Expand Down Expand Up @@ -504,6 +504,9 @@ pub(crate) async fn pip_compile(
}
}

// Commit the output to disk.
writer.commit().await?;

// Notify the user of any resolution diagnostics.
operations::diagnose_resolution(resolution.diagnostics(), printer)?;

Expand Down Expand Up @@ -600,42 +603,49 @@ fn cmd(
format!("uv {args}")
}

/// A multi-casting writer that writes to both the standard output and an output file, if present.
/// A multicasting writer that writes to both the standard output and an output file, if present.
#[allow(clippy::disallowed_types)]
struct OutputWriter {
struct OutputWriter<'a> {
stdout: Option<AutoStream<std::io::Stdout>>,
output_file: Option<StripStream<std::fs::File>>,
output_file: Option<&'a Path>,
buffer: Vec<u8>,
}

#[allow(clippy::disallowed_types)]
impl OutputWriter {
impl<'a> OutputWriter<'a> {
/// Create a new output writer.
fn new(include_stdout: bool, output_file: Option<&Path>) -> Result<Self> {
fn new(include_stdout: bool, output_file: Option<&'a Path>) -> Self {
let stdout = include_stdout.then(|| AutoStream::<std::io::Stdout>::auto(stdout()));
let output_file = output_file
.map(|output_file| -> Result<_, std::io::Error> {
let output_file = fs_err::File::create(output_file)?;
Ok(StripStream::new(output_file.into()))
})
.transpose()?;
Ok(Self {
Self {
stdout,
output_file,
})
buffer: Vec::new(),
}
}

/// Write the given arguments to both the standard output and the output file, if present.
/// Write the given arguments to both standard output and the output buffer, if present.
fn write_fmt(&mut self, args: std::fmt::Arguments<'_>) -> std::io::Result<()> {
use std::io::Write;

if let Some(output_file) = &mut self.output_file {
write!(output_file, "{args}")?;
// Write to the buffer.
if self.output_file.is_some() {
self.buffer.write_fmt(args)?;
}

// Write to standard output.
if let Some(stdout) = &mut self.stdout {
write!(stdout, "{args}")?;
}

Ok(())
}

/// Commit the buffer to the output file.
async fn commit(self) -> std::io::Result<()> {
if let Some(output_file) = self.output_file {
let stream = anstream::adapter::strip_bytes(&self.buffer).into_vec();
uv_fs::write_atomic(output_file, &stream).await?;
}
Ok(())
}
}

0 comments on commit 9892a4a

Please sign in to comment.