Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improvements to pretty-printing for AND/OR/CASE #482

Merged
merged 3 commits into from
Jul 31, 2024
Merged

Improvements to pretty-printing for AND/OR/CASE #482

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
88f61bd
Improvements to pretty-printing for AND/OR/CASE
jpschorr Jul 31, 2024
63b0023
Clean-up and some more minor improvements
jpschorr Jul 31, 2024
6b546aa
Add pretty-print roundtrip tests and Fix PIVOT pretty-print
jpschorr Jul 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

## [Unreleased]
### Changed

- partiql-ast: fixed pretty-printing of `PIVOT`
- partiql-ast: improved pretty-printing of `CASE` and various clauses
-
### Added

### Fixed
Expand Down
175 changes: 101 additions & 74 deletions partiql-ast/src/pretty.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,14 +213,24 @@ impl PrettyDoc for Select {
D::Doc: Clone,
A: Clone,
{
fn format<'b, C, D, A>(child: &'b C, arena: &'b D) -> DocBuilder<'b, D, A>
where
D: DocAllocator<'b, A>,
D::Doc: Clone,
A: Clone,
C: PrettyDoc,
{
child.pretty_doc(arena).group()
}

fn delegate<'b, C, D, A>(child: &'b Option<C>, arena: &'b D) -> Option<DocBuilder<'b, D, A>>
where
D: DocAllocator<'b, A>,
D::Doc: Clone,
A: Clone,
C: PrettyDoc,
{
child.as_ref().map(|inner| inner.pretty_doc(arena).group())
child.as_ref().map(|inner| format(inner, arena))
}

let Select {
Expand All @@ -232,25 +242,32 @@ impl PrettyDoc for Select {
group_by,
having,
} = self;
let clauses = [
Some(project.pretty_doc(arena).group()),
let mut clauses = [
Some(format(project, arena)),
delegate(exclude, arena),
from.as_ref().map(|inner| inner.pretty_doc(arena).group()),
from_let
.as_ref()
.map(|inner| inner.pretty_doc(arena).group()),
where_clause
.as_ref()
.map(|inner| inner.pretty_doc(arena).group()),
group_by
.as_ref()
.map(|inner| inner.pretty_doc(arena).group()),
having.as_ref().map(|inner| inner.pretty_doc(arena).group()),
delegate(from, arena),
delegate(from_let, arena),
delegate(where_clause, arena),
delegate(group_by, arena),
delegate(having, arena),
]
.into_iter()
.flatten();

arena.intersperse(clauses, arena.softline()).group()
let mut result = arena.nil();
let separator = arena.line();
if let Some(first) = clauses.next() {
let mut curr = first;

for clause in clauses {
result = result.append(curr.append(separator.clone()).group());
curr = clause;
}

result = result.append(curr);
}

result
}
}

Expand Down Expand Up @@ -282,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)
Expand Down Expand Up @@ -401,6 +418,7 @@ impl PrettyDoc for Expr {
unreachable!();
}
}
.group()
}
}

Expand Down Expand Up @@ -570,10 +588,13 @@ impl PrettyDoc for BinOp {
let op = arena.text(sym);
let lhs = lhs.pretty_doc(arena).nest(nest);
let rhs = rhs.pretty_doc(arena).nest(nest);
let sep = arena.space();
let sep = if nest == 0 {
arena.space()
} else {
arena.softline()
};
let expr = arena.intersperse([lhs, op, rhs], sep).group();
let paren_expr = [arena.text("("), expr, arena.text(")")];
arena.concat(paren_expr).group()
pretty_parenthesized_doc(expr, arena).group()
}
}

Expand Down Expand Up @@ -698,30 +719,9 @@ impl PrettyDoc for SimpleCase {
default,
} = self;

let kw_case = arena.text("CASE");
let search = expr.pretty_doc(arena);
let branches = cases.iter().map(|ExprPair { first, second }| {
let kw_when = arena.text("WHEN");
let test = first.pretty_doc(arena);
let kw_then = arena.text("THEN");
let then = second.pretty_doc(arena);
arena
.intersperse([kw_when, test, kw_then, then], arena.space())
.group()
});
let branches = arena
.intersperse(branches, arena.softline())
.group()
.nest(MINOR_NEST_INDENT);
let default = default
.as_ref()
.map(|d| arena.text("ELSE ").append(d.pretty_doc(arena)));

if let Some(default) = default {
arena.intersperse([kw_case, search, branches, default], arena.softline())
} else {
arena.intersperse([kw_case, search, branches], arena.softline())
}
let branches = case_branches(arena, cases, default);
pretty_seq_doc(branches, "CASE", Some(search), "END", " ", arena)
}
}

Expand All @@ -734,30 +734,37 @@ impl PrettyDoc for SearchedCase {
{
let SearchedCase { cases, default } = self;

let kw_case = arena.text("CASE");
let branches = cases.iter().map(|ExprPair { first, second }| {
let branches = case_branches(arena, cases, default);
pretty_seq_doc(branches, "CASE", None, "END", " ", arena)
}
}

fn case_branches<'b, D, A>(
arena: &'b D,
cases: &'b [ExprPair],
default: &'b Option<Box<Expr>>,
) -> impl Iterator<Item = DocBuilder<'b, D, A>>
where
D: DocAllocator<'b, A>,
D::Doc: Clone,
A: Clone + 'b,
{
cases
.iter()
.map(|ExprPair { first, second }| {
let kw_when = arena.text("WHEN");
let test = first.pretty_doc(arena);
let kw_then = arena.text("THEN");
let then = second.pretty_doc(arena);
arena
.intersperse([kw_when, test, kw_then, then], arena.space())
.group()
});
let branches = arena
.intersperse(branches, arena.softline())
.group()
.nest(MINOR_NEST_INDENT);
let default = default
.as_ref()
.map(|d| arena.text("ELSE ").append(d.pretty_doc(arena)));

if let Some(default) = default {
arena.intersperse([kw_case, branches, default], arena.softline())
} else {
arena.intersperse([kw_case, branches], arena.softline())
}
}
})
.chain(
default
.iter()
.map(|d| arena.text("ELSE ").append(d.pretty_doc(arena)).group()),
)
}

impl PrettyDoc for Struct {
Expand Down Expand Up @@ -1258,41 +1265,61 @@ where
D::Doc: Clone,
A: Clone,
{
arena
.text("(")
.append(arena.space())
.append(doc)
.append(arena.space())
.append(arena.text(")"))
.group()
arena.text("(").append(doc).append(arena.text(")")).group()
}

fn pretty_seq<'i, 'b, I, P, D, A>(
list: I,
fn pretty_seq_doc<'i, 'b, I, E, D, A>(
seq: I,
start: &'static str,
qualifier: Option<E>,
end: &'static str,
sep: &'static str,
arena: &'b D,
) -> DocBuilder<'b, D, A>
where
I: IntoIterator<Item = &'b P>,
P: PrettyDoc + 'b,
E: Pretty<'b, D, A>,
I: IntoIterator<Item = E>,
D: DocAllocator<'b, A>,
D::Doc: Clone,
A: Clone,
{
let start = arena.text(start);
let end = arena.text(end);
let sep = arena.text(sep).append(arena.line());
let seq = list.into_iter().map(|l| l.pretty_doc(arena));
let body = arena.line().append(arena.intersperse(seq, sep)).group();
let start = if let Some(qual) = qualifier {
start.append(arena.space()).append(qual)
} else {
start
};
let body = arena
.line()
.append(arena.intersperse(seq, sep))
.append(arena.line())
.group();
start
.append(body.nest(MINOR_NEST_INDENT))
.append(arena.line())
.append(end)
.group()
}

fn pretty_seq<'i, 'b, I, P, D, A>(
list: I,
start: &'static str,
end: &'static str,
sep: &'static str,
arena: &'b D,
) -> DocBuilder<'b, D, A>
where
I: IntoIterator<Item = &'b P>,
P: PrettyDoc + 'b,
D: DocAllocator<'b, A>,
D::Doc: Clone,
A: Clone,
{
let seq = list.into_iter().map(|l| l.pretty_doc(arena));
pretty_seq_doc(seq, start, None, end, sep, arena)
}

fn pretty_list<'b, I, P, D, A>(list: I, nest: isize, arena: &'b D) -> DocBuilder<'b, D, A>
where
I: IntoIterator<Item = &'b P>,
Expand Down
34 changes: 32 additions & 2 deletions partiql/tests/pretty.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use itertools::Itertools;
use partiql_ast::ast::{AstNode, TopLevelQuery};
use partiql_ast::pretty::ToPretty;
use partiql_parser::ParserResult;

Expand All @@ -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(&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");
Expand All @@ -33,6 +39,18 @@ fn pretty_print_test(name: &str, statement: &str) {
insta::assert_snapshot!(name, doc)
}

#[track_caller]
fn pretty_print_roundtrip_test(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(
Expand Down Expand Up @@ -169,6 +187,18 @@ fn pretty_pivot() {
);
}

#[test]
fn pretty_ands_and_ors() {
pretty_print_test(
"pretty_ands_and_ors",
"
SELECT *
FROM test_data AS test_data
WHERE ((((test_data.country_code <> 'Distinctio.') AND ((test_data.* < false) AND (NOT (test_data.description LIKE 'Esse solam.') AND NOT (test_data.transaction_id LIKE 'Esset accusata.')))) OR (test_data.test_address <> 'Potest. Sed.')) AND (test_data.* > -28.146858383543243))
",
);
}

#[test]
fn pretty_exclude() {
pretty_print_test(
Expand Down
Loading
Loading