More formatting options (#78)

lu-zero · web-flow · commit 31a2dfe950c5 · 2025-02-01T10:11:53.000-05:00
* chore: Upgrade to Winnow 0.6.26

* refactor: Resolve deprecations

* refactor: Switch from Parser to ModalParser

* chore: Upgrade to Winnow 0.7.0

* refactor: Remove use of ErrMode

* fix: Fix the span computation and usage

* feat: Extend ALTER TABLE actions support

* feat: Always keep track of the indentation level of blocks

This way is possible to have inlined blocks that break if
the argument limit is hit.

* feat!: Consider SELECT ALL/DISTINCT a single token

* feat: Give joins their own token kind

* fix: Allow inline case-end

* fix: Improve inlining of subqueries
diff --git a/src/formatter.rs b/src/formatter.rs
@@ -101,6 +101,10 @@ pub(crate) fn format(
                 formatter.format_newline_reserved_word(token, &mut formatted_query);
                 formatter.previous_reserved_word = Some(token);
             }
+            TokenKind::Join => {
+                formatter.format_newline_reserved_word(token, &mut formatted_query);
+                formatter.previous_reserved_word = Some(token);
+            }
             TokenKind::Reserved => {
                 formatter.format_with_spaces(token, &mut formatted_query);
                 formatter.previous_reserved_word = Some(token);
@@ -163,7 +167,14 @@ impl<'a> Formatter<'a> {
             indentation: Indentation::new(options),
             inline_block: InlineBlock::new(
                 options.max_inline_block,
-                options.max_inline_arguments.is_none(),
+                match (options.max_inline_arguments, options.max_inline_top_level) {
+                    (Some(max_inline_args), Some(max_inline_top)) => {
+                        max_inline_args.min(max_inline_top)
+                    }
+                    (Some(max_inline_args), None) => max_inline_args,
+                    (None, Some(max_inline_top)) => max_inline_top,
+                    (None, None) => 0,
+                },
             ),
             block_level: 0,
         }
@@ -204,28 +215,51 @@ impl<'a> Formatter<'a> {
         self.add_new_line(query);
     }
 
+    // if we are inside an inline block we decide our behaviour as if we are an inline argument
+    fn top_level_behavior(&self) -> (bool, bool) {
+        let span_len = self.top_level_tokens_span();
+        let block_len = self.inline_block.cur_len();
+        if block_len > 0 {
+            let limit = self.options.max_inline_arguments.unwrap_or(0);
+            (limit < block_len, limit < span_len)
+        } else {
+            (
+                true,
+                self.options
+                    .max_inline_top_level
+                    .map_or(true, |limit| limit < span_len),
+            )
+        }
+    }
+
     fn format_top_level_reserved_word(&mut self, token: &Token<'_>, query: &mut String) {
         let span_len = self.top_level_tokens_span();
-        self.indentation.decrease_top_level();
-        self.add_new_line(query);
-        self.indentation.increase_top_level(span_len);
+        let (newline_before, newline_after) = self.top_level_behavior();
+
+        if newline_before {
+            self.indentation.decrease_top_level();
+            self.add_new_line(query);
+        }
         query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value)));
-        if self
-            .options
-            .max_inline_top_level
-            .map_or(true, |limit| limit < span_len)
-        {
+        if newline_after {
+            self.indentation.increase_top_level(span_len);
             self.add_new_line(query);
         } else {
             query.push(' ');
         }
     }
 
     fn format_top_level_reserved_word_no_indent(&mut self, token: &Token<'_>, query: &mut String) {
-        self.indentation.decrease_top_level();
-        self.add_new_line(query);
+        let (newline_before, newline_after) = self.top_level_behavior();
+
+        if newline_before {
+            self.indentation.decrease_top_level();
+            self.add_new_line(query);
+        }
         query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value)));
-        self.add_new_line(query);
+        if newline_after {
+            self.add_new_line(query);
+        }
     }
 
     fn format_newline_reserved_word(&mut self, token: &Token<'_>, query: &mut String) {
@@ -297,8 +331,9 @@ impl<'a> Formatter<'a> {
 
         self.inline_block.begin_if_possible(self.tokens, self.index);
 
+        self.indentation.increase_block_level();
+
         if !self.inline_block.is_active() {
-            self.indentation.increase_block_level();
             self.add_new_line(query);
         }
     }
@@ -330,6 +365,8 @@ impl<'a> Formatter<'a> {
 
         token.value = &value;
 
+        self.indentation.decrease_block_level();
+
         if self.inline_block.is_active() {
             self.inline_block.end();
             if token.value.to_lowercase() == "end" {
@@ -340,7 +377,6 @@ impl<'a> Formatter<'a> {
                 self.format_with_space_after(&token, query);
             }
         } else {
-            self.indentation.decrease_block_level();
             self.add_new_line(query);
             self.format_with_spaces(&token, query);
         }
@@ -493,7 +529,6 @@ impl<'a> Formatter<'a> {
     }
 
     fn top_level_tokens_span(&self) -> usize {
-        assert_eq!(self.tokens[self.index].kind, TokenKind::ReservedTopLevel);
         let mut block_level = self.block_level;
 
         self.tokens[self.index..]
@@ -508,7 +543,9 @@ impl<'a> Formatter<'a> {
                     block_level = block_level.saturating_sub(1);
                     block_level > self.block_level
                 }
-                TokenKind::ReservedTopLevel => block_level != self.block_level,
+                TokenKind::ReservedTopLevel | TokenKind::ReservedTopLevelNoIndent => {
+                    block_level != self.block_level
+                }
                 _ => true,
             })
             .map(|token| token.value.len())
diff --git a/src/inline_block.rs b/src/inline_block.rs
@@ -1,86 +1,129 @@
 use crate::tokenizer::{Token, TokenKind};
 
+pub(crate) struct BlockInfo {
+    length: usize,
+    has_forbidden_tokens: bool,
+    top_level_token_span: usize,
+}
+
 pub(crate) struct InlineBlock {
     level: usize,
     inline_max_length: usize,
-    newline_on_reserved: bool,
+    newline_on_reserved_limit: usize,
+    info: Vec<BlockInfo>,
 }
 
 impl Default for InlineBlock {
     fn default() -> Self {
         InlineBlock {
+            info: Vec::new(),
             level: 0,
             inline_max_length: 50,
-            newline_on_reserved: true,
+            newline_on_reserved_limit: 0,
         }
     }
 }
 
 impl InlineBlock {
-    pub fn new(inline_max_length: usize, newline_on_reserved: bool) -> Self {
+    pub fn new(inline_max_length: usize, newline_on_reserved_limit: usize) -> Self {
         InlineBlock {
-            level: 0,
             inline_max_length,
-            newline_on_reserved,
+            newline_on_reserved_limit,
+            ..Default::default()
         }
     }
 
+    fn is_inline_block(&self, info: &BlockInfo) -> bool {
+        !info.has_forbidden_tokens
+            && info.length <= self.inline_max_length
+            && info.top_level_token_span <= self.newline_on_reserved_limit
+    }
+
     pub fn begin_if_possible(&mut self, tokens: &[Token<'_>], index: usize) {
-        if self.level == 0 && self.is_inline_block(tokens, index) {
+        let info = self.build_info(tokens, index);
+        if self.level == 0 && self.is_inline_block(&info) {
             self.level = 1;
         } else if self.level > 0 {
             self.level += 1;
         } else {
             self.level = 0;
         }
+        if self.level > 0 {
+            self.info.push(info);
+        }
     }
 
     pub fn end(&mut self) {
+        self.info.pop();
         self.level -= 1;
     }
 
     pub fn is_active(&self) -> bool {
         self.level > 0
     }
 
-    fn is_inline_block(&self, tokens: &[Token<'_>], index: usize) -> bool {
+    /// Get the current inline block length
+    pub fn cur_len(&self) -> usize {
+        self.info.last().map_or(0, |info| info.length)
+    }
+
+    fn build_info(&self, tokens: &[Token<'_>], index: usize) -> BlockInfo {
         let mut length = 0;
         let mut level = 0;
+        let mut top_level_token_span = 0;
+        let mut start_top_level = -1;
+        let mut start_span = 0;
+        let mut has_forbidden_tokens = false;
 
         for token in &tokens[index..] {
             length += token.value.len();
 
-            // Overran max length
-            if length > self.inline_max_length {
-                return false;
-            }
-            if token.kind == TokenKind::OpenParen {
-                level += 1;
-            } else if token.kind == TokenKind::CloseParen {
-                level -= 1;
-                if level == 0 {
-                    return true;
+            match token.kind {
+                TokenKind::ReservedTopLevel | TokenKind::ReservedTopLevelNoIndent => {
+                    if start_top_level != -1 {
+                        if start_top_level == level {
+                            top_level_token_span = top_level_token_span.max(length - start_span);
+                            start_top_level = -1;
+                        }
+                    } else {
+                        start_top_level = level;
+                        start_span = length;
+                    }
                 }
+                TokenKind::OpenParen => {
+                    level += 1;
+                }
+                TokenKind::CloseParen => {
+                    level -= 1;
+                    if level == 0 {
+                        break;
+                    }
+                }
+                _ => {}
             }
 
             if self.is_forbidden_token(token) {
-                return false;
+                has_forbidden_tokens = true;
             }
         }
 
-        false
+        // broken syntax let's try our best
+        BlockInfo {
+            length,
+            has_forbidden_tokens,
+            top_level_token_span,
+        }
     }
 
     fn is_forbidden_token(&self, token: &Token<'_>) -> bool {
-        token.kind == TokenKind::ReservedTopLevel
-            || token.kind == TokenKind::LineComment
+        token.kind == TokenKind::LineComment
             || token.kind == TokenKind::BlockComment
             || token.value == ";"
-            || if self.newline_on_reserved {
-                token.kind == TokenKind::ReservedNewline
+            || if self.newline_on_reserved_limit == 0 {
+                token.kind == TokenKind::ReservedTopLevel
+                    || token.kind == TokenKind::ReservedNewline
             } else {
                 false
             }
-            || ["case", "end"].contains(&token.value.to_lowercase().as_str())
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -2135,4 +2135,35 @@ from
 
         assert_eq!(format(input, &QueryParams::None, &options), expected);
     }
+
+    #[test]
+    fn it_formats_blocks_inline_or_not() {
+        let input = " UPDATE t SET o = ($5 + $6 + $7 + $8),a = CASE WHEN $2
+            THEN NULL ELSE COALESCE($3, b) END, b = CASE WHEN $4 THEN NULL ELSE
+            COALESCE($5, b) END, s = (SELECT true FROM bar WHERE bar.foo = $99),
+            c = CASE WHEN $6 THEN NULL ELSE COALESCE($7, c) END,
+            d = CASE WHEN $8 THEN NULL ELSE COALESCE($9, d) END,
+            e = (SELECT true FROM bar) WHERE id = $1";
+        let options = FormatOptions {
+            max_inline_arguments: Some(50),
+            max_inline_block: 100,
+            max_inline_top_level: Some(10),
+            ..Default::default()
+        };
+        let expected = indoc!(
+            "
+          UPDATE t
+          SET
+            o = ($5 + $6 + $7 + $8),
+            a = CASE WHEN $2 THEN NULL ELSE COALESCE($3, b) END,
+            b = CASE WHEN $4 THEN NULL ELSE COALESCE($5, b) END,
+            s = (SELECT true FROM bar WHERE bar.foo = $99),
+            c = CASE WHEN $6 THEN NULL ELSE COALESCE($7, c) END,
+            d = CASE WHEN $8 THEN NULL ELSE COALESCE($9, d) END,
+            e = (SELECT true FROM bar)
+          WHERE id = $1"
+        );
+
+        assert_eq!(format(input, &QueryParams::None, &options), expected);
+    }
 }
diff --git a/src/tokenizer.rs b/src/tokenizer.rs
