sbrd_gen/writer/
json.rs

1use crate::error::{IntoSbrdError, SchemaErrorKind, SchemaResult};
2use crate::generator::Randomizer;
3use crate::value::DataValue;
4use crate::writer::writer_base::{
5    GeneratedDisplayValues, SerializeWithGenerate, DUMMY_KEYS_NAME, DUMMY_VALUES_NAME,
6};
7use crate::writer::GeneratedValueWriterBase;
8use crate::Schema;
9use serde::ser::{SerializeMap, SerializeSeq};
10use serde::Serializer;
11use serde_json::ser::{CompactFormatter, PrettyFormatter};
12use serde_json::Serializer as JsonSerializer;
13use std::borrow::BorrowMut;
14use std::io;
15use std::marker::PhantomData;
16
17/// A writer that outputs as Json with no extra spaces for the key and the generated value
18pub struct CompactJsonWriter<W: io::Write> {
19    json_writer: JsonWriter<W, CompactFormatter>,
20}
21
22impl<W: io::Write> GeneratedValueWriterBase<W> for CompactJsonWriter<W> {
23    fn from_writer(writer: W) -> Self {
24        Self {
25            json_writer: JsonWriter::from_writer(writer),
26        }
27    }
28
29    fn into_inner(self) -> W {
30        self.json_writer.into_inner()
31    }
32
33    fn flush(&mut self) -> SchemaResult<()> {
34        self.json_writer.flush()
35    }
36
37    fn write_after_all_generated<R: Randomizer + ?Sized>(
38        &mut self,
39        use_key_header: bool,
40        schema: &Schema<R>,
41        rng: &mut R,
42        count: u64,
43    ) -> SchemaResult<()> {
44        self.json_writer
45            .write_after_all_generated(use_key_header, schema, rng, count)
46    }
47
48    fn write_with_generate<R: Randomizer + ?Sized>(
49        &mut self,
50        use_key_header: bool,
51        schema: &Schema<R>,
52        rng: &mut R,
53        count: u64,
54    ) -> SchemaResult<()> {
55        self.json_writer
56            .write_with_generate(use_key_header, schema, rng, count)
57    }
58}
59
60/// A writer that outputs as Json with human-readable for the key and the generated value
61pub struct PrettyJsonWriter<'a, W: io::Write> {
62    json_writer: JsonWriter<W, PrettyFormatter<'a>>,
63}
64
65impl<'a, W: io::Write> GeneratedValueWriterBase<W> for PrettyJsonWriter<'a, W> {
66    fn from_writer(writer: W) -> Self {
67        Self {
68            json_writer: JsonWriter::from_writer(writer),
69        }
70    }
71
72    fn into_inner(self) -> W {
73        self.json_writer.into_inner()
74    }
75
76    fn flush(&mut self) -> SchemaResult<()> {
77        self.json_writer.flush()
78    }
79
80    fn write_after_all_generated<R: Randomizer + ?Sized>(
81        &mut self,
82        use_key_header: bool,
83        schema: &Schema<R>,
84        rng: &mut R,
85        count: u64,
86    ) -> SchemaResult<()> {
87        self.json_writer
88            .write_after_all_generated(use_key_header, schema, rng, count)
89    }
90
91    fn write_with_generate<R: Randomizer + ?Sized>(
92        &mut self,
93        use_key_header: bool,
94        schema: &Schema<R>,
95        rng: &mut R,
96        count: u64,
97    ) -> SchemaResult<()> {
98        self.json_writer
99            .write_with_generate(use_key_header, schema, rng, count)
100    }
101}
102
103trait BuildJsonFormatter: serde_json::ser::Formatter {
104    fn build_formatter() -> Self;
105}
106
107impl BuildJsonFormatter for CompactFormatter {
108    fn build_formatter() -> Self {
109        CompactFormatter
110    }
111}
112
113impl<'a> BuildJsonFormatter for PrettyFormatter<'a> {
114    fn build_formatter() -> Self {
115        PrettyFormatter::new()
116    }
117}
118
119/// A writer that outputs as Json for the key and the generated value
120struct JsonWriter<W: io::Write, F: BuildJsonFormatter> {
121    writer: W,
122    formatter: PhantomData<F>,
123}
124
125impl<W: io::Write, F: BuildJsonFormatter> JsonWriter<W, F> {
126    fn from_writer(writer: W) -> Self {
127        Self {
128            writer,
129            formatter: PhantomData,
130        }
131    }
132
133    fn build_serializer(&mut self) -> JsonSerializer<&mut W, F> {
134        JsonSerializer::with_formatter(&mut self.writer, F::build_formatter())
135    }
136
137    fn flush(&mut self) -> SchemaResult<()> {
138        self.writer
139            .flush()
140            .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))
141    }
142
143    fn into_inner(self) -> W {
144        self.writer
145    }
146
147    fn write_after_all_generated<R: Randomizer + ?Sized>(
148        &mut self,
149        use_key_header: bool,
150        schema: &Schema<R>,
151        rng: &mut R,
152        count: u64,
153    ) -> SchemaResult<()> {
154        let mut values_list: Vec<GeneratedDisplayValues<String, DataValue>> = Vec::new();
155        for _ in 1..=count {
156            let generated = schema.generate(rng)?;
157            let values = generated.into_values_with_key()?;
158
159            let value_map = GeneratedDisplayValues::new(values);
160            values_list.push(value_map);
161        }
162
163        let mut serializer = self.build_serializer();
164        if use_key_header {
165            let mut map_state = serializer
166                .borrow_mut()
167                .serialize_map(Some(2))
168                .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
169            SerializeMap::serialize_entry(&mut map_state, DUMMY_KEYS_NAME, schema.get_keys())
170                .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
171            SerializeMap::serialize_entry(
172                &mut map_state,
173                DUMMY_VALUES_NAME,
174                values_list.as_slice(),
175            )
176            .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
177            SerializeMap::end(map_state)
178                .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
179        } else {
180            let mut seq_state = serializer
181                .borrow_mut()
182                .serialize_seq(Some(values_list.len()))
183                .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
184            for values in values_list.iter() {
185                SerializeSeq::serialize_element(&mut seq_state, values)
186                    .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
187            }
188            SerializeSeq::end(seq_state)
189                .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
190        }
191
192        self.flush()?;
193        Ok(())
194    }
195
196    fn write_with_generate<R: Randomizer + ?Sized>(
197        &mut self,
198        use_key_header: bool,
199        schema: &Schema<R>,
200        rng: &mut R,
201        count: u64,
202    ) -> SchemaResult<()> {
203        let mut serializer = self.build_serializer();
204        if use_key_header {
205            let mut map_state = serializer
206                .borrow_mut()
207                .serialize_map(Some(2))
208                .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
209            SerializeMap::serialize_entry(&mut map_state, DUMMY_KEYS_NAME, schema.get_keys())
210                .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
211            SerializeMap::serialize_entry(
212                &mut map_state,
213                DUMMY_VALUES_NAME,
214                &SerializeWithGenerate::new(schema, rng, &count),
215            )
216            .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
217            SerializeMap::end(map_state)
218                .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
219        } else {
220            let mut seq_state = serializer
221                .borrow_mut()
222                .serialize_seq(Some(count as usize))
223                .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
224            for _ in 1..=count {
225                let generated = schema.generate(rng)?;
226                let values = generated.into_values_with_key()?;
227
228                let value_map = GeneratedDisplayValues::new(values);
229                SerializeSeq::serialize_element(&mut seq_state, &value_map)
230                    .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
231            }
232            SerializeSeq::end(seq_state)
233                .map_err(|e| e.into_sbrd_gen_error(SchemaErrorKind::OutputError))?;
234        }
235
236        self.flush()?;
237        Ok(())
238    }
239}