[lexer] Supports arithmetic operators

Author: springcomp

jmespath/lexer.py

@@ -21,10 +21,16 @@ class Lexer(object):
         ')': 'rparen',
         '{': 'lbrace',
         '}': 'rbrace',
+        '+': 'plus',
+        '%': 'modulo',
+        u'\u2212': 'minus',
+        u'\u00d7': 'multiply',
+        u'\u00f7': 'divide',
     }
 
     def tokenize(self, expression):
         self._initialize_for_expression(expression)
+        print(self._current in self.SIMPLE_TOKENS)
         while self._current is not None:
             if self._current in self.SIMPLE_TOKENS:
                 yield {'type': self.SIMPLE_TOKENS[self._current],
@@ -68,16 +74,30 @@ def tokenize(self, expression):
                 yield {'type': 'number', 'value': int(buff),
                        'start': start, 'end': start + len(buff)}
             elif self._current == '-':
-                # Negative number.
-                start = self._position
-                buff = self._consume_number()
-                if len(buff) > 1:
-                    yield {'type': 'number', 'value': int(buff),
-                           'start': start, 'end': start + len(buff)}
+                if not self._peek_is_next_digit():
+                    self._next()
+                    yield {'type': 'minus', 'value': '-',
+                       'start': self._position - 1, 'end': self._position}
+                else:
+                    # Negative number.
+                    start = self._position
+                    buff = self._consume_number()
+                    if len(buff) > 1:
+                        yield {'type': 'number', 'value': int(buff),
+                            'start': start, 'end': start + len(buff)}
+                    else:
+                        raise LexerError(lexer_position=start,
+                                        lexer_value=buff,
+                                        message="Unknown token '%s'" % buff)
+            elif self._current == '/':
+                self._next()
+                if self._current == '/':
+                    self._next()
+                    yield {'type': 'div', 'value': '//',
+                        'start': self._position - 1, 'end': self._position}
                 else:
-                    raise LexerError(lexer_position=start,
-                                     lexer_value=buff,
-                                     message="Unknown token '%s'" % buff)
+                    yield {'type': 'divide', 'value': '/',
+                        'start': self._position, 'end': self._position + 1}
             elif self._current == '"':
                 yield self._consume_quoted_identifier()
             elif self._current == '<':
@@ -117,6 +137,13 @@ def _consume_number(self):
             buff += self._current
         return buff
 
+    def _peek_is_next_digit(self):
+        if (self._position == self._length - 1):
+            return False
+        else:
+            next = self._chars[self._position + 1]
+            return next in self.VALID_NUMBER
+
     def _initialize_for_expression(self, expression):
         if not expression:
             raise EmptyExpressionError()

tests/test_lexer.py

@@ -45,6 +45,50 @@ def test_negative_number(self):
         self.assert_tokens(tokens, [{'type': 'number',
                                      'value': -24}])
 
+    def test_plus(self):
+        tokens = list(self.lexer.tokenize('+'))
+        self.assert_tokens(tokens, [{'type': 'plus',
+                                     'value': '+'}])
+
+    def test_minus(self):
+        tokens = list(self.lexer.tokenize('-'))
+        self.assert_tokens(tokens, [{'type': 'minus',
+                                     'value': '-'}])
+    def test_minus_unicode(self):
+        tokens = list(self.lexer.tokenize(u'\u2212'))
+        self.assert_tokens(tokens, [{'type': 'minus',
+                                     'value': u'\u2212'}])
+
+    def test_multiplication(self):
+        tokens = list(self.lexer.tokenize('*'))
+        self.assert_tokens(tokens, [{'type': 'star',
+                                     'value': '*'}])
+
+    def test_multiplication_unicode(self):
+        tokens = list(self.lexer.tokenize(u'\u00d7'))
+        self.assert_tokens(tokens, [{'type': 'multiply',
+                                     'value': u'\u00d7'}])
+
+    def test_division(self):
+        tokens = list(self.lexer.tokenize('/'))
+        self.assert_tokens(tokens, [{'type': 'divide',
+                                     'value': '/'}])
+
+    def test_division_unicode(self):
+        tokens = list(self.lexer.tokenize('÷'))
+        self.assert_tokens(tokens, [{'type': 'divide',
+                                     'value': '÷'}])
+
+    def test_modulo(self):
+        tokens = list(self.lexer.tokenize('%'))
+        self.assert_tokens(tokens, [{'type': 'modulo',
+                                     'value': '%'}])
+
+    def test_integer_division(self):
+        tokens = list(self.lexer.tokenize('//'))
+        self.assert_tokens(tokens, [{'type': 'div',
+                                     'value': '//'}])
+
     def test_quoted_identifier(self):
         tokens = list(self.lexer.tokenize('"foobar"'))
         self.assert_tokens(tokens, [{'type': 'quoted_identifier',
@@ -151,9 +195,17 @@ def test_bad_first_character(self):
         with self.assertRaises(LexerError):
             tokens = list(self.lexer.tokenize('^foo[0]'))
 
-    def test_unknown_character_with_identifier(self):
-        with self.assertRaisesRegex(LexerError, "Unknown token"):
-            list(self.lexer.tokenize('foo-bar'))
+    def test_arithmetic_expression(self):
+        tokens = list(self.lexer.tokenize('foo-bar'))
+        self.assertEqual(
+            tokens,
+            [
+                {'type': 'unquoted_identifier', 'value': 'foo', 'start': 0, 'end': 3},
+                {'type': 'minus', 'value': '-', 'start': 3, 'end': 4},
+                {'type': 'unquoted_identifier', 'value': 'bar', 'start': 4, 'end': 7},
+                {'type': 'eof', 'value': '', 'start': 7, 'end': 7}
+            ]
+        )
 
 
 if __name__ == '__main__':