Merge pull request #72 from kaapstorm/find_or_create

JSONPath.update_or_create()

jsonpath_ng/jsonpath.py

@@ -3,6 +3,7 @@
 import six
 from six.moves import xrange
 from itertools import *  # noqa
+from .exceptions import JSONPathError
 
 # Get logger name
 logger = logging.getLogger(__name__)
@@ -11,6 +12,9 @@
 # ... could be a kwarg pervasively but uses are rare and simple today
 auto_id_field = None
 
+NOT_SET = object()
+LIST_KEY = object()
+
 
 class JSONPath(object):
     """
@@ -27,6 +31,9 @@ def find(self, data):
         """
         raise NotImplementedError()
 
+    def find_or_create(self, data):
+        return self.find(data)
+
     def update(self, data, val):
         """
         Returns `data` with the specified path replaced by `val`. Only updates
@@ -35,6 +42,9 @@ def update(self, data, val):
 
         raise NotImplementedError()
 
+    def update_or_create(self, data, val):
+        return self.update(data, val)
+
     def filter(self, fn, data):
         """
         Returns `data` with the specified path filtering nodes according
@@ -261,6 +271,23 @@ def update(self, data, val):
             self.right.update(datum.value, val)
         return data
 
+    def find_or_create(self, datum):
+        datum = DatumInContext.wrap(datum)
+        submatches = []
+        for subdata in self.left.find_or_create(datum):
+            if isinstance(subdata, AutoIdForDatum):
+                # Extra special case: auto ids do not have children,
+                # so cut it off right now rather than auto id the auto id
+                continue
+            for submatch in self.right.find_or_create(subdata):
+                submatches.append(submatch)
+        return submatches
+
+    def update_or_create(self, data, val):
+        for datum in self.left.find_or_create(data):
+            self.right.update_or_create(datum.value, val)
+        return _clean_list_keys(data)
+
     def filter(self, fn, data):
         for datum in self.left.find(data):
             self.right.filter(fn, datum.value)
@@ -497,15 +524,20 @@ class Fields(JSONPath):
     def __init__(self, *fields):
         self.fields = fields
 
-    def get_field_datum(self, datum, field):
+    @staticmethod
+    def get_field_datum(datum, field, create):
         if field == auto_id_field:
             return AutoIdForDatum(datum)
-        else:
-            try:
-                field_value = datum.value[field] # Do NOT use `val.get(field)` since that confuses None as a value and None due to `get`
-                return DatumInContext(value=field_value, path=Fields(field), context=datum)
-            except (TypeError, KeyError, AttributeError):
-                return None
+        try:
+            field_value = datum.value.get(field, NOT_SET)
+            if field_value is NOT_SET:
+                if create:
+                    datum.value[field] = field_value = {}
+                else:
+                    return None
+            return DatumInContext(field_value, path=Fields(field), context=datum)
+        except (TypeError, AttributeError):
+            return None
 
     def reified_fields(self, datum):
         if '*' not in self.fields:
@@ -518,15 +550,28 @@ def reified_fields(self, datum):
                 return ()
 
     def find(self, datum):
-        datum  = DatumInContext.wrap(datum)
+        return self._find_base(datum, create=False)
 
-        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]
+    def find_or_create(self, datum):
+        return self._find_base(datum, create=True)
+
+    def _find_base(self, datum, create):
+        datum = DatumInContext.wrap(datum)
+        field_data = [self.get_field_datum(datum, field, create)
+                      for field in self.reified_fields(datum)]
+        return [fd for fd in field_data if fd is not None]
 
     def update(self, data, val):
+        return self._update_base(data, val, create=False)
+
+    def update_or_create(self, data, val):
+        return self._update_base(data, val, create=True)
+
+    def _update_base(self, data, val, create):
         if data is not None:
             for field in self.reified_fields(DatumInContext.wrap(data)):
+                if field not in data and create:
+                    data[field] = {}
                 if field in data:
                     if hasattr(val, '__call__'):
                         val(data[field], data, field)
@@ -565,14 +610,33 @@ def __init__(self, index):
         self.index = index
 
     def find(self, datum):
-        datum = DatumInContext.wrap(datum)
+        return self._find_base(datum, create=False)
+
+    def find_or_create(self, datum):
+        return self._find_base(datum, create=True)
 
+    def _find_base(self, datum, create):
+        datum = DatumInContext.wrap(datum)
+        if create:
+            if datum.value == {}:
+                datum.value = _create_list_key(datum.value)
+            self._pad_value(datum.value)
         if datum.value and len(datum.value) > self.index:
             return [DatumInContext(datum.value[self.index], path=self, context=datum)]
         else:
             return []
 
     def update(self, data, val):
+        return self._update_base(data, val, create=False)
+
+    def update_or_create(self, data, val):
+        return self._update_base(data, val, create=True)
+
+    def _update_base(self, data, val, create):
+        if create:
+            if data == {}:
+                data = _create_list_key(data)
+            self._pad_value(data)
         if hasattr(val, '__call__'):
             val.__call__(data[self.index], data, self.index)
         elif len(data) > self.index:
@@ -590,6 +654,14 @@ def __eq__(self, other):
     def __str__(self):
         return '[%i]' % self.index
 
+    def __repr__(self):
+        return '%s(index=%r)' % (self.__class__.__name__, self.index)
+
+    def _pad_value(self, value):
+        if len(value) <= self.index:
+            pad = self.index - len(value) + 1
+            value += [{} for __ in range(pad)]
+
 
 class Slice(JSONPath):
     """
@@ -668,3 +740,32 @@ def __repr__(self):
 
     def __eq__(self, other):
         return isinstance(other, Slice) and other.start == self.start and self.end == other.end and other.step == self.step
+
+
+def _create_list_key(dict_):
+    """
+    Adds a list to a dictionary by reference and returns the list.
+
+    See `_clean_list_keys()`
+    """
+    dict_[LIST_KEY] = new_list = [{}]
+    return new_list
+
+
+def _clean_list_keys(dict_):
+    """
+    Replace {LIST_KEY: ['foo', 'bar']} with ['foo', 'bar'].
+
+    >>> _clean_list_keys({LIST_KEY: ['foo', 'bar']})
+    ['foo', 'bar']
+
+    """
+    for key, value in dict_.items():
+        if isinstance(value, dict):
+            dict_[key] = _clean_list_keys(value)
+        elif isinstance(value, list):
+            dict_[key] = [_clean_list_keys(v) if isinstance(v, dict) else v
+                          for v in value]
+    if LIST_KEY in dict_:
+        return dict_[LIST_KEY]
+    return dict_

tests/test_create.py

@@ -0,0 +1,177 @@
+import doctest
+from collections import namedtuple
+
+import pytest
+
+import jsonpath_ng
+from jsonpath_ng.ext import parse
+
+Params = namedtuple('Params', 'string initial_data insert_val target')
+
+
+@pytest.mark.parametrize('string, initial_data, insert_val, target', [
+
+    Params(string='$.foo',
+           initial_data={},
+           insert_val=42,
+           target={'foo': 42}),
+
+    Params(string='$.foo.bar',
+           initial_data={},
+           insert_val=42,
+           target={'foo': {'bar': 42}}),
+
+    Params(string='$.foo[0]',
+           initial_data={},
+           insert_val=42,
+           target={'foo': [42]}),
+
+    Params(string='$.foo[1]',
+           initial_data={},
+           insert_val=42,
+           target={'foo': [{}, 42]}),
+
+    Params(string='$.foo[0].bar',
+           initial_data={},
+           insert_val=42,
+           target={'foo': [{'bar': 42}]}),
+
+    Params(string='$.foo[1].bar',
+           initial_data={},
+           insert_val=42,
+           target={'foo': [{}, {'bar': 42}]}),
+
+    Params(string='$.foo[0][0]',
+           initial_data={},
+           insert_val=42,
+           target={'foo': [[42]]}),
+
+    Params(string='$.foo[1][1]',
+           initial_data={},
+           insert_val=42,
+           target={'foo': [{}, [{}, 42]]}),
+
+    Params(string='foo[0]',
+           initial_data={},
+           insert_val=42,
+           target={'foo': [42]}),
+
+    Params(string='foo[1]',
+           initial_data={},
+           insert_val=42,
+           target={'foo': [{}, 42]}),
+
+    Params(string='foo',
+           initial_data={},
+           insert_val=42,
+           target={'foo': 42}),
+
+    # Initial data can be a list if we expect a list back
+    Params(string='[0]',
+           initial_data=[],
+           insert_val=42,
+           target=[42]),
+
+    Params(string='[1]',
+           initial_data=[],
+           insert_val=42,
+           target=[{}, 42]),
+
+    # Converts initial data to a list if necessary
+    Params(string='[0]',
+           initial_data={},
+           insert_val=42,
+           target=[42]),
+
+    Params(string='[1]',
+           initial_data={},
+           insert_val=42,
+           target=[{}, 42]),
+
+    Params(string='foo[?bar="baz"].qux',
+           initial_data={'foo': [
+               {'bar': 'baz'},
+               {'bar': 'bizzle'},
+           ]},
+           insert_val=42,
+           target={'foo': [
+               {'bar': 'baz', 'qux': 42},
+               {'bar': 'bizzle'}
+           ]}),
+])
+def test_update_or_create(string, initial_data, insert_val, target):
+    jsonpath = parse(string)
+    result = jsonpath.update_or_create(initial_data, insert_val)
+    assert result == target
+
+
+@pytest.mark.parametrize('string, initial_data, insert_val, target', [
+    # Slice not supported
+    Params(string='foo[0:1]',
+           initial_data={},
+           insert_val=42,
+           target={'foo': [42, 42]}),
+    # result is {'foo': {}}
+
+    # Filter does not create items to meet criteria
+    Params(string='foo[?bar="baz"].qux',
+           initial_data={},
+           insert_val=42,
+           target={'foo': [{'bar': 'baz', 'qux': 42}]}),
+    # result is {'foo': {}}
+
+    # Does not convert initial data to a dictionary
+    Params(string='foo',
+           initial_data=[],
+           insert_val=42,
+           target={'foo': 42}),
+    # raises TypeError
+
+])
+@pytest.mark.xfail
+def test_unsupported_classes(string, initial_data, insert_val, target):
+    jsonpath = parse(string)
+    result = jsonpath.update_or_create(initial_data, insert_val)
+    assert result == target
+
+
+@pytest.mark.parametrize('string, initial_data, insert_val, target', [
+
+    Params(string='$.name[0].text',
+           initial_data={},
+           insert_val='Sir Michael',
+           target={'name': [{'text': 'Sir Michael'}]}),
+
+    Params(string='$.name[0].given[0]',
+           initial_data={'name': [{'text': 'Sir Michael'}]},
+           insert_val='Michael',
+           target={'name': [{'text': 'Sir Michael',
+                             'given': ['Michael']}]}),
+
+    Params(string='$.name[0].prefix[0]',
+           initial_data={'name': [{'text': 'Sir Michael',
+                                   'given': ['Michael']}]},
+           insert_val='Sir',
+           target={'name': [{'text': 'Sir Michael',
+                             'given': ['Michael'],
+                             'prefix': ['Sir']}]}),
+
+    Params(string='$.birthDate',
+           initial_data={'name': [{'text': 'Sir Michael',
+                                   'given': ['Michael'],
+                                   'prefix': ['Sir']}]},
+           insert_val='1943-05-05',
+           target={'name': [{'text': 'Sir Michael',
+                             'given': ['Michael'],
+                             'prefix': ['Sir']}],
+                   'birthDate': '1943-05-05'}),
+])
+def test_build_doc(string, initial_data, insert_val, target):
+    jsonpath = parse(string)
+    result = jsonpath.update_or_create(initial_data, insert_val)
+    assert result == target
+
+
+def test_doctests():
+    results = doctest.testmod(jsonpath_ng)
+    assert results.failed == 0