add include function

Author: wangjild

jsonpath_rw/jsonpath.py

@@ -267,6 +267,11 @@ def exclude(self, data):
             self.right.exclude(datum.value)
         return data
 
+    def include(self, data):
+        for datum in self.left.find(data):
+            self.right.include(datum.value)
+        return data
+
     def __eq__(self, other):
         return isinstance(other, Child) and self.left == other.left and self.right == other.right
 
@@ -524,6 +529,23 @@ def exclude(self, data):
                 del data[field]
         return data
 
+    def include(self, data):
+        datum = DatumInContext.wrap(data)
+
+        try:
+            all_fields = tuple(datum.value.keys())
+        except AttributeError:
+            all_fields = ()
+
+        path_fields = self.reified_fields(datum)
+        remove_fields = set(all_fields) - set(path_fields)
+
+        for field in remove_fields:
+            if field in data:
+                del data[field]
+
+        return data
+
     def __str__(self):
         return ','.join(map(str, self.fields))
 
@@ -564,6 +586,15 @@ def exclude(self, data):
             del data[self.index]
         return data
 
+    def include(self, data):
+        if data is None:
+            return None
+
+        if len(data) > self.index:
+            return [data[self.index]]
+
+        return []
+
     def __eq__(self, other):
         return isinstance(other, Index) and self.index == other.index
 
@@ -625,6 +656,18 @@ def exclude(self, data):
 
         return data
 
+    def include(self, data):
+
+        if not data:
+            return data
+
+        ret = []
+        for datum in self.find(data):
+            ret.append(datum.value)
+
+        data = ret
+        return data
+
     def __str__(self):
         if self.start == None and self.end == None and self.step == None:
             return '[*]'

tests/test_jsonpath.py

@@ -463,7 +463,88 @@ def test_exclude_slice(self):
 
     def check_include_cases(self, test_cases):
         for original, string, expected in test_cases:
-            logging.debug('parse("%s").exclude(%s) =?= %s' % (string, original, expected))
+            logging.debug('parse("%s").include(%s) =?= %s' % (string, original, expected))
             actual = parse(string).include(original)
+            print ('actual: ', actual, 'expected: ', expected)
             assert actual == expected
 
+    def test_include_fields(self):
+        self.check_include_cases([
+            ({'foo': 'baz'}, 'foo', {'foo': 'baz'}),
+            ({'foo': 1, 'baz': 2}, 'foo', {'foo': 1}),
+            ({'foo': 1, 'baz': 2}, 'foo,baz', {'foo': 1, 'baz': 2}),
+            ({'@foo': 1}, '@foo', {'@foo': 1}),
+            ({'@foo': 1, 'baz': 2}, '@foo', {'@foo': 1}),
+            ({'foo': 1, 'baz': 2}, '*', {'foo': 1, 'baz': 2}),
+        ])
+
+    def test_include_index(self):
+        self.check_include_cases([
+            ([42], '[0]', [42]),
+            ([42], '[5]', []),
+            ([34, 65, 29, 59], '[2]', [29]),
+            (None, '[0]', None),
+            ([], '[0]', []),
+            (['foo', 'bar', 'baz'], '[0]', ['foo']),
+        ])
+
+    def test_include_slice(self):
+        self.check_include_cases([
+            (['foo', 'bar', 'baz'], '[0:2]', ['foo', 'bar']),
+            (['foo', 'bar', 'baz'], '[0:1]', ['foo']),
+            (['foo', 'bar', 'baz'], '[0:]', ['foo', 'bar', 'baz']),
+            (['foo', 'bar', 'baz'], '[:2]', ['foo', 'bar']),
+            (['foo', 'bar', 'baz'], '[:3]', ['foo', 'bar', 'baz']),
+            (['foo', 'bar', 'baz'], '[0:0]', []),
+        ])
+
+    def test_include_root(self):
+        self.check_include_cases([
+            ('foo', '$', 'foo'),
+            ({}, '$', {}),
+            ({'foo': 1}, '$', {'foo': 1})
+        ])
+
+    def test_include_this(self):
+        self.check_include_cases([
+            ('foo', '`this`', 'foo'),
+            ({}, '`this`', {}),
+            ({'foo': 1}, '`this`', {'foo': 1}),
+            # TODO: fixme
+            #({'foo': 1}, 'foo.`this`', {}),
+            ({'foo': {'bar': 1}}, 'foo.`this`.bar', {'foo': {'bar': 1}}),
+            ({'foo': {'bar': 1, 'baz': 2}}, 'foo.`this`.bar', {'foo': {'bar': 1}})
+        ])
+
+    def test_include_child(self):
+        self.check_include_cases([
+            ({'foo': 'bar'}, '$.foo', {'foo': 'bar'}),
+            ({'foo': 'bar'}, 'foo', {'foo': 'bar'}),
+            ({'foo': {'bar': 1}}, 'foo.bar', {'foo': {'bar': 1}}),
+            ({'foo': {'bar': 1}}, 'foo.$.foo.bar', {'foo': {'bar': 1}}),
+            ({'foo': {'bar': 1, 'baz': 2}}, 'foo.$.foo.bar', {'foo': {'bar': 1}}),
+            ({'foo': {'bar': 1, 'baz': 2}}, '*', {'foo': {'bar': 1, 'baz': 2}}),
+            ({'foo': {'bar': 1, 'baz': 2}}, 'non', {}),
+        ])
+
+    """
+    def test_include_where(self):
+        self.check_include_cases([
+            #({'foo': {'bar': {'baz': 1}}, 'bar': {'baz': 2}},
+            # '*.bar where none', {}),
+
+            ({'foo': {'bar': {'baz': 1}}, 'bar': {'baz': 2}},
+             '*.bar where baz', {'foo': {'bar': {'baz': 1}}})
+        ])
+    """
+
+    """
+    def test_include_descendants(self):
+        self.check_include_cases([
+            ({'somefield': 1}, '$..somefield', {'somefield': 1}),
+            ({'outer': {'nestedfield': 1}}, '$..nestedfield', {'outer': {'nestedfield': 1}}),
+            ({'outs': {'bar': 1, 'ins': {'bar': 9}}, 'outs2': {'bar': 2}},
+             '$..bar',
+             {'outs': {'bar': 1, 'ins': {'bar': 9}}, 'outs2': {'bar': 2}})
+        ])
+    """