yt/
output.rs

1use std::borrow::Cow;
2use std::io::{self, IsTerminal as _, Write};
3use std::sync::Mutex;
4use std::sync::atomic::{AtomicBool, Ordering};
5
6use miette::IntoDiagnostic as _;
7use serde::Serialize;
8
9use crate::error::Result;
10
11/// How command results are printed on stdout.
12#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, usage_rs::ValueEnum)]
13pub enum Format {
14    /// Human-readable table
15    #[default]
16    Table,
17    /// One JSON document
18    Json,
19    /// One JSON object per line, written as results arrive
20    Ndjson,
21}
22
23pub struct Output {
24    format: Format,
25    tty: bool,
26    closed: AtomicBool,
27    out: Mutex<Box<dyn Write + Send>>,
28}
29
30impl Output {
31    pub fn stdout(format: Format) -> Self {
32        Self::new(format, io::stdout().is_terminal(), Box::new(io::stdout()))
33    }
34
35    pub fn new(format: Format, tty: bool, out: Box<dyn Write + Send>) -> Self {
36        Self {
37            format,
38            tty,
39            closed: AtomicBool::new(false),
40            out: Mutex::new(out),
41        }
42    }
43
44    pub fn format(&self) -> Format {
45        self.format
46    }
47
48    /// Whether stdout is a terminal (tables then get a header; summaries may add a total line).
49    pub fn is_terminal(&self) -> bool {
50        self.tty
51    }
52
53    /// True once the reader of stdout has gone away (`yt … | head`); callers stop producing output.
54    pub fn is_closed(&self) -> bool {
55        self.closed.load(Ordering::Relaxed)
56    }
57
58    /// One JSON document: pretty-printed, or on one line in ndjson mode.
59    pub fn json<T: Serialize + ?Sized>(&self, value: &T) -> Result<()> {
60        if self.format == Format::Ndjson {
61            return self.ndjson(value);
62        }
63        let text = serde_json::to_string_pretty(value).into_diagnostic()?;
64        self.emit(&escape_controls(&text))
65    }
66
67    pub fn ndjson<T: Serialize + ?Sized>(&self, value: &T) -> Result<()> {
68        let text = serde_json::to_string(value).into_diagnostic()?;
69        self.emit(&escape_controls(&text))
70    }
71
72    /// JSON text as the server sent it, indented; keys keep their order.
73    pub fn raw_json(&self, text: &str) -> Result<()> {
74        self.emit(&escape_controls(&indent_json(text)))
75    }
76
77    pub fn table(&self, table: &Table) -> Result<()> {
78        if table.is_empty() {
79            return Ok(());
80        }
81        self.write(&table.render(self.tty))
82    }
83
84    /// Writes `text` followed by a newline, with terminal control characters replaced.
85    pub fn write(&self, text: &str) -> Result<()> {
86        self.emit(&printable(text))
87    }
88
89    fn emit(&self, text: &str) -> Result<()> {
90        if self.is_closed() {
91            return Ok(());
92        }
93        let mut out = self
94            .out
95            .lock()
96            .unwrap_or_else(|poisoned| poisoned.into_inner());
97        let result = out
98            .write_all(text.as_bytes())
99            .and_then(|()| out.write_all(b"\n"))
100            .and_then(|()| out.flush());
101        match result {
102            Err(error) if error.kind() == io::ErrorKind::BrokenPipe => {
103                self.closed.store(true, Ordering::Relaxed);
104                Ok(())
105            }
106            other => other.into_diagnostic(),
107        }
108    }
109
110    /// Raw bytes (a download to stdout); a closed pipe stops output as elsewhere.
111    pub fn bytes(&self, data: &[u8]) -> Result<()> {
112        if self.is_closed() {
113            return Ok(());
114        }
115        let mut out = self
116            .out
117            .lock()
118            .unwrap_or_else(|poisoned| poisoned.into_inner());
119        match out.write_all(data).and_then(|()| out.flush()) {
120            Err(error) if error.kind() == io::ErrorKind::BrokenPipe => {
121                self.closed.store(true, Ordering::Relaxed);
122                Ok(())
123            }
124            other => other.into_diagnostic(),
125        }
126    }
127}
128
129pub struct Table {
130    header: Vec<String>,
131    rows: Vec<Vec<String>>,
132}
133
134impl Table {
135    pub fn new<S: Into<String>>(header: impl IntoIterator<Item = S>) -> Self {
136        Self {
137            header: header.into_iter().map(Into::into).collect(),
138            rows: Vec::new(),
139        }
140    }
141
142    pub fn row<S: Into<String>>(&mut self, cells: impl IntoIterator<Item = S>) -> &mut Self {
143        self.rows.push(cells.into_iter().map(Into::into).collect());
144        self
145    }
146
147    pub fn is_empty(&self) -> bool {
148        self.rows.is_empty()
149    }
150
151    fn render(&self, tty: bool) -> String {
152        let mut table = comfy_table::Table::new();
153        table.load_style(comfy_table::presets::NOTHING);
154        if tty {
155            table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
156            table.set_header(self.header.clone());
157        }
158        for row in &self.rows {
159            table.add_row(row.clone());
160        }
161        for column in table.column_iter_mut() {
162            column.set_padding((0, 2));
163        }
164        table
165            .to_string()
166            .lines()
167            .map(|line| line.trim_end().to_string())
168            .collect::<Vec<_>>()
169            .join("\n")
170    }
171}
172
173/// `text` safe to show in a terminal: control characters other than tab and newline become U+FFFD, and `\r\n` becomes `\n`.
174pub fn printable(text: &str) -> Cow<'_, str> {
175    if !text.chars().any(is_unsafe) {
176        return Cow::Borrowed(text);
177    }
178    let mut out = String::with_capacity(text.len());
179    let mut chars = text.chars().peekable();
180    while let Some(c) = chars.next() {
181        match c {
182            '\r' if chars.peek() == Some(&'\n') => {}
183            c if is_unsafe(c) => out.push('\u{fffd}'),
184            c => out.push(c),
185        }
186    }
187    Cow::Owned(out)
188}
189
190fn is_unsafe(c: char) -> bool {
191    c.is_control() && c != '\n' && c != '\t'
192}
193
194/// JSON text with raw control characters written as `\u` escapes, which decode to the same value.
195fn escape_controls(json: &str) -> Cow<'_, str> {
196    let raw = |c: char| c.is_control() && c != '\n';
197    if !json.chars().any(raw) {
198        return Cow::Borrowed(json);
199    }
200    let mut out = String::with_capacity(json.len() + 16);
201    for c in json.chars() {
202        if raw(c) {
203            out.push_str(&format!("\\u{:04x}", u32::from(c)));
204        } else {
205            out.push(c);
206        }
207    }
208    Cow::Owned(out)
209}
210
211/// Re-indents valid JSON text the way `serde_json::to_string_pretty` does, without reordering keys.
212fn indent_json(text: &str) -> String {
213    let mut out = String::with_capacity(text.len() * 2);
214    let mut depth = 0;
215    let mut in_string = false;
216    let mut escaped = false;
217    let mut chars = text.chars().peekable();
218    while let Some(c) = chars.next() {
219        if in_string {
220            out.push(c);
221            if escaped {
222                escaped = false;
223            } else if c == '\\' {
224                escaped = true;
225            } else if c == '"' {
226                in_string = false;
227            }
228            continue;
229        }
230        match c {
231            '"' => {
232                in_string = true;
233                out.push(c);
234            }
235            '{' | '[' => {
236                out.push(c);
237                while chars.next_if(|next| next.is_whitespace()).is_some() {}
238                if let Some(close) = chars.next_if(|next| matches!(next, '}' | ']')) {
239                    out.push(close);
240                } else {
241                    depth += 1;
242                    new_line(&mut out, depth);
243                }
244            }
245            '}' | ']' => {
246                depth -= 1;
247                new_line(&mut out, depth);
248                out.push(c);
249            }
250            ',' => {
251                out.push(c);
252                new_line(&mut out, depth);
253            }
254            ':' => out.push_str(": "),
255            c if c.is_whitespace() => {}
256            c => out.push(c),
257        }
258    }
259    out
260}
261
262fn new_line(out: &mut String, depth: usize) {
263    out.push('\n');
264    out.extend(std::iter::repeat_n("  ", depth));
265}
266
267#[cfg(test)]
268mod tests {
269    use std::io;
270    use std::sync::{Arc, Mutex};
271
272    use super::*;
273
274    #[derive(Clone, Default)]
275    struct Buffer(Arc<Mutex<Vec<u8>>>);
276
277    impl Write for Buffer {
278        fn write(&mut self, data: &[u8]) -> io::Result<usize> {
279            self.0.lock().unwrap().extend_from_slice(data);
280            Ok(data.len())
281        }
282        fn flush(&mut self) -> io::Result<()> {
283            Ok(())
284        }
285    }
286
287    impl Buffer {
288        fn text(&self) -> String {
289            String::from_utf8(self.0.lock().unwrap().clone()).unwrap()
290        }
291    }
292
293    struct ClosedPipe;
294
295    impl Write for ClosedPipe {
296        fn write(&mut self, _: &[u8]) -> io::Result<usize> {
297            Err(io::Error::from(io::ErrorKind::BrokenPipe))
298        }
299        fn flush(&mut self) -> io::Result<()> {
300            Ok(())
301        }
302    }
303
304    fn sample() -> Table {
305        let mut table = Table::new(["ID", "SUMMARY"]);
306        table.row(["WEB-1", "Login fails"]).row(["WEB-12", "Typo"]);
307        table
308    }
309
310    #[test]
311    fn json_is_pretty_printed() {
312        let buffer = Buffer::default();
313        let output = Output::new(Format::Json, false, Box::new(buffer.clone()));
314        output.json(&serde_json::json!({"id": "WEB-1"})).unwrap();
315        assert_eq!(buffer.text(), "{\n  \"id\": \"WEB-1\"\n}\n");
316    }
317
318    #[test]
319    fn ndjson_is_one_line_per_item() {
320        let buffer = Buffer::default();
321        let output = Output::new(Format::Ndjson, false, Box::new(buffer.clone()));
322        output.ndjson(&serde_json::json!({"id": "WEB-1"})).unwrap();
323        output.ndjson(&serde_json::json!({"id": "WEB-2"})).unwrap();
324        assert_eq!(buffer.text(), "{\"id\":\"WEB-1\"}\n{\"id\":\"WEB-2\"}\n");
325    }
326
327    #[test]
328    fn a_single_result_is_one_line_in_ndjson_mode() {
329        let buffer = Buffer::default();
330        let output = Output::new(Format::Ndjson, false, Box::new(buffer.clone()));
331        output
332            .json(&serde_json::json!({"operations": [{"kind": "command"}]}))
333            .unwrap();
334        assert_eq!(buffer.text(), "{\"operations\":[{\"kind\":\"command\"}]}\n");
335    }
336
337    #[test]
338    fn raw_json_is_indented_like_serde_without_reordering() {
339        let text = r#"{"z":1,"a":[],"o":{ },"s":"q\"{a, b: c}\\","n":[{"k":null},2.5e3,true]}"#;
340        let value: serde_json::Value = serde_json::from_str(text).unwrap();
341        let indented = indent_json(text);
342        assert_eq!(
343            serde_json::from_str::<serde_json::Value>(&indented).unwrap(),
344            value
345        );
346        assert_eq!(
347            indented,
348            "{\n  \"z\": 1,\n  \"a\": [],\n  \"o\": {},\n  \"s\": \"q\\\"{a, b: c}\\\\\",\n  \"n\": [\n    {\n      \"k\": null\n    },\n    2.5e3,\n    true\n  ]\n}"
349        );
350    }
351
352    #[test]
353    fn piped_tables_have_no_header() {
354        let buffer = Buffer::default();
355        Output::new(Format::Table, false, Box::new(buffer.clone()))
356            .table(&sample())
357            .unwrap();
358        assert_eq!(buffer.text(), "WEB-1   Login fails\nWEB-12  Typo\n");
359    }
360
361    #[test]
362    fn terminal_tables_have_a_header() {
363        let buffer = Buffer::default();
364        Output::new(Format::Table, true, Box::new(buffer.clone()))
365            .table(&sample())
366            .unwrap();
367        assert!(
368            buffer.text().starts_with("ID      SUMMARY"),
369            "{}",
370            buffer.text()
371        );
372    }
373
374    #[test]
375    fn printable_replaces_terminal_control_characters() {
376        assert_eq!(
377            printable("a\u{1b}[2Jb\u{7}c\u{9b}d\u{7f}e"),
378            "a\u{fffd}[2Jb\u{fffd}c\u{fffd}d\u{fffd}e"
379        );
380        assert_eq!(
381            printable("one\r\ntwo\tthree\rfour"),
382            "one\ntwo\tthree\u{fffd}four"
383        );
384        assert!(matches!(printable("plain ā\n"), Cow::Borrowed(_)));
385    }
386
387    #[test]
388    fn text_output_is_printable() {
389        let buffer = Buffer::default();
390        let output = Output::new(Format::Table, true, Box::new(buffer.clone()));
391        output.write("a\u{1b}]0;title\u{7}b").unwrap();
392        assert_eq!(buffer.text(), "a\u{fffd}]0;title\u{fffd}b\n");
393    }
394
395    #[test]
396    fn json_output_escapes_control_characters() {
397        let buffer = Buffer::default();
398        let output = Output::new(Format::Ndjson, false, Box::new(buffer.clone()));
399        output
400            .json(&serde_json::json!({"s": "a\u{9b}b\u{7f}"}))
401            .unwrap();
402        output.raw_json("{\"s\":\"c\u{1b}d\"}").unwrap();
403        assert_eq!(
404            buffer.text(),
405            "{\"s\":\"a\\u009bb\\u007f\"}\n{\n  \"s\": \"c\\u001bd\"\n}\n"
406        );
407    }
408
409    #[test]
410    fn a_closed_pipe_is_not_an_error_and_stops_output() {
411        let output = Output::new(Format::Ndjson, false, Box::new(ClosedPipe));
412        output.write("first").unwrap();
413        assert!(output.is_closed());
414        output.write("second").unwrap();
415    }
416}
yt/output.rs

yt/
output.rs
