Add update_or_create() method

Author: kaapstorm

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_