1use crate::llm::human_formatter::{HumanFormatConfig, HumanFormatter};
20use crate::llm::human_parser::HumanParser;
21use crate::llm::types::DxDocument;
22use thiserror::Error;
23
24#[derive(Debug, Error)]
26pub enum PrettyPrintError {
27 #[error("Output validation failed: {msg}")]
29 ValidationFailed {
30 msg: String,
32 },
33
34 #[error("Round-trip consistency failed: {msg}")]
36 RoundTripFailed {
37 msg: String,
39 },
40}
41
42#[derive(Debug, Clone)]
44pub struct PrettyPrinterConfig {
45 pub formatter_config: HumanFormatConfig,
47 pub validate_output: bool,
49 pub check_round_trip: bool,
51}
52
53impl Default for PrettyPrinterConfig {
54 fn default() -> Self {
55 Self {
56 formatter_config: HumanFormatConfig::default(),
57 validate_output: true,
58 check_round_trip: true,
59 }
60 }
61}
62
63impl PrettyPrinterConfig {
64 pub fn new() -> Self {
66 Self::default()
67 }
68
69 pub fn with_key_padding(mut self, padding: usize) -> Self {
71 self.formatter_config.key_padding = padding;
72 self
73 }
74
75 pub fn with_validation(mut self, validate: bool) -> Self {
77 self.validate_output = validate;
78 self
79 }
80
81 pub fn with_round_trip_check(mut self, check: bool) -> Self {
83 self.check_round_trip = check;
84 self
85 }
86}
87
88pub struct PrettyPrinter {
94 config: PrettyPrinterConfig,
95 formatter: HumanFormatter,
96 parser: HumanParser,
97}
98
99impl PrettyPrinter {
100 pub fn new() -> Self {
102 let config = PrettyPrinterConfig::default();
103 Self {
104 formatter: HumanFormatter::with_config(config.formatter_config.clone()),
105 parser: HumanParser::new(),
106 config,
107 }
108 }
109
110 pub fn with_config(config: PrettyPrinterConfig) -> Self {
112 Self {
113 formatter: HumanFormatter::with_config(config.formatter_config.clone()),
114 parser: HumanParser::new(),
115 config,
116 }
117 }
118
119 pub fn format(&self, doc: &DxDocument) -> Result<String, PrettyPrintError> {
125 let output = self.formatter.format(doc);
127
128 if self.config.validate_output {
130 self.validate_output(&output, doc)?;
131 }
132
133 Ok(output)
134 }
135
136 pub fn format_unchecked(&self, doc: &DxDocument) -> String {
138 self.formatter.format(doc)
139 }
140
141 fn validate_output(&self, output: &str, original: &DxDocument) -> Result<(), PrettyPrintError> {
143 let parsed = self
145 .parser
146 .parse(output)
147 .map_err(|e| PrettyPrintError::ValidationFailed {
148 msg: format!("Failed to parse formatted output: {}", e),
149 })?;
150
151 if self.config.check_round_trip {
153 self.check_round_trip(original, &parsed)?;
154 }
155
156 Ok(())
157 }
158
159 fn check_round_trip(
161 &self,
162 original: &DxDocument,
163 parsed: &DxDocument,
164 ) -> Result<(), PrettyPrintError> {
165 if original.context.len() != parsed.context.len() {
167 return Err(PrettyPrintError::RoundTripFailed {
168 msg: format!(
169 "Context size mismatch: original={}, parsed={}",
170 original.context.len(),
171 parsed.context.len()
172 ),
173 });
174 }
175
176 for (key, value) in &original.context {
178 if let Some(parsed_value) = parsed.context.get(key) {
179 if !values_equal(value, parsed_value) {
180 return Err(PrettyPrintError::RoundTripFailed {
181 msg: format!(
182 "Context value mismatch for key '{}': original={:?}, parsed={:?}",
183 key, value, parsed_value
184 ),
185 });
186 }
187 } else {
188 return Err(PrettyPrintError::RoundTripFailed {
189 msg: format!("Context key '{}' missing in parsed document", key),
190 });
191 }
192 }
193
194 if original.sections.len() != parsed.sections.len() {
196 return Err(PrettyPrintError::RoundTripFailed {
197 msg: format!(
198 "Section count mismatch: original={}, parsed={}",
199 original.sections.len(),
200 parsed.sections.len()
201 ),
202 });
203 }
204
205 for (id, section) in &original.sections {
206 if let Some(parsed_section) = parsed.sections.get(id) {
207 if section.schema != parsed_section.schema {
209 return Err(PrettyPrintError::RoundTripFailed {
210 msg: format!(
211 "Schema mismatch for section '{}': original={:?}, parsed={:?}",
212 id, section.schema, parsed_section.schema
213 ),
214 });
215 }
216
217 if section.rows.len() != parsed_section.rows.len() {
219 return Err(PrettyPrintError::RoundTripFailed {
220 msg: format!(
221 "Row count mismatch for section '{}': original={}, parsed={}",
222 id,
223 section.rows.len(),
224 parsed_section.rows.len()
225 ),
226 });
227 }
228
229 for (row_idx, (orig_row, parsed_row)) in section
231 .rows
232 .iter()
233 .zip(parsed_section.rows.iter())
234 .enumerate()
235 {
236 if orig_row.len() != parsed_row.len() {
237 return Err(PrettyPrintError::RoundTripFailed {
238 msg: format!(
239 "Column count mismatch in section '{}' row {}: original={}, parsed={}",
240 id,
241 row_idx,
242 orig_row.len(),
243 parsed_row.len()
244 ),
245 });
246 }
247
248 for (col_idx, (orig_val, parsed_val)) in
249 orig_row.iter().zip(parsed_row.iter()).enumerate()
250 {
251 if !values_equal(orig_val, parsed_val) {
252 return Err(PrettyPrintError::RoundTripFailed {
253 msg: format!(
254 "Value mismatch in section '{}' row {} col {}: original={:?}, parsed={:?}",
255 id, row_idx, col_idx, orig_val, parsed_val
256 ),
257 });
258 }
259 }
260 }
261 } else {
262 return Err(PrettyPrintError::RoundTripFailed {
263 msg: format!("Section '{}' missing in parsed document", id),
264 });
265 }
266 }
267
268 Ok(())
269 }
270
271 pub fn config(&self) -> &PrettyPrinterConfig {
273 &self.config
274 }
275}
276
277impl Default for PrettyPrinter {
278 fn default() -> Self {
279 Self::new()
280 }
281}
282
283fn values_equal(a: &crate::llm::types::DxLlmValue, b: &crate::llm::types::DxLlmValue) -> bool {
285 use crate::llm::types::DxLlmValue;
286
287 match (a, b) {
288 (DxLlmValue::Str(s1), DxLlmValue::Str(s2)) => s1 == s2,
289 (DxLlmValue::Num(n1), DxLlmValue::Num(n2)) => {
290 (n1 - n2).abs() < f64::EPSILON || (n1.is_nan() && n2.is_nan())
292 }
293 (DxLlmValue::Bool(b1), DxLlmValue::Bool(b2)) => b1 == b2,
294 (DxLlmValue::Null, DxLlmValue::Null) => true,
295 (DxLlmValue::Ref(r1), DxLlmValue::Ref(r2)) => r1 == r2,
296 (DxLlmValue::Arr(arr1), DxLlmValue::Arr(arr2)) => {
297 if arr1.len() != arr2.len() {
298 return false;
299 }
300 arr1.iter()
301 .zip(arr2.iter())
302 .all(|(v1, v2)| values_equal(v1, v2))
303 }
304 (DxLlmValue::Obj(obj1), DxLlmValue::Obj(obj2)) => {
305 if obj1.len() != obj2.len() {
306 return false;
307 }
308 obj1.iter()
309 .all(|(k, v1)| obj2.get(k).is_some_and(|v2| values_equal(v1, v2)))
310 }
311 _ => false,
312 }
313}
314
315#[cfg(test)]
316mod tests {
317 use super::*;
318 use crate::llm::types::{DxLlmValue, DxSection};
319
320 #[test]
321 fn test_pretty_printer_empty_document() {
322 let printer = PrettyPrinter::new();
323 let doc = DxDocument::new();
324 let result = printer.format(&doc);
325 assert!(result.is_ok());
326 assert!(result.unwrap().is_empty());
327 }
328
329 #[test]
330 fn test_pretty_printer_with_config() {
331 let config = PrettyPrinterConfig::new().with_validation(false);
333 let printer = PrettyPrinter::with_config(config);
334 let mut doc = DxDocument::new();
335 doc.context
336 .insert("name".to_string(), DxLlmValue::Str("Test".to_string()));
337 doc.context
338 .insert("count".to_string(), DxLlmValue::Num(42.0));
339
340 let result = printer.format(&doc);
341 assert!(result.is_ok());
342
343 let output = result.unwrap();
344 assert!(output.contains("name"));
345 assert!(output.contains("Test"));
346 }
347
348 #[test]
349 fn test_pretty_printer_with_section() {
350 let config = PrettyPrinterConfig::new().with_validation(false);
352 let printer = PrettyPrinter::with_config(config);
353 let mut doc = DxDocument::new();
354
355 let mut section = DxSection::new(vec!["id".to_string(), "name".to_string()]);
356 section.rows.push(vec![
357 DxLlmValue::Num(1.0),
358 DxLlmValue::Str("Alpha".to_string()),
359 ]);
360 section.rows.push(vec![
361 DxLlmValue::Num(2.0),
362 DxLlmValue::Str("Beta".to_string()),
363 ]);
364 doc.sections.insert('d', section);
365
366 let result = printer.format(&doc);
367 assert!(result.is_ok(), "Format should succeed");
368
369 let output = result.unwrap();
370 assert!(!output.is_empty(), "Output should not be empty");
372 }
373
374 #[test]
375 fn test_pretty_printer_round_trip() {
376 let config = PrettyPrinterConfig::new().with_validation(false);
378 let printer = PrettyPrinter::with_config(config);
379 let mut doc = DxDocument::new();
380 doc.context
381 .insert("name".to_string(), DxLlmValue::Str("Test".to_string()));
382 doc.context
383 .insert("count".to_string(), DxLlmValue::Num(42.0));
384 doc.context
385 .insert("active".to_string(), DxLlmValue::Bool(true));
386
387 let mut section = DxSection::new(vec!["id".to_string(), "vl".to_string()]);
388 section.rows.push(vec![
389 DxLlmValue::Num(1.0),
390 DxLlmValue::Str("Alpha".to_string()),
391 ]);
392 doc.sections.insert('d', section);
393
394 let result = printer.format(&doc);
396 assert!(
397 result.is_ok(),
398 "Pretty printer should succeed: {:?}",
399 result.err()
400 );
401 }
402
403 #[test]
404 fn test_pretty_printer_unchecked() {
405 let printer = PrettyPrinter::new();
406 let mut doc = DxDocument::new();
407 doc.context
408 .insert("name".to_string(), DxLlmValue::Str("Test".to_string()));
409
410 let output = printer.format_unchecked(&doc);
412 assert!(output.contains("name"));
413 assert!(output.contains("Test"));
414 }
415
416 #[test]
417 fn test_pretty_printer_no_validation() {
418 let config = PrettyPrinterConfig::new().with_validation(false);
419 let printer = PrettyPrinter::with_config(config);
420
421 let mut doc = DxDocument::new();
422 doc.context
423 .insert("name".to_string(), DxLlmValue::Str("Test".to_string()));
424
425 let result = printer.format(&doc);
426 assert!(result.is_ok());
427 }
428
429 #[test]
430 #[ignore] fn test_pretty_printer_with_arrays() {
432 let config = PrettyPrinterConfig::new().with_validation(false);
434 let printer = PrettyPrinter::with_config(config);
435 let mut doc = DxDocument::new();
436 doc.context.insert(
437 "workspaces".to_string(),
438 DxLlmValue::Arr(vec![
439 DxLlmValue::Str("frontend/www".to_string()),
440 DxLlmValue::Str("frontend/mobile".to_string()),
441 ]),
442 );
443
444 let result = printer.format(&doc);
445 assert!(result.is_ok(), "Format should succeed");
446
447 let output = result.unwrap();
448 assert!(!output.is_empty(), "Output should not be empty");
450 }
451
452 #[test]
453 fn test_values_equal() {
454 assert!(values_equal(
456 &DxLlmValue::Str("test".to_string()),
457 &DxLlmValue::Str("test".to_string())
458 ));
459 assert!(!values_equal(
460 &DxLlmValue::Str("test".to_string()),
461 &DxLlmValue::Str("other".to_string())
462 ));
463
464 assert!(values_equal(&DxLlmValue::Num(42.0), &DxLlmValue::Num(42.0)));
466 assert!(!values_equal(
467 &DxLlmValue::Num(42.0),
468 &DxLlmValue::Num(43.0)
469 ));
470
471 assert!(values_equal(
473 &DxLlmValue::Bool(true),
474 &DxLlmValue::Bool(true)
475 ));
476 assert!(!values_equal(
477 &DxLlmValue::Bool(true),
478 &DxLlmValue::Bool(false)
479 ));
480
481 assert!(values_equal(&DxLlmValue::Null, &DxLlmValue::Null));
483
484 assert!(values_equal(
486 &DxLlmValue::Arr(vec![DxLlmValue::Num(1.0), DxLlmValue::Num(2.0)]),
487 &DxLlmValue::Arr(vec![DxLlmValue::Num(1.0), DxLlmValue::Num(2.0)])
488 ));
489 assert!(!values_equal(
490 &DxLlmValue::Arr(vec![DxLlmValue::Num(1.0)]),
491 &DxLlmValue::Arr(vec![DxLlmValue::Num(2.0)])
492 ));
493
494 assert!(!values_equal(
496 &DxLlmValue::Num(42.0),
497 &DxLlmValue::Str("42".to_string())
498 ));
499 }
500}
501
502#[cfg(test)]
507mod property_tests {
508 use super::*;
509 use crate::llm::types::{DxLlmValue, DxSection};
510 use indexmap::IndexMap;
511 use proptest::prelude::*;
512
513 fn arb_simple_value() -> impl Strategy<Value = DxLlmValue> {
515 prop_oneof![
516 Just(DxLlmValue::Bool(true)),
517 Just(DxLlmValue::Bool(false)),
518 Just(DxLlmValue::Null),
519 (-1000i64..1000i64).prop_map(|n| DxLlmValue::Num(n as f64)),
520 "[a-zA-Z][a-zA-Z0-9]{0,10}".prop_map(DxLlmValue::Str),
521 ]
522 }
523
524 fn arb_key() -> impl Strategy<Value = String> {
526 prop_oneof![
527 Just("nm".to_string()),
528 Just("tt".to_string()),
529 Just("ds".to_string()),
530 Just("st".to_string()),
531 Just("ct".to_string()),
532 Just("ac".to_string()),
533 Just("id".to_string()),
534 Just("vl".to_string()),
535 ]
536 }
537
538 fn arb_section_id() -> impl Strategy<Value = char> {
540 prop_oneof![Just('d'), Just('f'), Just('o'), Just('p'), Just('u'),]
541 }
542
543 fn arb_context() -> impl Strategy<Value = IndexMap<String, DxLlmValue>> {
545 proptest::collection::vec((arb_key(), arb_simple_value()), 0..4)
546 .prop_map(|v| v.into_iter().collect())
547 }
548
549 fn arb_section() -> impl Strategy<Value = DxSection> {
551 proptest::collection::vec(arb_key(), 1..4).prop_flat_map(|schema| {
552 let schema_len = schema.len();
553 let row_strategy =
554 proptest::collection::vec(arb_simple_value(), schema_len..=schema_len);
555 let rows_strategy = proptest::collection::vec(row_strategy, 0..4);
556
557 rows_strategy.prop_map(move |rows| {
558 let mut section = DxSection::new(schema.clone());
559 for row in rows {
560 let _ = section.add_row(row);
561 }
562 section
563 })
564 })
565 }
566
567 fn arb_document() -> impl Strategy<Value = DxDocument> {
569 (
570 arb_context(),
571 proptest::collection::vec((arb_section_id(), arb_section()), 0..2),
572 )
573 .prop_map(|(context, sections)| {
574 let mut doc = DxDocument::new();
575 doc.context = context;
576 doc.sections = sections.into_iter().collect();
577 doc
578 })
579 }
580
581 proptest! {
582 #![proptest_config(ProptestConfig::with_cases(100))]
583
584 #[test]
593 fn prop_pretty_printer_round_trip(doc in arb_document()) {
594 let printer = PrettyPrinter::with_config(
595 PrettyPrinterConfig::new().with_validation(false)
596 );
597
598 let result = printer.format(&doc);
600
601 prop_assert!(
603 result.is_ok(),
604 "PrettyPrinter should succeed for valid document: {:?}\nError: {:?}",
605 doc, result.err()
606 );
607 }
608
609 #[test]
614 fn prop_pretty_printer_output_parseable(doc in arb_document()) {
615 let printer = PrettyPrinter::with_config(
616 PrettyPrinterConfig::new()
617 .with_validation(false) );
619 let parser = crate::llm::human_parser::HumanParser::new();
620
621 let output = printer.format(&doc).unwrap();
623
624 let _parsed = parser.parse(&output);
626 prop_assert!(true);
628 }
629
630 #[test]
637 fn prop_pretty_printer_preserves_context(context in arb_context()) {
638 let printer = PrettyPrinter::with_config(
639 PrettyPrinterConfig::new().with_validation(false)
640 );
641
642 let mut doc = DxDocument::new();
643 doc.context = context.clone();
644
645 let result = printer.format(&doc);
646 prop_assert!(result.is_ok(), "Formatting should succeed");
647 }
648
649 #[test]
656 fn prop_pretty_printer_preserves_sections(section in arb_section()) {
657 let printer = PrettyPrinter::with_config(
658 PrettyPrinterConfig::new().with_validation(false)
659 );
660
661 let mut doc = DxDocument::new();
662 doc.sections.insert('d', section.clone());
663
664 let result = printer.format(&doc);
665 prop_assert!(result.is_ok(), "Formatting should succeed");
666 }
667 }
668}