Merge pull request #4 from jmespath-community/jep/arithmetic-expressions

JEP-16 Arithmetic Expressions

Author: springcomp

jmespath/ast.py

@@ -2,6 +2,14 @@
 # {"type": <node type>", children: [], "value": ""}
 
 
+def arithmetic_unary(operator, expression):
+    return {'type': 'arithmetic_unary', 'children': [expression], 'value': operator}
+
+
+def arithmetic(operator, left, right):
+    return {'type': 'arithmetic', 'children': [left, right], 'value': operator}
+
+
 def comparator(name, first, second):
     return {'type': 'comparator', 'children': [first, second], 'value': name}
 

jmespath/lexer.py

@@ -21,6 +21,11 @@ class Lexer(object):
         ')': 'rparen',
         '{': 'lbrace',
         '}': 'rbrace',
+        '+': 'plus',
+        '%': 'modulo',
+        u'\u2212': 'minus',
+        u'\u00d7': 'multiply',
+        u'\u00f7': 'divide',
     }
 
     def tokenize(self, expression):
@@ -68,16 +73,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 +136,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()

jmespath/parser.py

@@ -57,6 +57,12 @@ class Parser(object):
         'gte': 5,
         'lte': 5,
         'ne': 5,
+        'minus': 6,
+        'plus': 6,
+        'div': 7,
+        'divide': 7,
+        'modulo': 7,
+        'multiply': 7,
         'flatten': 9,
         # Everything above stops a projection.
         'star': 20,
@@ -170,6 +176,12 @@ def _token_nud_lparen(self, token):
         self._match('rparen')
         return expression
 
+    def _token_nud_minus(self, token):
+        return self._parse_arithmetic_unary(token)
+
+    def _token_nud_plus(self, token):
+        return self._parse_arithmetic_unary(token)
+
     def _token_nud_flatten(self, token):
         left = ast.flatten(ast.identity())
         right = self._parse_projection_rhs(
@@ -318,6 +330,27 @@ def _token_led_lt(self, left):
     def _token_led_lte(self, left):
         return self._parse_comparator(left, 'lte')
 
+    def _token_led_div(self, left):
+        return self._parse_arithmetic(left, 'div')
+
+    def _token_led_divide(self, left):
+        return self._parse_arithmetic(left, 'divide')
+
+    def _token_led_minus(self, left):
+        return self._parse_arithmetic(left, 'minus')
+
+    def _token_led_modulo(self, left):
+        return self._parse_arithmetic(left, 'modulo')
+
+    def _token_led_multiply(self, left):
+        return self._parse_arithmetic(left, 'multiply')
+
+    def _token_led_plus(self, left):
+        return self._parse_arithmetic(left, 'plus')
+
+    def _token_led_star(self, left):
+        return self._parse_arithmetic(left, 'multiply')
+
     def _token_led_flatten(self, left):
         left = ast.flatten(left)
         right = self._parse_projection_rhs(
@@ -356,6 +389,14 @@ def _parse_comparator(self, left, comparator):
         right = self._expression(self.BINDING_POWER[comparator])
         return ast.comparator(comparator, left, right)
 
+    def _parse_arithmetic_unary(self, token):
+        expression = self._expression(self.BINDING_POWER[token['type']])
+        return ast.arithmetic_unary(token['type'], expression)
+
+    def _parse_arithmetic(self, left, operator):
+        right = self._expression(self.BINDING_POWER[operator])
+        return ast.arithmetic(operator, left, right)
+
     def _parse_multi_select_list(self):
         expressions = []
         while True:

jmespath/visitor.py

@@ -107,6 +107,18 @@ class TreeInterpreter(Visitor):
         'gte': operator.ge
     }
     _EQUALITY_OPS = ['eq', 'ne']
+    _ARITHMETIC_UNARY_FUNC = {
+        'minus': operator.neg,
+        'plus': lambda x: x
+    }
+    _ARITHMETIC_FUNC = {
+        'div': operator.floordiv,
+        'divide': operator.truediv,
+        'minus': operator.sub,
+        'modulo': operator.mod,
+        'multiply': operator.mul,
+        'plus': operator.add,
+    }
     MAP_TYPE = dict
 
     def __init__(self, options=None):
@@ -159,6 +171,19 @@ def visit_comparator(self, node, value):
                 return None
             return comparator_func(left, right)
 
+    def visit_arithmetic_unary(self, node, value):
+        operation = self._ARITHMETIC_UNARY_FUNC[node['value']]
+        return operation(
+            self.visit(node['children'][0], value)
+        )
+
+    def visit_arithmetic(self, node, value):
+        operation = self._ARITHMETIC_FUNC[node['value']]
+        return operation(
+            self.visit(node['children'][0], value),
+            self.visit(node['children'][1], value)
+        )
+
     def visit_current(self, node, value):
         return value
 

tests/compliance/arithmetic.json

@@ -0,0 +1,62 @@
+[
+  {
+    "given": {
+      "a": {
+        "b": 1
+      },
+      "c": {
+        "d": 2
+      }
+    },
+    "cases": [
+      {
+        "expression": "`1` + `2`",
+        "result": 3.0
+      },
+      {
+        "expression": "`1` - `2`",
+        "result": -1.0
+      },
+      {
+        "expression": "`2` * `4`",
+        "result": 8.0
+      },
+      {
+        "expression": "`2` × `4`",
+        "result": 8.0
+      },
+      {
+        "expression": "`2` / `3`",
+        "result": 0.6666666666666666
+      },
+      {
+        "expression": "`2` ÷ `3`",
+        "result": 0.6666666666666666
+      },
+      {
+        "expression": "`10` % `3`",
+        "result": 1.0
+      },
+      {
+        "expression": "`10` // `3`",
+        "result": 3.0
+      },
+      {
+        "expression": "-`1` - + `2`",
+        "result": -3.0
+      },
+      {
+        "expression": "{ ab: a.b, cd: c.d } | ab + cd",
+        "result": 3.0
+      },
+      {
+        "expression": "{ ab: a.b, cd: c.d } | ab + cd × cd",
+        "result": 5.0
+      },
+      {
+        "expression": "{ ab: a.b, cd: c.d } | (ab + cd) × cd",
+        "result": 6.0
+      }
+    ]
+  }
+]

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__':

tests/test_parser.py

@@ -3,6 +3,7 @@
 import random
 import string
 import threading
+from jmespath.ast import arithmetic
 from tests import unittest, OrderedDict
 
 from jmespath import parser
@@ -77,6 +78,55 @@ def test_or_repr(self):
         self.assert_parsed_ast('foo || bar', ast.or_expression(ast.field('foo'),
                                                                ast.field('bar')))
 
+    def test_arithmetic_expressions(self):
+        operations = {
+            '+': 'plus',
+            '-': 'minus',
+            '//': 'div',
+            '/': 'divide',
+            '%': 'modulo',
+            u'\u2212': 'minus',
+            u'\u00d7': 'multiply',
+            u'\u00f7': 'divide',
+        }
+        for sign in operations:
+            operation = operations[sign]
+            expression = 'foo {} bar'.format(sign)
+            print(expression)
+            self.assert_parsed_ast(
+                expression,
+                ast.arithmetic(
+                    operation,
+                    ast.field('foo'),
+                    ast.field('bar')
+                ))
+
+    def test_arithmetic_unary(self):
+        operations = {
+            '+': 'plus',
+            '-': 'minus',
+            u'\u2212': 'minus',
+        }
+        for sign in operations:
+            operation = operations[sign]
+            expression = '{} foo'.format(sign)
+            print(expression)
+            self.assert_parsed_ast(
+                expression,
+                ast.arithmetic_unary(
+                    operation,
+                    ast.field('foo'),
+                ))
+
+    def test_arithmetic_multiplication(self):
+        self.assert_parsed_ast(
+            'foo * bar',
+            ast.arithmetic(
+                'multiply',
+                ast.field('foo'),
+                ast.field('bar')
+            ))
+
     def test_unicode_literals_escaped(self):
         self.assert_parsed_ast(r'`"\u2713"`', ast.literal(u'\u2713'))