Support json query filter

sileht · sileht · commit ca5ff71ac985 · 2015-08-20T10:34:29.000+02:00
Related bug: kennknowles#8
diff --git a/jsonpath_rw/jsonpath.py b/jsonpath_rw/jsonpath.py
@@ -4,6 +4,7 @@
 from six.moves import xrange
 from itertools import *
 import functools
+import operator
 
 logger = logging.getLogger(__name__)
 
@@ -14,7 +15,7 @@
 class JSONPath(object):
     """
     The base class for JSONPath abstract syntax; those
-    methods stubbed here are the interface to supported 
+    methods stubbed here are the interface to supported
     JSONPath semantics.
     """
 
@@ -53,8 +54,8 @@ class DatumInContext(object):
     """
     Represents a datum along a path from a context.
 
-    Essentially a zipper but with a structure represented by JsonPath, 
-    and where the context is more of a parent pointer than a proper 
+    Essentially a zipper but with a structure represented by JsonPath,
+    and where the context is more of a parent pointer than a proper
     representation of the context.
 
     For quick-and-dirty work, this proxies any non-special attributes
@@ -115,17 +116,17 @@ class AutoIdForDatum(DatumInContext):
     """
     This behaves like a DatumInContext, but the value is
     always the path leading up to it, not including the "id",
-    and with any "id" fields along the way replacing the prior 
+    and with any "id" fields along the way replacing the prior
     segment of the path
 
     For example, it will make "foo.bar.id" return a datum
     that behaves like DatumInContext(value="foo.bar", path="foo.bar.id").
 
     This is disabled by default; it can be turned on by
     settings the `auto_id_field` global to a value other
-    than `None`. 
+    than `None`.
     """
-    
+
     def __init__(self, datum, id_field=None):
         """
         Invariant is that datum.path is the path from context to datum. The auto id
@@ -212,7 +213,7 @@ class Child(JSONPath):
     JSONPath that first matches the left, then the right.
     Concrete syntax is <left> '.' <right>
     """
-    
+
     def __init__(self, left, right):
         self.left = left
         self.right = right
@@ -222,7 +223,7 @@ def find(self, datum):
         Extra special case: auto ids do not have children,
         so cut it off right now rather than auto id the auto id
         """
-        
+
         return [submatch
                 for subdata in self.left.find(datum)
                 if not isinstance(subdata, AutoIdForDatum)
@@ -256,7 +257,7 @@ def __str__(self):
 
     def __repr__(self):
         return 'Parent()'
-        
+
 
 class Where(JSONPath):
     """
@@ -267,7 +268,7 @@ class Where(JSONPath):
     WARNING: Subject to change. May want to have "contains"
     or some other better word for it.
     """
-    
+
     def __init__(self, left, right):
         self.left = left
         self.right = right
@@ -286,7 +287,7 @@ class Descendants(JSONPath):
     JSONPath that matches first the left expression then any descendant
     of it which matches the right expression.
     """
-    
+
     def __init__(self, left, right):
         self.left = left
         self.right = right
@@ -295,7 +296,7 @@ def find(self, datum):
         # <left> .. <right> ==> <left> . (<right> | *..<right> | [*]..<right>)
         #
         # With with a wonky caveat that since Slice() has funky coercions
-        # we cannot just delegate to that equivalence or we'll hit an 
+        # we cannot just delegate to that equivalence or we'll hit an
         # infinite loop. So right here we implement the coercion-free version.
 
         # Get all left matches into a list
@@ -321,12 +322,12 @@ def match_recursively(datum):
                 recursive_matches = []
 
             return right_matches + list(recursive_matches)
-                
+
         # TODO: repeatable iterator instead of list?
         return [submatch
                 for left_match in left_matches
                 for submatch in match_recursively(left_match)]
-            
+
     def is_singular():
         return False
 
@@ -385,7 +386,7 @@ class Fields(JSONPath):
     WARNING: If '*' is any of the field names, then they will
     all be returned.
     """
-    
+
     def __init__(self, *fields):
         self.fields = fields
 
@@ -411,7 +412,7 @@ def reified_fields(self, datum):
 
     def find(self, datum):
         datum  = DatumInContext.wrap(datum)
-        
+
         return  [field_datum
                  for field_datum in [self.get_field_datum(datum, field) for field in self.reified_fields(datum)]
                  if field_datum is not None]
@@ -429,7 +430,7 @@ def __eq__(self, other):
 class Index(JSONPath):
     """
     JSONPath that matches indices of the current datum, or none if not large enough.
-    Concrete syntax is brackets. 
+    Concrete syntax is brackets.
 
     WARNING: If the datum is not long enough, it will not crash but will not match anything.
     NOTE: For the concrete syntax of `[*]`, the abstract syntax is a Slice() with no parameters (equiv to `[:]`
@@ -440,7 +441,7 @@ def __init__(self, index):
 
     def find(self, datum):
         datum = DatumInContext.wrap(datum)
-        
+
         if len(datum.value) > self.index:
             return [DatumInContext(datum.value[self.index], path=self, context=datum)]
         else:
@@ -454,15 +455,15 @@ def __str__(self):
 
 class Slice(JSONPath):
     """
-    JSONPath matching a slice of an array. 
+    JSONPath matching a slice of an array.
 
     Because of a mismatch between JSON and XML when schema-unaware,
     this always returns an iterable; if the incoming data
     was not a list, then it returns a one element list _containing_ that
     data.
 
     Consider these two docs, and their schema-unaware translation to JSON:
-    
+
     <a><b>hello</b></a> ==> {"a": {"b": "hello"}}
     <a><b>hello</b><b>goodbye</b></a> ==> {"a": {"b": ["hello", "goodbye"]}}
 
@@ -480,10 +481,10 @@ def __init__(self, start=None, end=None, step=None):
         self.start = start
         self.end = end
         self.step = step
-    
+
     def find(self, datum):
         datum = DatumInContext.wrap(datum)
-        
+
         # Here's the hack. If it is a dictionary or some kind of constant,
         # put it in a single-element list
         if (isinstance(datum.value, dict) or isinstance(datum.value, six.integer_types) or isinstance(datum.value, six.string_types)):
@@ -500,7 +501,7 @@ def __str__(self):
         if self.start == None and self.end == None and self.step == None:
             return '[*]'
         else:
-            return '[%s%s%s]' % (self.start or '', 
+            return '[%s%s%s]' % (self.start or '',
                                    ':%d'%self.end if self.end else '',
                                    ':%d'%self.step if self.step else '')
 
@@ -559,3 +560,90 @@ def __repr__(self):
 
     def __str__(self):
         return '[?%s]' % self.expressions
+
+
+OPERATOR_MAP = {
+    '!=': operator.ne,
+    '==': operator.eq,
+    '=': operator.eq,
+    '<=': operator.le,
+    '<': operator.lt,
+    '>=': operator.ge,
+    '>': operator.gt,
+}
+
+
+class Filter(JSONPath):
+    """The JSONQuery filter"""
+
+    def __init__(self, expressions):
+        self.expressions = expressions
+
+    def find(self, datum):
+        if not self.expressions:
+            return []
+
+        datum = DatumInContext.wrap(datum)
+        return [DatumInContext(datum.value[i],
+                               path=Index(i),
+                               context=datum)
+                for i in xrange(0, len(datum.value))
+                if (len(self.expressions) ==
+                    len(list(filter(lambda x: x.find(datum.value[i]),
+                                    self.expressions))))]
+
+    def __repr__(self):
+        return '%s(%r)' % (self.__class__.__name__, self.expressions)
+
+    def __str__(self):
+        return '[?%s]' % self.expressions
+
+
+class FilterExpression(JSONPath):
+    """The JSONQuery expression"""
+
+    def __init__(self, target, op, value):
+        self.target = target
+        self.op = op
+        self.value = value
+
+    def find(self, datum):
+        datum = self.target.find(DatumInContext.wrap(datum))
+
+        if not datum:
+            return []
+        if self.op is None:
+            return datum
+
+        found = []
+        for data in datum:
+            value = data.value
+            if isinstance(self.value, int):
+                try:
+                    value = int(value)
+                except ValueError:
+                    continue
+
+            if OPERATOR_MAP[self.op](value, self.value):
+                found.append(data)
+
+        return found
+
+    def __eq__(self, other):
+        return (isinstance(other, Filter) and
+                self.target == other.target and
+                self.op == other.op and
+                self.value == other.value)
+
+    def __repr__(self):
+        if self.op is None:
+            return '%s(%r)' % (self.__class__.__name__, self.target)
+        else:
+            return '%s(%r %s %r)' % (self.__class__.__name__,
+                                     self.target, self.op, self.value)
+
+    def __str__(self):
+        if self.op is None:
+            return '%s' % self.target
+        else:
+            return '%s %s %s' % (self.target, self.op, self.value)
diff --git a/jsonpath_rw/lexer.py b/jsonpath_rw/lexer.py
@@ -46,11 +46,12 @@ def tokenize(self, string):
     #
     # Anyhow, it is pythonic to give some rope to hang oneself with :-)
 
-    literals = ['*', '.', '[', ']', '(', ')', '$', ',', ':', '|', '&']
+    literals = ['*', '.', '[', ']', '(', ')', '$', ',', ':', '|', '&', '@', '?']
 
     reserved_words = { 'where': 'WHERE' }
 
-    tokens = ['DOUBLEDOT', 'NUMBER', 'ID', 'NAMED_OPERATOR', 'SORT_DIRECTION'] + list(reserved_words.values())
+    tokens = ['DOUBLEDOT', 'NUMBER', 'ID', 'NAMED_OPERATOR', 'SORT_DIRECTION',
+              'FILTER_OP'] + list(reserved_words.values())
 
     states = [ ('singlequote', 'exclusive'),
                ('doublequote', 'exclusive'),
@@ -59,9 +60,10 @@ def tokenize(self, string):
     # Normal lexing, rather easy
     t_DOUBLEDOT = r'\.\.'
     t_ignore = ' \t'
+    t_FILTER_OP = r'(==?|<=|>=|!=|<|>)'
 
     def t_ID(self, t):
-        r'[a-zA-Z_@][a-zA-Z0-9_@\-]*'
+        r'@?[a-zA-Z_][a-zA-Z0-9_@\-]*'
         t.type = self.reserved_words.get(t.value, 'ID')
         return t
 
diff --git a/jsonpath_rw/parser.py b/jsonpath_rw/parser.py
@@ -187,6 +187,41 @@ def p_jsonpath_sort(self, p):
         sort = SortedThis(p[3])
         p[0] = Child(p[1], sort)
 
+    def p_jsonpath_this(self, p):
+        "jsonpath : '@'"
+        p[0] = This()
+
+    def p_expression(self, p):
+        """expression : jsonpath
+                      | jsonpath FILTER_OP ID
+                      | jsonpath FILTER_OP NUMBER
+        """
+        if len(p) == 2:
+            left, op, right = p[1], None, None
+        else:
+            __, left, op, right = p
+        p[0] = FilterExpression(left, op, right)
+
+    def p_expressions_expression(self, p):
+        "expressions : expression"
+        p[0] = [p[1]]
+
+    def p_expressions_and(self, p):
+        "expressions : expressions '&' expressions"
+        p[0] = p[1] + p[3]
+
+    def p_expressions_parens(self, p):
+        "expressions : '(' expressions ')'"
+        p[0] = p[2]
+
+    def p_filter(self, p):
+        "filter : '?' expressions "
+        p[0] = Filter(p[2])
+
+    def p_jsonpath_filter(self, p):
+        "jsonpath : jsonpath '[' filter ']'"
+        p[0] = Child(p[1], p[3])
+
 
 class IteratorToTokenStream(object):
     def __init__(self, iterator):
diff --git a/tests/test_jsonpath.py b/tests/test_jsonpath.py
@@ -127,6 +127,19 @@ def test_sort_value(self):
              [{'cat': {'dog': 2, 'cow': 1}}, {'cat': {'dog': 2, 'bow': 2}}, {'cat': {'dog': 3, 'cow': 2}}, {'cat': {'dog': 1, 'bow': 3}}]),
         ])
 
+    def test_filter_value(self):
+        jsonpath.auto_id_field = None
+        self.check_cases([
+            ('objects[?cow]', {'objects': [{'cow': 'moo'}, {'cat': 'neigh'}]}, [{'cow': 'moo'}]),
+            ('objects[?@.cow]', {'objects': [{'cow': 'moo'}, {'cat': 'neigh'}]}, [{'cow': 'moo'}]),
+            ('objects[?(@.cow)]', {'objects': [{'cow': 'moo'}, {'cat': 'neigh'}]}, [{'cow': 'moo'}]),
+            ('objects[?(@."cow!?cat")]', {'objects': [{'cow!?cat': 'moo'}, {'cat': 'neigh'}]}, [{'cow!?cat': 'moo'}]),
+            ('objects[?cow="moo"]', {'objects': [{'cow': 'moo'}, {'cow': 'neigh'}, {'cat': 'neigh'}]}, [{'cow': 'moo'}]),
+            ('objects[?(@.["cow"]="moo")]', {'objects': [{'cow': 'moo'}, {'cow': 'neigh'}, {'cat': 'neigh'}]}, [{'cow': 'moo'}]),
+            ('objects[?cow=="moo"]', {'objects': [{'cow': 'moo'}, {'cow': 'neigh'}, {'cat': 'neigh'}]}, [{'cow': 'moo'}]),
+            ('objects[?cow>5]', {'objects': [{'cow': 8}, {'cow': 7}, {'cow': 5}, {'cow': 'neigh'}]}, [{'cow': 8}, {'cow': 7}]),
+            ('objects[?cow>5&cat=2]', {'objects': [{'cow': 8, 'cat': 2}, {'cow': 7, 'cat': 2}, {'cow': 2, 'cat': 2}, {'cow': 5, 'cat': 3}, {'cow': 8, 'cat': 3}]}, [{'cow': 8, 'cat': 2}, {'cow': 7, 'cat': 2}]),
+        ])
 
     def test_root_value(self):
         jsonpath.auto_id_field = None
diff --git a/tests/test_lexer.py b/tests/test_lexer.py
