Merge pull request #7 from jmespath-community/jep/lexical-scoping

JEP-11 Lexical Scoping

Author: springcomp

jmespath/ast.py

@@ -18,6 +18,10 @@ def current_node():
     return {'type': 'current', 'children': []}
 
 
+def root_node():
+    return {'type': 'root', 'children': []}
+
+
 def expref(expression):
     return {'type': 'expref', 'children': [expression]}
 

jmespath/functions.py

@@ -1,5 +1,7 @@
 import math
 import json
+import inspect
+from pydoc import resolve
 
 from collections import OrderedDict
 
@@ -73,7 +75,7 @@ class Functions(metaclass=FunctionRegistry):
     FUNCTION_TABLE = {
     }
 
-    def call_function(self, function_name, resolved_args):
+    def call_function(self, function_name, resolved_args, *args, **kwargs):
         try:
             spec = self.FUNCTION_TABLE[function_name]
         except KeyError:
@@ -82,6 +84,13 @@ def call_function(self, function_name, resolved_args):
         function = spec['function']
         signature = spec['signature']
         self._validate_arguments(resolved_args, signature, function_name)
+
+        # supply extra arguments only if the function expects them
+
+        parameters = [parameter.name for parameter in inspect.signature(function).parameters.values()]
+        if ('kwargs' in parameters):
+            return function(self, *resolved_args, *args, scopes = kwargs.get('scopes'))
+
         return function(self, *resolved_args)
 
     def _validate_arguments(self, args, signature, function_name):
@@ -531,6 +540,12 @@ def _func_max_by(self, array, expref):
         else:
             return None
 
+    @signature({'types': ['object']}, {'types': ['expref']})
+    def _func_let(self, scope, expref, *args, **kwargs):
+        if 'scopes' in kwargs:
+            kwargs.get('scopes').pushScope(scope)
+        return expref.visit(expref.expression, expref.context, *args, **kwargs)
+
     @signature({'types': ['array'], 'variadic': True})
     def _func_zip(self, *arguments):
         return list(map(list, zip(*arguments)))

jmespath/lexer.py

@@ -17,6 +17,7 @@ class Lexer(object):
         ',': 'comma',
         ':': 'colon',
         '@': 'current',
+        '$': 'root',
         '(': 'lparen',
         ')': 'rparen',
         '{': 'lbrace',

jmespath/parser.py

@@ -27,6 +27,7 @@
 """
 import random
 
+
 from jmespath import lexer
 from jmespath.compat import with_repr_method
 from jmespath import ast
@@ -46,6 +47,7 @@ class Parser(object):
         'rbrace': 0,
         'number': 0,
         'current': 0,
+        'root': 0,
         'expref': 0,
         'colon': 0,
         'pipe': 1,
@@ -251,6 +253,9 @@ def _parse_slice_expression(self):
     def _token_nud_current(self, token):
         return ast.current_node()
 
+    def _token_nud_root(self, token):
+        return ast.root_node()
+
     def _token_nud_expref(self, token):
         expression = self._expression(self.BINDING_POWER['expref'])
         return ast.expref(expression)
@@ -546,9 +551,8 @@ def __init__(self, expression, parsed):
         self.parsed = parsed
 
     def search(self, value, options=None):
-        interpreter = visitor.TreeInterpreter(options)
-        result = interpreter.visit(self.parsed, value)
-        return result
+        evaluator = visitor.ScopedInterpreter(options)
+        return evaluator.evaluate(self.parsed, value)
 
     def _render_dot_file(self):
         """Render the parsed AST as a dot file.
@@ -565,4 +569,4 @@ def _render_dot_file(self):
         return contents
 
     def __repr__(self):
-        return repr(self.parsed)
+        return repr(self.parsed)

jmespath/visitor.py

@@ -2,6 +2,7 @@
 
 from jmespath import functions
 from jmespath.compat import string_type
+from jmespath.compat import with_str_method
 from numbers import Number
 
 
@@ -72,14 +73,14 @@ def __init__(self, dict_cls=None, custom_functions=None):
 
 
 class _Expression(object):
-    def __init__(self, expression, interpreter):
+    def __init__(self, expression, interpreter, context):
         self.expression = expression
         self.interpreter = interpreter
+        self.context = context
 
     def visit(self, node, *args, **kwargs):
         return self.interpreter.visit(node, *args, **kwargs)
 
-
 class Visitor(object):
     def __init__(self):
         self._method_cache = {}
@@ -96,7 +97,6 @@ def visit(self, node, *args, **kwargs):
     def default_visit(self, node, *args, **kwargs):
         raise NotImplementedError("default_visit")
 
-
 class TreeInterpreter(Visitor):
     COMPARATOR_FUNC = {
         'eq': _equals,
@@ -145,11 +145,24 @@ def visit_subexpression(self, node, value):
                 return None
         return result
 
-    def visit_field(self, node, value):
+    def visit_field(self, node, value, *args, **kwargs):
+        identifier = node['value']
+        scopes = kwargs.get('scopes')
+
         try:
-            return value.get(node['value'])
+            result = value.get(identifier) 
+            if result == None:
+                result = self._get_from_scopes(
+                    identifier, *args, scopes=scopes)
+            return result
         except AttributeError:
-            return None
+            return self._get_from_scopes(
+                identifier, *args, scopes=scopes)
+
+    def _get_from_scopes(self, identifier, *args, **kwargs):
+        if 'scopes' in kwargs:
+            return kwargs['scopes'].getValue(identifier) 
+        return None
 
     def visit_comparator(self, node, value):
         # Common case: comparator is == or !=
@@ -187,15 +200,20 @@ def visit_arithmetic(self, node, value):
     def visit_current(self, node, value):
         return value
 
+    def visit_root(self, *args, **kwargs):
+        if 'scopes' in kwargs:
+            return kwargs['scopes'].getValue('$')
+        return None
+
     def visit_expref(self, node, value):
-        return _Expression(node['children'][0], self)
+        return _Expression(node['children'][0], self, value)
 
-    def visit_function_expression(self, node, value):
+    def visit_function_expression(self, node, value, *args, **kwargs):
         resolved_args = []
         for child in node['children']:
             current = self.visit(child, value)
             resolved_args.append(current)
-        return self._functions.call_function(node['value'], resolved_args)
+        return self._functions.call_function(node['value'], resolved_args, scopes = kwargs.get('scopes'))
 
     def visit_filter_projection(self, node, value):
         base = self.visit(node['children'][0], value)
@@ -366,3 +384,45 @@ def _visit(self, node, current):
             self._count += 1
             self._lines.append('  %s -> %s' % (current, child_name))
             self._visit(child, child_name)
+
+
+@with_str_method
+class Scopes:
+    def __init__(self):
+        self._scopes = []
+
+    def pushScope(self, scope):
+        self._scopes.append(scope)
+
+    def popScope(self):
+        if len(self._scopes) > 0:
+            self._scopes.pop()
+
+    def getValue(self, identifier):
+        for scope in self._scopes[::-1]:
+            if scope.get(identifier) != None:
+                return scope[identifier]
+        return None
+
+    def __str__(self):
+        return '{}'.format(self._scopes)
+
+
+class ScopedInterpreter(TreeInterpreter):
+    def __init__(self, options = None):
+        super().__init__(options)
+        self._scopes = Scopes()
+
+    def evaluate(self, ast, root_scope):
+        self._scopes.pushScope({'$': root_scope})
+        return self.visit(ast, root_scope)
+
+    def visit(self, node, *args, **kwargs):
+        scoped_types = ['field', 'function_expression', 'root']
+        if (node['type'] in scoped_types):
+            kwargs.update({'scopes': self._scopes})
+        else:
+            if 'scopes' in kwargs:
+                kwargs.pop('scopes')
+
+        return super().visit(node, *args, **kwargs)