Skip to main content

writer/
json_writer.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use crate::{Format, Result, TestBuffers, ToolIO, Writer};
6use serde::Serialize;
7use std::fmt::Display;
8use std::io::Write;
9
10/// Type-safe machine output implementation of [`crate::ToolIO`]
11pub struct JsonWriter<T> {
12    format: Option<Format>,
13    simple_writer: Writer,
14    _p: std::marker::PhantomData<fn(T)>,
15}
16
17impl<T> JsonWriter<T> {
18    /// Create a new writer that doesn't support machine output at all, with the
19    /// given streams underlying it.
20    pub fn new_buffers<'a, O, E>(format: Option<Format>, stdout: O, stderr: E) -> Self
21    where
22        O: Write + 'static,
23        E: Write + 'static,
24    {
25        let simple_writer = Writer::new_buffers(stdout, stderr);
26        let _p = Default::default();
27        Self { format, simple_writer, _p }
28    }
29
30    /// Create a new Writer with the specified format.
31    ///
32    /// Passing None for format implies no output via the machine function.
33    pub fn new(format: Option<Format>) -> Self {
34        let simple_writer = Writer::new();
35        let _p = Default::default();
36        Self { format, simple_writer, _p }
37    }
38
39    /// Returns a writer backed by string buffers that can be extracted after
40    /// the writer is done with
41    pub fn new_test(format: Option<Format>, test_buffers: &TestBuffers) -> Self {
42        Self::new_buffers(format, test_buffers.stdout.clone(), test_buffers.stderr.clone())
43    }
44
45    /// Allow using this writer as a SimpleWriter.
46    /// This is useful for calling other tools that do not use JsonWriter.
47    pub fn simple_writer(self) -> Writer {
48        self.simple_writer
49    }
50
51    pub fn stderr(&mut self) -> &mut dyn Write {
52        self.simple_writer.stderr()
53    }
54
55    // Always returns Some(Json|JsonPretty) or None
56    fn get_machine_format(&self) -> Option<Format> {
57        match self.format {
58            None | Some(Format::Raw) => None,
59            j => j,
60        }
61    }
62}
63
64impl<T> JsonWriter<T>
65where
66    T: Serialize,
67{
68    /// Write the items from the iterable object to standard output.
69    ///
70    /// This is a no-op if `is_machine` returns false.
71    pub fn machine_many<I: IntoIterator<Item = T>>(&mut self, output: I) -> Result<()> {
72        if self.is_machine() {
73            for output in output {
74                self.machine(&output)?;
75            }
76        }
77        Ok(())
78    }
79
80    /// Write the item to standard output.
81    ///
82    /// This is a no-op if `is_machine` returns false.
83    pub fn machine(&mut self, output: &T) -> Result<()> {
84        if let Some(format) = self.get_machine_format() {
85            format_output(format, &mut self.simple_writer, output)
86        } else {
87            Ok(())
88        }
89    }
90
91    /// If this object is outputting machine output, print the item's machine
92    /// representation to stdout. Otherwise, print the display item given.
93    pub fn machine_or<D: Display>(&mut self, value: &T, or: D) -> Result<()> {
94        if let Some(format) = self.get_machine_format() {
95            format_output(format, &mut self.simple_writer, value)?
96        } else {
97            writeln!(self, "{or}")?
98        }
99        Ok(())
100    }
101
102    /// If this object is outputting machine output, print the item's machine
103    /// representation to stdout. Otherwise, `write!` the display item given.
104    pub fn machine_or_write<D: Display>(&mut self, value: &T, or: D) -> Result<()> {
105        if let Some(format) = self.get_machine_format() {
106            format_output(format, &mut self.simple_writer, value)?
107        } else {
108            write!(self, "{or}")?
109        }
110        Ok(())
111    }
112
113    /// If this object is outputting machine output, prints the item's machine
114    /// representation to stdout. Otherwise, call the closure with the object
115    /// and print the result.
116    pub fn machine_or_else<F, R>(&mut self, value: &T, f: F) -> Result<()>
117    where
118        F: FnOnce() -> R,
119        R: Display,
120    {
121        match self.format {
122            Some(format) => format_output(format, &mut self.simple_writer, value)?,
123            None => writeln!(self, "{}", f())?,
124        }
125        Ok(())
126    }
127
128    pub fn format(&self) -> Option<Format> {
129        self.format
130    }
131
132    pub fn formatted<J: Serialize>(&mut self, output: &J) -> Result<()> {
133        if let Some(format) = self.get_machine_format() {
134            format_output(format, &mut self.simple_writer, output)
135        } else {
136            Ok(())
137        }
138    }
139}
140
141pub fn format_output<W: Write, T>(format: Format, mut out: W, output: &T) -> Result<()>
142where
143    T: Serialize,
144{
145    match format {
146        Format::Json => {
147            serde_json::to_writer(&mut out, output)?;
148            writeln!(out)?;
149            out.flush()?;
150        }
151        Format::JsonPretty => {
152            serde_json::to_writer_pretty(&mut out, output)?;
153        }
154        Format::Raw => {}
155    }
156    Ok(())
157}
158
159impl<T> ToolIO for JsonWriter<T>
160where
161    T: Serialize,
162{
163    type OutputItem = T;
164
165    fn supports_structured_output() -> bool {
166        true
167    }
168
169    fn is_machine(&self) -> bool {
170        match self.format {
171            Some(Format::Json) | Some(Format::JsonPretty) => true,
172            _ => false,
173        }
174    }
175
176    fn stderr(&mut self) -> &mut dyn Write {
177        self.stderr()
178    }
179
180    fn item(&mut self, value: &T) -> Result<()>
181    where
182        T: Display,
183    {
184        if self.is_machine() { self.machine(value) } else { self.line(value) }
185    }
186}
187
188impl<T> Write for JsonWriter<T>
189where
190    T: Serialize,
191{
192    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
193        if self.is_machine() { Ok(buf.len()) } else { self.simple_writer.write(buf) }
194    }
195
196    fn flush(&mut self) -> std::io::Result<()> {
197        if self.is_machine() { Ok(()) } else { self.simple_writer.flush() }
198    }
199}
200#[cfg(test)]
201mod test {
202    use super::*;
203
204    #[test]
205    fn test_not_machine_is_ok() {
206        let test_buffers = TestBuffers::default();
207        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
208        let res = writer.machine(&"ehllo");
209        assert!(res.is_ok());
210    }
211
212    #[test]
213    fn test_machine_valid_json_is_ok() {
214        let test_buffers = TestBuffers::default();
215        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
216        let res = writer.machine(&"ehllo");
217        assert!(res.is_ok());
218    }
219
220    #[test]
221    fn writer_implements_write() {
222        let test_buffers = TestBuffers::default();
223        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
224        writer.write_all(b"foobar").unwrap();
225
226        let (stdout, stderr) = test_buffers.into_strings();
227        assert_eq!(stdout, "foobar");
228        assert_eq!(stderr, "");
229    }
230
231    #[test]
232    fn writer_implements_write_ignored_on_machine() {
233        let test_buffers = TestBuffers::default();
234        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
235        writer.write_all(b"foobar").unwrap();
236
237        let (stdout, stderr) = test_buffers.into_strings();
238        assert_eq!(stdout, "");
239        assert_eq!(stderr, "");
240    }
241
242    #[test]
243    fn test_item_for_test_as_machine() {
244        let test_buffers = TestBuffers::default();
245        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
246        writer.item(&"hello").unwrap();
247        writer.machine_or(&"hello again", "but what if").unwrap();
248        writer.machine_or_else(&"hello forever", || "but what if else").unwrap();
249
250        assert_eq!(
251            test_buffers.into_stdout_str(),
252            "\"hello\"\n\"hello again\"\n\"hello forever\"\n"
253        );
254    }
255
256    #[test]
257    fn test_item_for_test_as_not_machine() {
258        let test_buffers = TestBuffers::default();
259        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
260        writer.item(&"hello").unwrap();
261        writer.machine_or(&"hello again", "but what if").unwrap();
262        writer.machine_or_else(&"hello forever", || "but what if else").unwrap();
263
264        assert_eq!(test_buffers.into_stdout_str(), "hello\nbut what if\nbut what if else\n");
265    }
266
267    #[test]
268    fn test_machine_for_test() {
269        let test_buffers = TestBuffers::default();
270        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
271        writer.machine(&"hello").unwrap();
272
273        assert_eq!(test_buffers.into_stdout_str(), "\"hello\"\n");
274    }
275
276    #[test]
277    fn test_not_machine_for_test_is_empty() {
278        let test_buffers = TestBuffers::default();
279        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
280        writer.machine(&"hello").unwrap();
281
282        assert!(test_buffers.into_stdout_str().is_empty());
283    }
284
285    #[test]
286    fn test_machine_makes_is_machine_true() {
287        let test_buffers = TestBuffers::default();
288        let writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
289        assert!(writer.is_machine());
290    }
291
292    #[test]
293    fn test_not_machine_makes_is_machine_false() {
294        let test_buffers = TestBuffers::default();
295        let writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
296        assert!(!writer.is_machine());
297    }
298
299    #[test]
300    fn line_writer_for_machine_is_ok() {
301        let test_buffers = TestBuffers::default();
302        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
303        writer.line("hello").unwrap();
304
305        let (stdout, stderr) = test_buffers.into_strings();
306        assert_eq!(stdout, "");
307        assert_eq!(stderr, "");
308    }
309
310    #[test]
311    fn writer_write_for_machine_is_ok() {
312        let test_buffers = TestBuffers::default();
313        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
314        writer.print("foobar").unwrap();
315
316        let (stdout, stderr) = test_buffers.into_strings();
317        assert_eq!(stdout, "");
318        assert_eq!(stderr, "");
319    }
320
321    #[test]
322    fn writer_print_output_has_no_newline() {
323        let test_buffers = TestBuffers::default();
324        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
325        writer.print("foobar").unwrap();
326
327        let (stdout, stderr) = test_buffers.into_strings();
328        assert_eq!(stdout, "foobar");
329        assert_eq!(stderr, "");
330    }
331
332    #[test]
333    fn writing_errors_goes_to_the_right_stream() {
334        let test_buffers = TestBuffers::default();
335        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
336        writeln!(writer.stderr(), "hello").unwrap();
337
338        let (stdout, stderr) = test_buffers.into_strings();
339        assert_eq!(stdout, "");
340        assert_eq!(stderr, "hello\n");
341    }
342
343    #[test]
344    fn test_machine_writes_pretty_json() {
345        let test_buffers = TestBuffers::default();
346        let mut writer: JsonWriter<serde_json::Value> =
347            JsonWriter::new_test(Some(Format::JsonPretty), &test_buffers);
348        let test_input = serde_json::json!({
349            "object1": {
350                "line1": "hello",
351                "line2": "foobar"
352            }
353        });
354        writer.machine(&test_input).unwrap();
355
356        assert_eq!(
357            test_buffers.into_stdout_str(),
358            r#"{
359  "object1": {
360    "line1": "hello",
361    "line2": "foobar"
362  }
363}"#
364        );
365    }
366}