[parser] Supports arithmetic expressions

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/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:

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'))