squawk_fmt/
lib.rs

1use itertools::Itertools;
2use squawk_syntax::ast::{self, AstNode};
3use tiny_pretty::Doc;
4use tiny_pretty::{PrintOptions, print};
5
6fn build_source_file(source_file: &ast::SourceFile) -> Doc<'_> {
7    let mut doc = Doc::nil();
8    for stmt in source_file.stmts() {
9        match stmt {
10            ast::Stmt::Select(select) => {
11                doc = doc.append(build_select_doc(select));
12            }
13            ast::Stmt::CreateTable(create_table) => {
14                doc = doc.append(build_create_table(create_table))
15            }
16            _ => (),
17        }
18        doc = doc
19            .append(Doc::text(";"))
20            .append(Doc::empty_line())
21            .append(Doc::empty_line());
22    }
23    doc
24}
25
26fn build_create_table<'a>(create_table: ast::CreateTable) -> Doc<'a> {
27    Doc::text("create")
28        .append(Doc::space())
29        .append(Doc::text("table"))
30        .append(Doc::space())
31        .append(Doc::text(
32            create_table.path().map(|x| x.syntax().to_string()).unwrap(),
33        ))
34        .append(Doc::text("("))
35        .append(
36            Doc::line_or_nil()
37                .append(Doc::list(
38                    Itertools::intersperse(
39                        create_table
40                            .table_arg_list()
41                            .unwrap()
42                            .args()
43                            .map(build_table_arg),
44                        Doc::text(",").append(Doc::line_or_space()),
45                    )
46                    .collect(),
47                ))
48                .nest(2)
49                .append(Doc::line_or_nil())
50                .group(),
51        )
52        .append(Doc::text(")"))
53}
54
55fn build_table_arg<'a>(create_table: ast::TableArg) -> Doc<'a> {
56    match create_table {
57        ast::TableArg::Column(column) => Doc::text(column.name().unwrap().syntax().to_string())
58            .append(Doc::space())
59            .append(Doc::text(column.ty().unwrap().syntax().to_string())),
60        ast::TableArg::LikeClause(_like_clause) => todo!(),
61        ast::TableArg::TableConstraint(_table_constraint) => todo!(),
62    }
63}
64
65fn build_select_doc<'a>(select: ast::Select) -> Doc<'a> {
66    let mut doc = Doc::text("select").append(Doc::space());
67
68    if let Some(targets) = select
69        .select_clause()
70        .and_then(|x| x.target_list())
71        .map(|x| x.targets())
72    {
73        doc = doc
74            .append(
75                Doc::line_or_nil().append(Doc::list(
76                    Itertools::intersperse(
77                        targets.flat_map(|x| Some(Doc::text(x.expr()?.syntax().to_string()))),
78                        Doc::text(",").append(Doc::line_or_space()),
79                    )
80                    .collect(),
81                )),
82            )
83            .nest(2);
84    }
85
86    if let Some(from) = &select.from_clause() {
87        doc = doc.append(
88            Doc::line_or_space()
89                .append(Doc::text("from"))
90                .append(Doc::space())
91                .append(Doc::text(
92                    from.from_items().next().unwrap().syntax().to_string(),
93                )),
94        );
95    }
96
97    if let Some(group) = &select.group_by_clause() {
98        doc = doc.append(
99            Doc::line_or_space()
100                .append(Doc::text("group by"))
101                .append(Doc::space())
102                .append(Doc::text(
103                    group.group_by_list().unwrap().syntax().to_string(),
104                )),
105        );
106    }
107
108    doc.group()
109}
110
111pub fn fmt(text: &str) -> String {
112    let parse = ast::SourceFile::parse(text);
113    let file = parse.tree();
114    let doc = build_source_file(&file);
115    print(&doc, &PrintOptions::default())
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use insta::assert_snapshot;
122
123    #[test]
124    fn select() {
125        assert_snapshot!(fmt("
126select a(), date_trunc(1, 2), foo(), avg(a - b), bar(carrot), buzz(potato), foo.b from t group by c;
127create table t(a int, b text);
128"), @r"
129        select 
130          a(),
131          date_trunc(1, 2),
132          foo(),
133          avg(a - b),
134          bar(carrot),
135          buzz(potato),
136          foo.b
137        from t
138        group by c;
139
140        create table t(a int, b text);
141        ");
142    }
143}
squawk_fmt/lib.rs

squawk_fmt/
lib.rs
