Function that captures output and prints lines

code review
1

I have a function, capture_output(), that captures command output and prints lines in "real-time" in the terminal:

use std::{
    io::{BufRead, BufReader, Read},
    process::{Command, Output, Stdio},
    thread::JoinHandle,
};

pub struct CmdRunner {
    pub cmd: Command,
}

fn capture_output<T: Read + Send + 'static>(pipe: T) -> JoinHandle<String> {
    std::thread::spawn(move || {
        let mut capture = String::new();
        let reader = BufReader::new(pipe);

        for line in reader.lines() {
            let line = line.unwrap();

            capture.push_str(&line);
            print!("{}\r\n", line);
        }

        capture
    })
}

impl CmdRunner {
    pub fn new(cmd_str: &str) -> CmdRunner {
        let mut cmd = Command::new("script");

        cmd.arg("-qec").arg(cmd_str).arg("/dev/null");
        CmdRunner { cmd }
    }

    pub fn run_with_output(&mut self) -> Result<Output, std::io::Error> {
        self.cmd.stdin(Stdio::null());
        self.cmd.stdout(Stdio::piped());
        self.cmd.stderr(Stdio::piped());

        let mut child = self.cmd.spawn().expect("failed to spawn command");
        let stdout_pipe = child.stdout.take().unwrap();
        let stdout_thread = capture_output(stdout_pipe);
        let stderr_pipe = child.stderr.take().unwrap();
        let stderr_thread = capture_output(stderr_pipe);
        let stdout_output = stdout_thread.join().expect("failed to join stdout thread");
        let stderr_output = stderr_thread.join().expect("failed to join stderr thread");
        let exit_status = child.wait()?;

        Ok(Output {
            stdout: stdout_output.into(),
            stderr: stderr_output.into(),
            status: exit_status,
        })
    }
}

fn main () {
    let mut cmd_runner = CmdRunner::new("echo 'Hello, World!'");
    cmd_runner.run_with_output().unwrap();
}

I think that function shouldn't be doing two things? Maybe I should do the printing in the caller function? If so, how can I accomplish that?

2

This is traditionally called tee or split, and the more principled interface is to accept an additional writer instead of printing directly, so "printing" is not hard-coded into the function. This is especially valuable because as it stands, your current code intermingles standard output and standard error (it always prints everything to stdout).

Here's an alternative implementation.

2 Likes
3

Thanks a lot for the suggestion!

With your alternative, the function doesn't print echo in real-time in the terminal, though.

"Hello, World!" will only appear after the command has finished executing.

4

If you want it to print lines as they are received, you might need to add a writer.flush()? after the writer.write_all()? so it forces each line to be written to the terminal. Otherwise the data will be written to stdout, but your OS may not decide to pass those bytes through to whoever is connected to stdout (e.g. your terminal) until later on.

I believe print!() and println!() do this automatically when the content being printed contains a \n because stdout is internally wrapped in a std::io::LineWriter.

Alternatively, you could leave it up to the caller to make sure their writer does the flushing.

use std::io::LineWriter;

tee(BufReader::new(stdout_pipe), LineWriter::new(io::stdout()));
2 Likes
5

I tried adding writer.flush()? after writer.write_all()?, but the problem persists: Rust Playground

6

It looks like you need to manually write a newline because the reader.lines() stripped it away. When you do writer.write_all(line.bytes())? the line variable only contains the bytes for "Hello, World!".

This code works as expected:

writer.write(line.as_bytes())?;
writer.write(b"\r\n")?;
writer.flush()?;

(playground)

3 Likes
7

(but you should probably still use write_all() instead of write().)

1 Like
8

It works. Thanks a lot!

9

Why do you suggest using write_all() instead of write() in this case?

10

Because write() isn't guaranteed to write the whole buffer.

If you want to use write(), then you need to call it in a loop, always checking its return value, and if it's smaller than the length of the slice, then retry. That's what write_all() does, so you might as well use it instead of reinventing the retry logic yourself.

1 Like
11

Thanks for the explanation!

In the rest of my app, I'm using write!(). I've never had any issues. I think this is a different situation? For example:

#[derive(Debug)]
pub enum InputError {
    NotUTF8(Vec<u8>),
    EmptyString,
    IoError(io::Error),
}

impl From<io::Error> for InputError {
    fn from(error: io::Error) -> Self {
        InputError::IoError(error)
    }
}

impl fmt::Display for InputError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            InputError::NotUTF8(bytes) => write!(
                f,
                "Input contained non-UTF8 bytes: {:?}",
                bytes
                    .iter()
                    .map(|b| format!("0x{:X}", b))
                    .collect::<Vec<_>>()
            ),
            InputError::EmptyString => write!(f, "Input was empty"),
            InputError::IoError(e) => write!(f, "I/O Error: {}", e),
        }
    }
}
12

write! macro (not to be confused with write method) essentially uses write_all internally.

2 Likes
13

In more detail, write! uses the write_fmt method[1], which in turn uses the write_all method (as its documentation explains, too).

The basic idea is that write!(writer, format_string, args…) simply becomes writer.write_fmt(format_args!(format_string, args…)). It’s in fact one macro with very easy to understand source code :slight_smile:.

macro_rules! write {
    ($dst:expr, $($arg:tt)*) => {
        $dst.write_fmt($crate::format_args!($($arg)*))
    };
}

  1. or in other use cases it could use this one depending on which of the 2 Write traits is in scope and implemented for the type at hand ↩︎

2 Likes
15

write!() is a macro. I'm talking about the Write::write() method.

1 Like
16

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.