Merge pull request #45 from mindsdb/add-operators

Add operators `<<`, `>>`, `:=`, `&`, `|`, `^`

Author: StpMax

mindsdb_sql_parser/__init__.py

@@ -13,7 +13,8 @@ def __init__(self, lexer, parser):
         self.parser = parser
         self.lexer = lexer
 
-    def process(self, error_info):
+    def process(self) -> str:
+        error_info = self.parser.error_info
         self.tokens = [t for t in error_info['tokens'] if t is not None]
         self.bad_token = error_info['bad_token']
         self.expected_tokens = error_info['expected_tokens']
@@ -48,7 +49,11 @@ def error_location(self):
             else:
                 line = line.ljust(token.index)
 
-            line += token.value
+            if token.type == 'VARIABLE':
+                line += f'@{token.value}'
+            else:
+                line += token.value
+
             lines_idx[token.lineno] = line
 
         msgs = []
@@ -175,7 +180,7 @@ def parse_sql(sql, dialect=None):
     if ast is None:
 
         eh = ErrorHandling(lexer, parser)
-        message = eh.process(parser.error_info)
+        message = eh.process()
 
         raise ParsingException(message)
 

mindsdb_sql_parser/lexer.py

@@ -71,8 +71,8 @@ class MindsDBLexer(Lexer):
         LBRACE, RBRACE, LBRACKET, RBRACKET, COLON, SEMICOLON,
 
         # Operators
-        PLUS, MINUS, MATCH, NOT_MATCH, DIVIDE, MODULO,
-        EQUALS, NEQUALS, GREATER, GEQ, LESS, LEQ,
+        PLUS, MINUS, MATCH, NOT_MATCH, DIVIDE, MODULO, BIT_AND, BIT_OR, BIT_XOR,
+        ASSIGN_COLON, RIGHT_SHIFT, LEFT_SHIFT, EQUALS, NEQUALS, GREATER, GEQ, LESS, LEQ,
         AND, OR, NOT, IS, IS_NOT, TYPECAST,
         IN, NOT_IN, LIKE, NOT_LIKE, CONCAT, BETWEEN, WINDOW, OVER, PARTITION_BY,
         VECT_L2, VECT_L1, VECT_INNER, VECT_COS, VECT_HAMM, VECT_JACC,
@@ -272,6 +272,9 @@ class MindsDBLexer(Lexer):
     LPAREN = r'\('
     RPAREN = r'\)'
     PARAMETER = r'\?'
+
+    ASSIGN_COLON = r':='
+
     # json
     LBRACE = r'\{'
     RBRACE = r'\}'
@@ -296,6 +299,10 @@ class MindsDBLexer(Lexer):
     NOT_MATCH = r'!~'
     DIVIDE = r'/'
     MODULO = r'%'
+    BIT_AND = r'&'
+    BIT_XOR = r'\^'
+    RIGHT_SHIFT = r'>>'
+    LEFT_SHIFT = r'<<'
     EQUALS = r'='
     NEQUALS = r'(!=|<>)'
     GEQ = r'>='
@@ -313,6 +320,7 @@ class MindsDBLexer(Lexer):
     IN = r'\bIN\b'
     CAST = r'\bCAST\b'
     CONCAT = r'\|\|'
+    BIT_OR = r'\|'
     BETWEEN = r'\bBETWEEN\b'
     INTERVAL = r'\bINTERVAL\b'
     WINDOW = r'\bWINDOW\b'

mindsdb_sql_parser/parser.py

@@ -46,10 +46,11 @@ class MindsDBParser(Parser):
         ('left', OR),
         ('left', AND),
         ('right', UNOT),
-        ('left', EQUALS, NEQUALS),
-        ('nonassoc', LESS, LEQ, GREATER, GEQ, IN, NOT_IN, BETWEEN, IS, IS_NOT, NOT_LIKE, LIKE),
+        ('left', EQUALS, NEQUALS, ASSIGN_COLON),
+        ('nonassoc', LESS, LEQ, GREATER, GEQ, IN, NOT_IN, BETWEEN, IS, IS_NOT, NOT_LIKE, LIKE, RIGHT_SHIFT, LEFT_SHIFT),
         ('left', JSON_GET),
         ('left', PLUS, MINUS),
+        ('left', BIT_AND, BIT_XOR, BIT_OR),
         ('left', STAR, DIVIDE, TYPECAST, MODULO),
         ('right', UMINUS),  # Unary minus operator, unary not
 
@@ -1631,8 +1632,14 @@ def star(self, p):
        'expr STAR expr',
        'expr DIVIDE expr',
        'expr MODULO expr',
+       'expr BIT_AND expr',
+       'expr BIT_XOR expr',
+       'expr BIT_OR expr',
+       'expr ASSIGN_COLON expr',
        'expr EQUALS expr',
        'expr NEQUALS expr',
+       'expr RIGHT_SHIFT expr',
+       'expr LEFT_SHIFT expr',
        'expr GEQ expr',
        'expr GREATER expr',
        'expr GEQ LAST',
@@ -1657,6 +1664,10 @@ def star(self, p):
        'expr NOT_IN expr',
        'expr IN expr',)
     def expr(self, p):
+        if p[1] == ':=':
+            # allow 'assign' only for variables
+            if not isinstance(p[0], Variable):
+                raise ParsingException("Assignment `:=` allowed only to variables (@var := ...)")
         if hasattr(p, 'LAST'):
             arg1 = Last()
         else:

tests/test_base_sql/test_base_lexer.py

@@ -120,6 +120,12 @@ def test_binary_ops(self):
             ('>=', 'GEQ'),
             ('<', 'LESS'),
             ('<=', 'LEQ'),
+            ('>>', 'RIGHT_SHIFT'),
+            ('<<', 'LEFT_SHIFT'),
+            (':=', 'ASSIGN_COLON'),
+            ('&', 'BIT_AND'),
+            ('|', 'BIT_OR'),
+            ('^', 'BIT_XOR'),
             ('AND', 'AND'),
             ('OR', 'OR'),
             ('IS', 'IS'),

tests/test_base_sql/test_select_operations.py

@@ -6,8 +6,8 @@
 
 class TestOperations:
     def test_select_binary_operations(self):
-        for op in ['+', '-', '/', '*', '%', '=', '!=', '>', '<', '>=', '<=',
-                   'is', 'IS NOT', 'like', 'in', 'and', 'or', '||']:
+        for op in ['+', '-', '/', '*', '%', '=', '!=', '>', '<', '>=', '<=', '>>', '<<',
+                   'is', 'IS NOT', 'like', 'in', 'and', 'or', '||', '&', '|', '^']:
             sql = f'SELECT column1 {op.upper()} column2 FROM tab'
             ast = parse_sql(sql)