Skip to content

Commit

Permalink
Use join instead
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh committed Sep 2, 2024
1 parent 7294b8d commit d06414d
Showing 1 changed file with 29 additions and 59 deletions.
88 changes: 29 additions & 59 deletions crates/uv-build/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1016,6 +1016,22 @@ impl PythonRunner {
environment_variables: &FxHashMap<OsString, OsString>,
modified_path: &OsString,
) -> Result<PythonRunnerOutput, Error> {
/// Read lines from a reader and store them in a buffer.
async fn read_from(
mut reader: tokio::io::Lines<tokio::io::BufReader<impl tokio::io::AsyncRead + Unpin>>,
buffer: &mut Vec<String>,
) -> io::Result<()> {
loop {
match reader.next_line().await? {
Some(line) => {
debug!("{line}");
buffer.push(line);
}
None => return Ok(()),
}
}
}

let _permit = self.control.acquire().await.unwrap();

let mut child = Command::new(venv.python_executable())
Expand All @@ -1031,71 +1047,25 @@ impl PythonRunner {
.map_err(|err| Error::CommandFailed(venv.python_executable().to_path_buf(), err))?;

// Create buffers to capture `stdout` and `stderr`.
let mut stdout_buf = Vec::new();
let mut stderr_buf = Vec::new();
let mut stdout_buf = Vec::with_capacity(1024);
let mut stderr_buf = Vec::with_capacity(1024);

// Create separate readers for `stdout` and `stderr`.
let stdout_reader = tokio::io::BufReader::new(child.stdout.take().unwrap()).lines();
let stderr_reader = tokio::io::BufReader::new(child.stderr.take().unwrap()).lines();

// Asynchronously read from the in-memory pipes.
tokio::pin!(stdout_reader, stderr_reader);

loop {
tokio::select! {
line = stdout_reader.next_line() => {
match line {
Ok(Some(line)) => {
debug!("{line}");
stdout_buf.push(line);
},
Ok(None) => break,
Err(err) => return Err(Error::CommandFailed(venv.python_executable().to_path_buf(), err)),
}
},
line = stderr_reader.next_line() => {
match line {
Ok(Some(line)) => {
debug!("{line}");
stderr_buf.push(line);
},
Ok(None) => break,
Err(err) => return Err(Error::CommandFailed(venv.python_executable().to_path_buf(), err)),
}
},
}
}

// Continue reading until both streams are closed.
loop {
match stdout_reader.next_line().await {
Ok(Some(line)) => {
debug!("{line}");
stdout_buf.push(line);
}
Ok(None) => break,
Err(err) => {
return Err(Error::CommandFailed(
venv.python_executable().to_path_buf(),
err,
))
}
}
}

loop {
match stderr_reader.next_line().await {
Ok(Some(line)) => {
debug!("{line}");
stderr_buf.push(line);
}
Ok(None) => break,
Err(err) => {
return Err(Error::CommandFailed(
venv.python_executable().to_path_buf(),
err,
))
}
let result = tokio::join!(
read_from(stdout_reader, &mut stdout_buf),
read_from(stderr_reader, &mut stderr_buf),
);
match result {
(Ok(()), Ok(())) => {}
(Err(err), _) | (_, Err(err)) => {
return Err(Error::CommandFailed(
venv.python_executable().to_path_buf(),
err,
))
}
}

Expand Down

0 comments on commit d06414d

Please sign in to comment.