jqrutils 0.1.0

Small Unix utilities: quiet, errfail, prefix
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

use std::env;
use std::io::{self, BufRead, BufReader, Read, Write};
use std::process::{Command, ExitCode, Stdio};
use std::thread;

fn prefix_lines(pipe: impl Read, prefix: &str, mut out: impl Write) {
    let reader = BufReader::new(pipe);
    for line in reader.split(b'\n') {
        match line {
            Ok(data) => {
                let _ = out.write_all(prefix.as_bytes());
                let _ = out.write_all(&data);
                let _ = out.write_all(b"\n");
                let _ = out.flush();
            }
            Err(_) => break,
        }
    }
}

fn passthrough(pipe: impl Read, mut out: impl Write) {
    let mut buf = [0u8; 8192];
    let mut pipe = pipe;
    loop {
        match pipe.read(&mut buf) {
            Ok(0) | Err(_) => break,
            Ok(n) => {
                let _ = out.write_all(&buf[..n]);
            }
        }
    }
}

fn main() -> ExitCode {
    let args: Vec<String> = env::args().skip(1).collect();

    let mut do_stdout = false;
    let mut do_stderr = false;
    let mut rest = &args[..];

    while let Some(arg) = rest.first() {
        match arg.as_str() {
            "--stdout" => { do_stdout = true; rest = &rest[1..]; }
            "--stderr" => { do_stderr = true; rest = &rest[1..]; }
            _ => break,
        }
    }

    if !do_stdout && !do_stderr {
        do_stdout = true;
        do_stderr = true;
    }

    if rest.len() < 2 {
        eprintln!("Usage: prefix [--stdout] [--stderr] <prefix> [--] <command> [args...]");
        return ExitCode::from(1);
    }

    let prefix = format!("{} ", &rest[0]);
    let cmd_args = if rest.get(1).map(|s| s.as_str()) == Some("--") {
        &rest[2..]
    } else {
        &rest[1..]
    };

    if cmd_args.is_empty() {
        eprintln!("Usage: prefix [--stdout] [--stderr] <prefix> [--] <command> [args...]");
        return ExitCode::from(1);
    }

    let mut child = match Command::new(&cmd_args[0])
        .args(&cmd_args[1..])
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
    {
        Ok(c) => c,
        Err(e) => {
            eprintln!("prefix: {}: {}", cmd_args[0], e);
            return ExitCode::from(127);
        }
    };

    let stdout = child.stdout.take().unwrap();
    let stderr = child.stderr.take().unwrap();

    let prefix_clone = prefix.clone();
    let t1 = thread::spawn(move || {
        if do_stdout {
            prefix_lines(stdout, &prefix_clone, io::stdout().lock());
        } else {
            passthrough(stdout, io::stdout().lock());
        }
    });

    let t2 = thread::spawn(move || {
        if do_stderr {
            prefix_lines(stderr, &prefix, io::stderr().lock());
        } else {
            passthrough(stderr, io::stderr().lock());
        }
    });

    let status = child.wait().expect("failed to wait on child");
    t1.join().unwrap();
    t2.join().unwrap();

    #[cfg(unix)]
    let code = {
        use std::os::unix::process::ExitStatusExt;
        status
            .code()
            .unwrap_or_else(|| status.signal().map(|s| 128 + s).unwrap_or(1))
    };
    #[cfg(not(unix))]
    let code = status.code().unwrap_or(1);

    ExitCode::from(code as u8)
}