Skip to content

Commit

Permalink
feat: More formatting options (#74)
Browse files Browse the repository at this point in the history
* chore: Add pretty_assertions

* feat: provide an option to format to a single line

* feat: Make the inline blocks limit tunable

* feat: Add an option to keep arguments inline

* feat: Add an option to inline over the top level tokens
  • Loading branch information
lu-zero authored Jan 17, 2025
1 parent 211fbb8 commit d49ed44
Show file tree
Hide file tree
Showing 4 changed files with 281 additions and 15 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ winnow = { version = "0.6.20", features = ["simd"] }
[dev-dependencies]
criterion = "0.4"
indoc = "2.0"
pretty_assertions = "1.4.1"

[[bench]]
name = "bench"
Expand Down
76 changes: 67 additions & 9 deletions src/formatter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -90,10 +90,12 @@ pub(crate) fn format(
TokenKind::ReservedTopLevel => {
formatter.format_top_level_reserved_word(token, &mut formatted_query);
formatter.previous_reserved_word = Some(token);
formatter.previous_top_level_reserved_word = Some(token);
}
TokenKind::ReservedTopLevelNoIndent => {
formatter.format_top_level_reserved_word_no_indent(token, &mut formatted_query);
formatter.previous_reserved_word = Some(token);
formatter.previous_top_level_reserved_word = Some(token);
}
TokenKind::ReservedNewline => {
formatter.format_newline_reserved_word(token, &mut formatted_query);
Expand Down Expand Up @@ -140,27 +142,34 @@ pub(crate) fn format(
struct Formatter<'a> {
index: usize,
previous_reserved_word: Option<&'a Token<'a>>,
previous_top_level_reserved_word: Option<&'a Token<'a>>,
tokens: &'a [Token<'a>],
params: Params<'a>,
options: &'a FormatOptions<'a>,
indentation: Indentation<'a>,
inline_block: InlineBlock,
line_start: usize,
}

impl<'a> Formatter<'a> {
fn new(tokens: &'a [Token<'a>], params: &'a QueryParams, options: &'a FormatOptions) -> Self {
Formatter {
index: 0,
previous_reserved_word: None,
previous_top_level_reserved_word: None,
tokens,
params: Params::new(params),
options,
indentation: Indentation::new(options),
inline_block: InlineBlock::new(),
inline_block: InlineBlock::new(
options.max_inline_block,
options.max_inline_arguments.is_none(),
),
line_start: 0,
}
}

fn format_line_comment(&self, token: &Token<'_>, query: &mut String) {
fn format_line_comment(&mut self, token: &Token<'_>, query: &mut String) {
let is_whitespace_followed_by_special_token =
self.next_token(1).map_or(false, |current_token| {
current_token.kind == TokenKind::Whitespace
Expand Down Expand Up @@ -189,7 +198,7 @@ impl<'a> Formatter<'a> {
self.trim_all_spaces_end(query);
query.push_str("::");
}
fn format_block_comment(&self, token: &Token<'_>, query: &mut String) {
fn format_block_comment(&mut self, token: &Token<'_>, query: &mut String) {
self.add_new_line(query);
query.push_str(&self.indent_comment(token.value));
self.add_new_line(query);
Expand All @@ -200,7 +209,16 @@ impl<'a> Formatter<'a> {
self.add_new_line(query);
self.indentation.increase_top_level();
query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value)));
self.add_new_line(query);
let len = self.top_level_tokens_span();
if self
.options
.max_inline_top_level
.map_or(true, |limit| limit < len)
{
self.add_new_line(query);
} else {
query.push(' ');
}
}

fn format_top_level_reserved_word_no_indent(&mut self, token: &Token<'_>, query: &mut String) {
Expand All @@ -210,8 +228,17 @@ impl<'a> Formatter<'a> {
self.add_new_line(query);
}

fn format_newline_reserved_word(&self, token: &Token<'_>, query: &mut String) {
self.add_new_line(query);
fn format_newline_reserved_word(&mut self, token: &Token<'_>, query: &mut String) {
if self
.options
.max_inline_arguments
.map_or(true, |limit| limit < self.line_len_next(query))
{
self.add_new_line(query);
} else {
self.trim_spaces_end(query);
query.push(' ');
}
query.push_str(&self.equalize_whitespace(&self.format_reserved_word(token.value)));
query.push(' ');
}
Expand Down Expand Up @@ -303,7 +330,13 @@ impl<'a> Formatter<'a> {

if self.inline_block.is_active() {
self.inline_block.end();
self.format_with_space_after(&token, query);
if token.value.to_lowercase() == "end" {
self.trim_spaces_end(query);
query.push(' ');
self.format_with_spaces(&token, query);
} else {
self.format_with_space_after(&token, query);
}
} else {
self.indentation.decrease_block_level();
self.add_new_line(query);
Expand All @@ -317,7 +350,7 @@ impl<'a> Formatter<'a> {
}

// Commas start a new line (unless within inline parentheses or SQL "LIMIT" clause)
fn format_comma(&self, token: &Token<'_>, query: &mut String) {
fn format_comma(&mut self, token: &Token<'_>, query: &mut String) {
self.trim_spaces_end(query);
query.push_str(token.value);
query.push(' ');
Expand All @@ -332,6 +365,12 @@ impl<'a> Formatter<'a> {
{
return;
}

if matches!((self.previous_top_level_reserved_word, self.options.max_inline_arguments),
(Some(word), Some(limit)) if word.value.to_lowercase() == "select" && limit > self.line_len_next(query))
{
return;
}
self.add_new_line(query);
}

Expand All @@ -355,11 +394,16 @@ impl<'a> Formatter<'a> {
}
}

fn add_new_line(&self, query: &mut String) {
fn add_new_line(&mut self, query: &mut String) {
self.trim_spaces_end(query);
if self.options.inline {
query.push(' ');
return;
}
if !query.ends_with('\n') {
query.push('\n');
}
self.line_start = query.len();
query.push_str(&self.indentation.get_indent());
}

Expand Down Expand Up @@ -446,6 +490,20 @@ impl<'a> Formatter<'a> {
}
}

fn line_len_next(&self, query: &str) -> usize {
query.len() - self.line_start + self.next_token(1).map_or(0, |t| t.value.len())
}

fn top_level_tokens_span(&self) -> usize {
assert_eq!(self.tokens[self.index].kind, TokenKind::ReservedTopLevel);

self.tokens[self.index..]
.iter()
.skip(1)
.map(|token| token.value.len())
.sum()
}

fn format_no_change(&self, token: &Token<'_>, query: &mut String) {
query.push_str(token.value);
}
Expand Down
31 changes: 25 additions & 6 deletions src/inline_block.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,27 @@ use crate::tokenizer::{Token, TokenKind};

pub(crate) struct InlineBlock {
level: usize,
inline_max_length: usize,
newline_on_reserved: bool,
}

impl Default for InlineBlock {
fn default() -> Self {
InlineBlock {
level: 0,
inline_max_length: 50,
newline_on_reserved: true,
}
}
}

impl InlineBlock {
pub fn new() -> Self {
InlineBlock { level: 0 }
pub fn new(inline_max_length: usize, newline_on_reserved: bool) -> Self {
InlineBlock {
level: 0,
inline_max_length,
newline_on_reserved,
}
}

pub fn begin_if_possible(&mut self, tokens: &[Token<'_>], index: usize) {
Expand All @@ -28,16 +44,14 @@ impl InlineBlock {
}

fn is_inline_block(&self, tokens: &[Token<'_>], index: usize) -> bool {
const INLINE_MAX_LENGTH: usize = 50;

let mut length = 0;
let mut level = 0;

for token in &tokens[index..] {
length += token.value.len();

// Overran max length
if length > INLINE_MAX_LENGTH {
if length > self.inline_max_length {
return false;
}
if token.kind == TokenKind::OpenParen {
Expand All @@ -59,9 +73,14 @@ impl InlineBlock {

fn is_forbidden_token(&self, token: &Token<'_>) -> bool {
token.kind == TokenKind::ReservedTopLevel
|| token.kind == TokenKind::ReservedNewline
|| token.kind == TokenKind::LineComment
|| token.kind == TokenKind::BlockComment
|| token.value == ";"
|| if self.newline_on_reserved {
token.kind == TokenKind::ReservedNewline
} else {
false
}
|| token.value.to_lowercase() == "case"
}
}
Loading

0 comments on commit d49ed44

Please sign in to comment.