diff --git a/Cargo.toml b/Cargo.toml index 582efbf..98fbf9a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,8 @@ categories = ["development-tools"] [dependencies] nom = "7.0.0" unicode_categories = "0.1.1" +once_cell = "1" +regex = "=1.6" [dev-dependencies] criterion = "0.5" diff --git a/src/formatter.rs b/src/formatter.rs index 66126d4..3f57504 100644 --- a/src/formatter.rs +++ b/src/formatter.rs @@ -1,3 +1,4 @@ +use regex::Regex; use std::borrow::Cow; use crate::indentation::Indentation; @@ -6,12 +7,40 @@ use crate::params::Params; use crate::tokenizer::{Token, TokenKind}; use crate::{FormatOptions, QueryParams}; +use once_cell::sync::Lazy; + +static RE: Lazy = Lazy::new(|| Regex::new(r"(?i)^(--|/\*)\s*fmt\s*:\s*(off|on)").unwrap()); + +pub(crate) fn check_fmt_off(s: &str) -> Option { + RE.captures(s)? + .get(2) + .map(|matched| matched.as_str().eq_ignore_ascii_case("off")) +} + pub(crate) fn format(tokens: &[Token<'_>], params: &QueryParams, options: FormatOptions) -> String { let mut formatter = Formatter::new(tokens, params, options); let mut formatted_query = String::new(); + let mut is_fmt_enabled = true; + let mut is_prev_token_fmt_switch = false; for (index, token) in tokens.iter().enumerate() { + if is_prev_token_fmt_switch { + is_prev_token_fmt_switch = false; + continue; + } + if matches!(token.kind, TokenKind::LineComment | TokenKind::BlockComment) { + if let Some(is_fmt_off) = check_fmt_off(token.value) { + is_fmt_enabled = !is_fmt_off; + is_prev_token_fmt_switch = true; + continue; + } + } formatter.index = index; + if !is_fmt_enabled { + formatter.format_no_change(token, &mut formatted_query); + continue; + } + if token.kind == TokenKind::Whitespace { // ignore (we do our own whitespace formatting) } else if token.kind == TokenKind::LineComment { @@ -79,16 +108,7 @@ impl<'a> Formatter<'a> { self.next_token(1).map_or(false, |current_token| { current_token.kind == TokenKind::Whitespace && self.next_token(2).map_or(false, |next_token| { - matches!( - next_token.kind, - TokenKind::Number - | TokenKind::String - | TokenKind::Word - | TokenKind::ReservedTopLevel - | TokenKind::ReservedTopLevelNoIndent - | TokenKind::ReservedNewline - | TokenKind::Reserved - ) + !matches!(next_token.kind, TokenKind::Operator) }) }); @@ -314,4 +334,8 @@ impl<'a> Formatter<'a> { None } } + + fn format_no_change(&self, token: &Token<'_>, query: &mut String) { + query.push_str(token.value); + } } diff --git a/src/lib.rs b/src/lib.rs index ad5b166..e89ddc5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1658,4 +1658,43 @@ mod tests { assert_eq!(format(input, &QueryParams::None, options), expected); } + + #[test] + fn it_recognizes_fmt_off() { + let input = indoc!( + "SELECT * FROM sometable + WHERE + -- comment test here + -- fmt: off + first_key.second_key = 1 + -- json:first_key.second_key = 1 + -- fmt: on + AND + -- fm1t: off + first_key.second_key = 1 + -- json:first_key.second_key = 1 + -- fmt:on" + ); + let options = FormatOptions { + indent: Indent::Spaces(4), + ..Default::default() + }; + let expected = indoc!( + " + SELECT + * + FROM + sometable + WHERE + -- comment test here + first_key.second_key = 1 + -- json:first_key.second_key = 1 + AND + -- fm1t: off + first_key.second_key = 1 + -- json:first_key.second_key = 1" + ); + + assert_eq!(format(input, &QueryParams::None, options), expected); + } }