allow_any_combination_of_operators (#57)

* allow_any_combination_of_operators * restore formatting
shssoichiro · Nov 2, 2024 · e4217f2 · e4217f2
1 parent f72f2ee
commit e4217f2
Show file tree

Hide file tree

Showing 2 changed files with 118 additions and 27 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -1759,6 +1759,87 @@ mod tests {
               and colb = 3"
         );
 
+        assert_eq!(format(input, &QueryParams::None, &options), expected);
+    }
+    #[test]
+    fn it_correctly_parses_all_operators() {
+        let operators = [
+            "!!", "!~~*", "!~~", "!~*", "!~", "##", "#>>", "#>", "#-", "&<|", "&<", "&>", "&&",
+            "*<>", "*<=", "*>=", "*>", "*=", "*<", "<<|", "<<=", "<<", "<->", "<@", "<^", "<=",
+            "<>", "<", ">=", ">>=", ">>", ">^", "->>", "->", "-|-", "-", "+", "/", "=", "%", "?||",
+            "?|", "?-|", "?-", "?#", "?&", "?", "@@@", "@@", "@>", "@?", "@-@", "@", "^@", "^",
+            "|&>", "|>>", "|/", "|", "||/", "||", "~>=~", "~>~", "~<=~", "~<~", "~=", "~*", "~~*",
+            "~~", "~",
+        ];
+
+        // Test each operator individually
+        for &operator in &operators {
+            let input = format!("left {} right", operator);
+            let expected = format!("left {} right", operator);
+            let options = FormatOptions {
+                uppercase: None,
+                ..FormatOptions::default()
+            };
+
+            assert_eq!(
+                format(&input, &QueryParams::None, &options),
+                expected,
+                "Failed to parse operator: {}",
+                operator
+            );
+        }
+    }
+    #[test]
+    fn it_correctly_splits_operators() {
+        let input = "
+  SELECT
+  left <@ right,
+  left << right,
+  left >> right,
+  left &< right,
+  left &> right,
+  left -|- right,
+  @@ left,
+  @-@ left,
+  left <-> right,
+  left <<| right,
+  left |>> right,
+  left &<| right,
+  left |>& right,
+  left <^ right,
+  left >^ right,
+  ?- left,
+  left ?-| right,
+  left ?|| right,
+  left ~= right";
+        let options = FormatOptions {
+            uppercase: None,
+            ..FormatOptions::default()
+        };
+        let expected = indoc!(
+            "
+SELECT
+  left <@ right,
+  left << right,
+  left >> right,
+  left &< right,
+  left &> right,
+  left -|- right,
+  @@ left,
+  @-@ left,
+  left <-> right,
+  left <<| right,
+  left |>> right,
+  left &<| right,
+  left |>& right,
+  left <^ right,
+  left >^ right,
+  ?- left,
+  left ?-| right,
+  left ?|| right,
+  left ~= right"
+        );
+
         assert_eq!(format(input, &QueryParams::None, &options), expected);
     }
 }
diff --git a/src/tokenizer.rs b/src/tokenizer.rs
@@ -1,10 +1,10 @@
 use nom::branch::alt;
 use nom::bytes::complete::{tag, tag_no_case, take, take_until, take_while1};
 use nom::character::complete::{anychar, char, digit0, digit1, not_line_ending};
-use nom::combinator::{eof, opt, peek, recognize, verify};
+use nom::combinator::{eof, map, opt, peek, recognize, verify};
 use nom::error::ParseError;
 use nom::error::{Error, ErrorKind};
-use nom::multi::many0;
+use nom::multi::{many0, many_m_n};
 use nom::sequence::{terminated, tuple};
 use nom::{AsChar, Err, IResult};
 use std::borrow::Cow;
@@ -107,9 +107,10 @@ fn get_next_token<'a>(
                 last_reserved_top_level_token,
             )
         })
+        .or_else(|_| get_operator_token(input))
         .or_else(|_| get_placeholder_token(input, named_placeholders))
         .or_else(|_| get_word_token(input))
-        .or_else(|_| get_operator_token(input))
+        .or_else(|_| get_any_other_char(input))
 }
 
 fn get_whitespace_token(input: &str) -> IResult<&str, Token<'_>> {
@@ -998,30 +999,39 @@ fn get_word_token(input: &str) -> IResult<&str, Token<'_>> {
 }
 
 fn get_operator_token(input: &str) -> IResult<&str, Token<'_>> {
-    alt((
-        tag("!="),
-        tag("<>"),
-        tag("=="),
-        tag("<="),
-        tag(">="),
-        tag("!<"),
-        tag("!>"),
-        tag("||"),
-        tag("::"),
-        tag("->>"),
-        tag("->"),
-        tag("~~*"),
-        tag("~~"),
-        tag("!~~*"),
-        tag("!~~"),
-        tag("~*"),
-        tag("!~*"),
-        tag("!~"),
-        tag(":="),
-        recognize(verify(take(1usize), |token: &str| {
-            token != "\n" && token != "\r"
-        })),
-    ))(input)
+    // Define the allowed operator characters
+    let allowed_operators = alt((
+        tag("!"),
+        tag("<"),
+        tag(">"),
+        tag("="),
+        tag("|"),
+        tag(":"),
+        tag("-"),
+        tag("~"),
+        tag("*"),
+        tag("&"),
+        tag("@"),
+        tag("^"),
+        tag("?"),
+        tag("#"),
+        tag("/"),
+    ));
+
+    map(
+        recognize(many_m_n(2, 5, allowed_operators)),
+        |token: &str| Token {
+            kind: TokenKind::Operator,
+            value: token,
+            key: None,
+        },
+    )(input)
+}
+fn get_any_other_char(input: &str) -> IResult<&str, Token<'_>> {
+    alt((recognize(verify(
+        nom::character::complete::anychar,
+        |&token: &char| token != '\n' && token != '\r',
+    )),))(input)
     .map(|(input, token)| {
         (
             input,