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 + Send + Sync + 'static,
23        E: Write + Send + Sync + '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 + Send + Sync) {
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            let mut buf = serde_json::to_vec(output)?;
148            buf.push(b'\n');
149            out.write_all(&buf)?;
150            out.flush().map_err(Into::into)
151        }
152        Format::JsonPretty => {
153            let buf = serde_json::to_vec_pretty(output)?;
154            out.write_all(&buf)?;
155            out.flush().map_err(Into::into)
156        }
157        Format::Raw => Ok(()),
158    }
159}
160
161impl<T> ToolIO for JsonWriter<T>
162where
163    T: Serialize,
164{
165    type OutputItem = T;
166
167    fn supports_structured_output() -> bool {
168        true
169    }
170
171    fn is_machine(&self) -> bool {
172        match self.format {
173            Some(Format::Json) | Some(Format::JsonPretty) => true,
174            _ => false,
175        }
176    }
177
178    fn stderr(&mut self) -> &mut (dyn Write + Send + Sync) {
179        self.stderr()
180    }
181
182    fn item(&mut self, value: &T) -> Result<()>
183    where
184        T: Display,
185    {
186        if self.is_machine() { self.machine(value) } else { self.line(value) }
187    }
188}
189
190impl<T> Write for JsonWriter<T>
191where
192    T: Serialize,
193{
194    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
195        if self.is_machine() { Ok(buf.len()) } else { self.simple_writer.write(buf) }
196    }
197
198    fn flush(&mut self) -> std::io::Result<()> {
199        if self.is_machine() { Ok(()) } else { self.simple_writer.flush() }
200    }
201}
202#[cfg(test)]
203mod test {
204    use super::*;
205
206    #[test]
207    fn test_not_machine_is_ok() {
208        let test_buffers = TestBuffers::default();
209        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
210        let res = writer.machine(&"ehllo");
211        assert!(res.is_ok());
212    }
213
214    #[test]
215    fn test_machine_valid_json_is_ok() {
216        let test_buffers = TestBuffers::default();
217        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
218        let res = writer.machine(&"ehllo");
219        assert!(res.is_ok());
220    }
221
222    #[test]
223    fn writer_implements_write() {
224        let test_buffers = TestBuffers::default();
225        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
226        writer.write_all(b"foobar").unwrap();
227
228        let (stdout, stderr) = test_buffers.into_strings();
229        assert_eq!(stdout, "foobar");
230        assert_eq!(stderr, "");
231    }
232
233    #[test]
234    fn writer_implements_write_ignored_on_machine() {
235        let test_buffers = TestBuffers::default();
236        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
237        writer.write_all(b"foobar").unwrap();
238
239        let (stdout, stderr) = test_buffers.into_strings();
240        assert_eq!(stdout, "");
241        assert_eq!(stderr, "");
242    }
243
244    #[test]
245    fn test_item_for_test_as_machine() {
246        let test_buffers = TestBuffers::default();
247        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
248        writer.item(&"hello").unwrap();
249        writer.machine_or(&"hello again", "but what if").unwrap();
250        writer.machine_or_else(&"hello forever", || "but what if else").unwrap();
251
252        assert_eq!(
253            test_buffers.into_stdout_str(),
254            "\"hello\"\n\"hello again\"\n\"hello forever\"\n"
255        );
256    }
257
258    #[test]
259    fn test_item_for_test_as_not_machine() {
260        let test_buffers = TestBuffers::default();
261        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
262        writer.item(&"hello").unwrap();
263        writer.machine_or(&"hello again", "but what if").unwrap();
264        writer.machine_or_else(&"hello forever", || "but what if else").unwrap();
265
266        assert_eq!(test_buffers.into_stdout_str(), "hello\nbut what if\nbut what if else\n");
267    }
268
269    #[test]
270    fn test_machine_for_test() {
271        let test_buffers = TestBuffers::default();
272        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
273        writer.machine(&"hello").unwrap();
274
275        assert_eq!(test_buffers.into_stdout_str(), "\"hello\"\n");
276    }
277
278    #[test]
279    fn test_not_machine_for_test_is_empty() {
280        let test_buffers = TestBuffers::default();
281        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
282        writer.machine(&"hello").unwrap();
283
284        assert!(test_buffers.into_stdout_str().is_empty());
285    }
286
287    #[test]
288    fn test_machine_makes_is_machine_true() {
289        let test_buffers = TestBuffers::default();
290        let writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
291        assert!(writer.is_machine());
292    }
293
294    #[test]
295    fn test_not_machine_makes_is_machine_false() {
296        let test_buffers = TestBuffers::default();
297        let writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
298        assert!(!writer.is_machine());
299    }
300
301    #[test]
302    fn line_writer_for_machine_is_ok() {
303        let test_buffers = TestBuffers::default();
304        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
305        writer.line("hello").unwrap();
306
307        let (stdout, stderr) = test_buffers.into_strings();
308        assert_eq!(stdout, "");
309        assert_eq!(stderr, "");
310    }
311
312    #[test]
313    fn writer_write_for_machine_is_ok() {
314        let test_buffers = TestBuffers::default();
315        let mut writer: JsonWriter<&str> = JsonWriter::new_test(Some(Format::Json), &test_buffers);
316        writer.print("foobar").unwrap();
317
318        let (stdout, stderr) = test_buffers.into_strings();
319        assert_eq!(stdout, "");
320        assert_eq!(stderr, "");
321    }
322
323    #[test]
324    fn writer_print_output_has_no_newline() {
325        let test_buffers = TestBuffers::default();
326        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
327        writer.print("foobar").unwrap();
328
329        let (stdout, stderr) = test_buffers.into_strings();
330        assert_eq!(stdout, "foobar");
331        assert_eq!(stderr, "");
332    }
333
334    #[test]
335    fn writing_errors_goes_to_the_right_stream() {
336        let test_buffers = TestBuffers::default();
337        let mut writer: JsonWriter<&str> = JsonWriter::new_test(None, &test_buffers);
338        writeln!(writer.stderr(), "hello").unwrap();
339
340        let (stdout, stderr) = test_buffers.into_strings();
341        assert_eq!(stdout, "");
342        assert_eq!(stderr, "hello\n");
343    }
344
345    #[test]
346    fn test_machine_writes_pretty_json() {
347        let test_buffers = TestBuffers::default();
348        let mut writer: JsonWriter<serde_json::Value> =
349            JsonWriter::new_test(Some(Format::JsonPretty), &test_buffers);
350        let test_input = serde_json::json!({
351            "object1": {
352                "line1": "hello",
353                "line2": "foobar"
354            }
355        });
356        writer.machine(&test_input).unwrap();
357
358        assert_eq!(
359            test_buffers.into_stdout_str(),
360            r#"{
361  "object1": {
362    "line1": "hello",
363    "line2": "foobar"
364  }
365}"#
366        );
367    }
368
369    #[test]
370    fn test_concurrent_format_output_does_not_interleave_json_lines() {
371        use std::sync::{Arc, Barrier, Mutex};
372
373        #[derive(Clone)]
374        struct YieldingBuffer(Arc<Mutex<Vec<u8>>>);
375
376        impl Write for YieldingBuffer {
377            fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
378                self.0.lock().unwrap().extend_from_slice(buf);
379                std::thread::yield_now();
380                Ok(buf.len())
381            }
382
383            fn flush(&mut self) -> std::io::Result<()> {
384                Ok(())
385            }
386        }
387
388        let shared = YieldingBuffer(Arc::new(Mutex::new(Vec::new())));
389        let num_threads = 8;
390        let items_per_thread = 50;
391        let barrier = Arc::new(Barrier::new(num_threads));
392        let mut handles = Vec::with_capacity(num_threads);
393
394        for tid in 0..num_threads {
395            let mut out = shared.clone();
396            let barrier = Arc::clone(&barrier);
397            handles.push(std::thread::spawn(move || {
398                barrier.wait();
399                for seq in 0..items_per_thread {
400                    let payload = serde_json::json!({
401                        "thread": tid,
402                        "seq": seq,
403                        "nested": {
404                            "alpha": "value_a",
405                            "beta": "value_b",
406                            "gamma": [1, 2, 3, 4, 5],
407                        }
408                    });
409                    format_output(Format::Json, &mut out, &payload).unwrap();
410                }
411            }));
412        }
413
414        for h in handles {
415            h.join().unwrap();
416        }
417
418        let raw = String::from_utf8(shared.0.lock().unwrap().clone()).unwrap();
419        let lines: Vec<&str> = raw.lines().collect();
420        assert_eq!(lines.len(), num_threads * items_per_thread);
421        for line in lines {
422            let parsed: serde_json::Value =
423                serde_json::from_str(line).expect("each line must be a complete valid JSON object");
424            assert!(parsed.get("thread").is_some());
425            assert!(parsed.get("seq").is_some());
426            assert_eq!(parsed["nested"]["alpha"], "value_a");
427        }
428    }
429}
writer/json_writer.rs

writer/
json_writer.rs
