Use atomic write for pip compile output (#6274)

## Summary

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

Closes https://github.com/astral-sh/uv/issues/6239.
This commit is contained in:
Charlie Marsh 2024-08-20 16:36:42 -04:00 committed by GitHub
parent f10ccc488e
commit 9892a4ab50
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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;
@ -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!(
@ -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)?;
@ -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(())
}
}