option to keep siblings when filtering by keyval

jsonextended/edict.py

@@ -847,25 +847,37 @@ def is_in(a,b):
     flatd = {k:v for k,v in flatd.items() if is_in(v,vals)}
     return unflatten(flatd,list_of_dicts=list_of_dicts)
 
-def filter_keyvals(d,vals=None,list_of_dicts=False):
+def filter_keyvals(d,vals=None,
+                  keep_siblings=False,
+                  list_of_dicts=False):
     """ filters leaf nodes key:value pairs of nested dictionary
 
     Parameters
     ----------
     d : dict
     vals : list of tuples
         (key,value) to filter by
+    keep_siblings : bool
+        keep all sibling paths
     list_of_dicts: bool
         treat list of dicts as additional branches
 
     Examples
     --------
 
     >>> from pprint import pprint
+
     >>> d = {1:{6:'a'},3:{7:'a'},2:{6:"b"},4:{5:{6:'a'}}}
     >>> pprint(filter_keyvals(d,[(6,'a')]))
     {1: {6: 'a'}, 4: {5: {6: 'a'}}}
 
+    >>> d2 = {'a':{'b':1,'c':2}, 'd':3}
+    >>> pprint(filter_keyvals(d2,[('b',1)],keep_siblings=False))
+    {'a': {'b': 1}}
+
+    >>> pprint(filter_keyvals(d2,[('b',1)],keep_siblings=True))
+    {'a': {'b': 1, 'c': 2}}
+
     """
     vals = [] if vals is None else vals
     list_of_dicts = '__list__' if list_of_dicts else None
@@ -877,8 +889,18 @@ def is_in(a,b):
         except:
             return False
 
-    flatd = {k:v for k,v in flatd.items() if is_in((k[-1],v),vals)}
-    return unflatten(flatd,list_of_dicts=list_of_dicts)
+    if not keep_siblings:
+        newd = {k:v for k,v in flatd.items() if is_in((k[-1],v),vals)}
+    else:
+        prune = [tuple(k[:-1]) for k,v in flatd.items() if is_in((k[-1],v),vals)]
+        newd = {}
+        for k,v in flatd.items():
+            for p in prune:
+                if tuple(k[:len(p)]) == p:
+                    newd[k] = v
+                    break
+
+    return unflatten(newd,list_of_dicts=list_of_dicts)
 
 def filter_keys(d, keys, use_wildcards=False,
                 list_of_dicts=False):

jsonextended/units/core.py

@@ -28,7 +28,7 @@ def get_in_units(value,units):
 
 def apply_unitschema(data, uschema, as_quantity=True,
                     raise_outerr=False, convert_base=False,
-                    use_wildcards=False):
+                    use_wildcards=False,list_of_dicts=False):
     """ apply the unit schema to the data 
     
     Parameters
@@ -44,6 +44,8 @@ def apply_unitschema(data, uschema, as_quantity=True,
         rescale units to base units
     use_wildcards : bool
         if true, can use * (matches everything) and ? (matches any single character)
+    list_of_dicts: bool
+        treat list of dicts as additional branches
     
     Examples
     --------
@@ -81,12 +83,13 @@ def apply_unitschema(data, uschema, as_quantity=True,
         _Quantity
     except NameError:
         raise ImportError('please install pint to use this module')
+    list_of_dicts = '__list__' if list_of_dicts else None
 
     # flatten edict
     uschema_flat = flatten(uschema,key_as_tuple=True)
     # sorted by longest key size, to get best match first
     uschema_keys = sorted(uschema_flat, key=len, reverse=True)
-    data_flat = flatten(data,key_as_tuple=True)
+    data_flat = flatten(data,key_as_tuple=True,list_of_dicts=list_of_dicts)
 
     for dkey, dvalue in data_flat.items():
         converted = False
@@ -124,7 +127,7 @@ def apply_unitschema(data, uschema, as_quantity=True,
         if not converted and raise_outerr:
             raise KeyError('could not find units for {}'.format(dkey))
 
-    return unflatten(data_flat)
+    return unflatten(data_flat,list_of_dicts=list_of_dicts)
 
 def split_quantities(data,units='units',magnitude='magnitude',
                     list_of_dicts=False):