add dot nesting for inputs in JSONSink

Author: oesteban

nipype/interfaces/io.py

@@ -1860,7 +1860,19 @@ def _list_outputs(self):
 
 class JSONFileSinkInputSpec(DynamicTraitedSpec, BaseInterfaceInputSpec):
     out_file = File(desc='JSON sink file')
-    in_dict = traits.Dict(desc='input JSON dictionary')
+    in_dict = traits.Dict(value={}, usedefault=True,
+                          desc='input JSON dictionary')
+    _outputs = traits.Dict(value={}, usedefault=True)
+
+    def __setattr__(self, key, value):
+        if key not in self.copyable_trait_names():
+            if not isdefined(value):
+                super(JSONFileSinkInputSpec, self).__setattr__(key, value)
+            self._outputs[key] = value
+        else:
+            if key in self._outputs:
+                self._outputs[key] = value
+            super(JSONFileSinkInputSpec, self).__setattr__(key, value)
 
 
 class JSONFileSinkOutputSpec(TraitedSpec):
@@ -1869,7 +1881,10 @@ class JSONFileSinkOutputSpec(TraitedSpec):
 
 class JSONFileSink(IOBase):
 
-    """ Very simple frontend for storing values into a JSON file.
+    """
+    Very simple frontend for storing values into a JSON file.
+    Entries already existing in in_dict will be overridden by matching
+    entries dynamically added as inputs.
 
         .. warning::
 
@@ -1896,34 +1911,52 @@ class JSONFileSink(IOBase):
     input_spec = JSONFileSinkInputSpec
     output_spec = JSONFileSinkOutputSpec
 
-    def __init__(self, input_names=[], **inputs):
+    def __init__(self, infields=[], force_run=True, **inputs):
         super(JSONFileSink, self).__init__(**inputs)
-        self._input_names = filename_to_list(input_names)
-        add_traits(self.inputs, [name for name in self._input_names])
+        self._input_names = infields
+
+        undefined_traits = {}
+        for key in infields:
+            self.inputs.add_trait(key, traits.Any)
+            self.inputs._outputs[key] = Undefined
+            undefined_traits[key] = Undefined
+        self.inputs.trait_set(trait_change_notify=False, **undefined_traits)
+
+        if force_run:
+            self._always_run = True
+
+    def _process_name(self, name, val):
+        if '.' in name:
+            newkeys = name.split('.')
+            name = newkeys.pop(0)
+            nested_dict = {newkeys.pop(): val}
+
+            for nk in reversed(newkeys):
+                nested_dict = {nk: nested_dict}
+            val = nested_dict
+
+        return name, val
 
     def _list_outputs(self):
         import json
         import os.path as op
+
         if not isdefined(self.inputs.out_file):
             out_file = op.abspath('datasink.json')
         else:
             out_file = self.inputs.out_file
 
-        out_dict = dict()
+        out_dict = self.inputs.in_dict
 
-        if isdefined(self.inputs.in_dict):
-            if isinstance(self.inputs.in_dict, dict):
-                out_dict = self.inputs.in_dict
-        else:
-            for name in self._input_names:
-                val = getattr(self.inputs, name)
-                val = val if isdefined(val) else 'undefined'
-                out_dict[name] = val
+        # Overwrite in_dict entries automatically
+        for key, val in self.inputs._outputs.items():
+            if not isdefined(val) or key == 'trait_added':
+                continue
+            key, val = self._process_name(key, val)
+            out_dict[key] = val
 
         with open(out_file, 'w') as f:
             json.dump(out_dict, f)
         outputs = self.output_spec().get()
         outputs['out_file'] = out_file
         return outputs
-
-

nipype/interfaces/tests/test_io.py

@@ -238,3 +238,41 @@ def test_freesurfersource():
     yield assert_equal, fss.inputs.hemi, 'both'
     yield assert_equal, fss.inputs.subject_id, Undefined
     yield assert_equal, fss.inputs.subjects_dir, Undefined
+
+
+def test_jsonsink():
+    import json
+    import os
+
+    ds = nio.JSONFileSink()
+    yield assert_equal, ds.inputs._outputs, {}
+    ds = nio.JSONFileSink(in_dict={'foo': 'var'})
+    yield assert_equal, ds.inputs.in_dict, {'foo': 'var'}
+    ds = nio.JSONFileSink(infields=['test'])
+    yield assert_true, 'test' in ds.inputs.copyable_trait_names()
+
+    curdir = os.getcwd()
+    outdir = mkdtemp()
+    os.chdir(outdir)
+    js = nio.JSONFileSink(infields=['test'], in_dict={'foo': 'var'})
+    js.inputs.new_entry = 'someValue'
+    setattr(js.inputs, 'contrasts.alt', 'someNestedValue')
+    res = js.run()
+
+    with open(res.outputs['out_file'], 'r') as f:
+        data = json.load(f)
+    yield assert_true, data == {"contrasts": {"alt": "someNestedValue"}, "foo": "var", "new_entry": "someValue"}
+
+    js = nio.JSONFileSink(infields=['test'], in_dict={'foo': 'var'})
+    js.inputs.new_entry = 'someValue'
+    js.inputs.test = 'testInfields'
+    setattr(js.inputs, 'contrasts.alt', 'someNestedValue')
+    res = js.run()
+
+    with open(res.outputs['out_file'], 'r') as f:
+        data = json.load(f)
+    yield assert_true, data == {"test": "testInfields", "contrasts": {"alt": "someNestedValue"}, "foo": "var", "new_entry": "someValue"}
+
+    os.chdir(curdir)
+    shutil.rmtree(outdir)
+