Add pretty-print roundtrip tests and Fix PIVOT pretty-print

partiql · Jul 31, 2024 · 519d540 · 519d540
1 parent 63b0023
commit 519d540
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 13 deletions.
diff --git a/partiql-ast/src/pretty.rs b/partiql-ast/src/pretty.rs
@@ -299,9 +299,9 @@ impl PrettyDoc for ProjectionKind {
             }
             ProjectionKind::ProjectPivot(ProjectPivot { key, value }) => {
                 let parts = [
-                    key.pretty_doc(arena),
-                    arena.text("AT"),
                     value.pretty_doc(arena),
+                    arena.text("AT"),
+                    key.pretty_doc(arena),
                 ];
                 let decl = arena.intersperse(parts, arena.space()).group();
                 pretty_annotated_doc("PIVOT", decl, arena)
@@ -741,7 +741,7 @@ impl PrettyDoc for SearchedCase {
 
 fn case_branches<'b, D, A>(
     arena: &'b D,
-    cases: &'b Vec<ExprPair>,
+    cases: &'b [ExprPair],
     default: &'b Option<Box<Expr>>,
 ) -> impl Iterator<Item = DocBuilder<'b, D, A>>
 where

diff --git a/partiql/tests/pretty.rs b/partiql/tests/pretty.rs
@@ -1,4 +1,5 @@
 use itertools::Itertools;
+use partiql_ast::ast::{AstNode, TopLevelQuery};
 use partiql_ast::pretty::ToPretty;
 use partiql_parser::ParserResult;
 
@@ -15,13 +16,18 @@ fn pretty_print_test(name: &str, statement: &str) {
     assert!(res.is_ok());
     let res = res.unwrap();
 
-    // TODO https://github.com/partiql/partiql-lang-rust/issues/473
+    pretty_print_output_test(name, statement, &res.ast);
+    pretty_print_roundtrip_test(name, statement, &res.ast);
+}
 
+#[track_caller]
+fn pretty_print_output_test(name: &str, statement: &str, statement_ast: &AstNode<TopLevelQuery>) {
+    // TODO https://github.com/partiql/partiql-lang-rust/issues/473
     let doc = [180, 120, 80, 40, 30, 20, 10]
         .into_iter()
         .map(|w| {
             let header = format!("{:-<w$}", "");
-            let ast = format!("{}\n", res.ast.to_pretty_string(w).unwrap());
+            let ast = format!("{}\n", statement_ast.to_pretty_string(w).unwrap());
             format!("{header}\n{ast}")
         })
         .join("\n");
@@ -33,6 +39,22 @@ fn pretty_print_test(name: &str, statement: &str) {
     insta::assert_snapshot!(name, doc)
 }
 
+#[track_caller]
+fn pretty_print_roundtrip_test(
+    name: &str,
+    statement: &str,
+    statement_ast: &AstNode<TopLevelQuery>,
+) {
+    let pretty = statement_ast.to_pretty_string(40).unwrap();
+
+    let reparsed = parse(pretty.as_str());
+    assert!(reparsed.is_ok());
+
+    let pretty2 = reparsed.unwrap().ast.to_pretty_string(40).unwrap();
+
+    assert_eq!(pretty, pretty2);
+}
+
 #[test]
 fn pretty() {
     pretty_print_test(

diff --git a/partiql/tests/snapshots/pretty__pretty_pivot.snap b/partiql/tests/snapshots/pretty__pretty_pivot.snap
@@ -1,5 +1,5 @@
 ---
-source: partiql-ast/tests/pretty.rs
+source: partiql/tests/pretty.rs
 expression: doc
 ---
 ========================================================================================================================================================================================================
@@ -12,36 +12,36 @@ expression: doc
 ========================================================================================================================================================================================================
 
 ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
-PIVOT foo.b AT foo.a FROM << { 'a': 1, 'b': 'I' }, { 'a': 2, 'b': 'II' }, { 'a': 3, 'b': 'III' } >> AS foo ORDER BY a LIMIT 1 OFFSET 1
+PIVOT foo.a AT foo.b FROM << { 'a': 1, 'b': 'I' }, { 'a': 2, 'b': 'II' }, { 'a': 3, 'b': 'III' } >> AS foo ORDER BY a LIMIT 1 OFFSET 1
 
 ------------------------------------------------------------------------------------------------------------------------
-PIVOT foo.b AT foo.a FROM << { 'a': 1, 'b': 'I' }, { 'a': 2, 'b': 'II' }, { 'a': 3, 'b': 'III' } >> AS foo ORDER BY a
+PIVOT foo.a AT foo.b FROM << { 'a': 1, 'b': 'I' }, { 'a': 2, 'b': 'II' }, { 'a': 3, 'b': 'III' } >> AS foo ORDER BY a
 LIMIT 1 OFFSET 1
 
 --------------------------------------------------------------------------------
-PIVOT foo.b AT foo.a FROM <<
+PIVOT foo.a AT foo.b FROM <<
     { 'a': 1, 'b': 'I' },
     { 'a': 2, 'b': 'II' },
     { 'a': 3, 'b': 'III' }
   >> AS foo ORDER BY a LIMIT 1 OFFSET 1
 
 ----------------------------------------
-PIVOT foo.b AT foo.a FROM <<
+PIVOT foo.a AT foo.b FROM <<
     { 'a': 1, 'b': 'I' },
     { 'a': 2, 'b': 'II' },
     { 'a': 3, 'b': 'III' }
   >> AS foo ORDER BY a LIMIT 1 OFFSET 1
 
 ------------------------------
-PIVOT foo.b AT foo.a FROM <<
+PIVOT foo.a AT foo.b FROM <<
     { 'a': 1, 'b': 'I' },
     { 'a': 2, 'b': 'II' },
     { 'a': 3, 'b': 'III' }
   >> AS foo ORDER BY a LIMIT 1
 OFFSET 1
 
 --------------------
-PIVOT foo.b AT foo.a
+PIVOT foo.a AT foo.b
 FROM <<
     {
       'a': 1,
@@ -60,7 +60,7 @@ ORDER BY a LIMIT 1
 OFFSET 1
 
 ----------
-PIVOT foo.b AT foo.a
+PIVOT foo.a AT foo.b
 FROM <<
     {
       'a': 1,