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
17pub 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
60pub 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
119struct 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}