Merge pull request #551 from mwaskom/enh/function_imports

ENH Increased flexibility in Function interface

Author: chrisgorgo

CHANGES

@@ -9,6 +9,9 @@ Next release
 * ENH: Minor improvements to FSL's FUGUE and FLIRT interfaces
 * ENH: Added optional dilation of parcels in cmtk.Parcellate 
 * ENH: Interpolation mode added to afni.Resample
+* ENH: Function interface can accept a list of strings containing import statements
+       that allow external functions to run without their imports defined in the
+       function body
 
 * FIX: SpecifyModel works with 3D files correctly now.
 

doc/users/function_interface.rst

@@ -71,6 +71,11 @@ be imported within the function itself::
 Without explicitly importing Nibabel in the body of the function, this
 would fail.
 
+Alternatively, it is possible to provide a list of strings corresponding
+to the imports needed to execute a function as a parameter of the `Function`
+constructor. This allows for the use of external functions that do not
+import all external definitions inside the function body.
+
 Hello World - Function interface in a workflow
 ----------------------------------------------
 

nipype/interfaces/tests/test_utility.py

@@ -4,7 +4,8 @@
 import shutil
 from tempfile import mkdtemp
 
-from nipype.testing import assert_equal, assert_true
+import numpy as np
+from nipype.testing import assert_equal, assert_true, assert_raises
 from nipype.interfaces import utility
 import nipype.pipeline.engine as pe
 
@@ -66,3 +67,49 @@ def increment_array(in_array):
     # Clean up
     os.chdir(origdir)
     shutil.rmtree(tempdir)
+
+
+def make_random_array(size):
+
+    return np.random.randn(size, size)
+
+
+def should_fail():
+
+    tempdir = os.path.realpath(mkdtemp())
+    origdir = os.getcwd()
+    os.chdir(tempdir)
+
+    node = pe.Node(utility.Function(input_names=["size"],
+                                    output_names=["random_array"],
+                                    function=make_random_array),
+                   name="should_fail")
+    try:
+        node.inputs.size = 10
+        node.run()
+    finally:
+        os.chdir(origdir)
+        shutil.rmtree(tempdir)
+
+
+assert_raises(NameError, should_fail)
+
+
+def test_function_with_imports():
+
+    tempdir = os.path.realpath(mkdtemp())
+    origdir = os.getcwd()
+    os.chdir(tempdir)
+
+    node = pe.Node(utility.Function(input_names=["size"],
+                                    output_names=["random_array"],
+                                    function=make_random_array,
+                                    imports=["import numpy as np"]),
+                   name="should_not_fail")
+    print node.inputs.function_str
+    try:
+        node.inputs.size = 10
+        node.run()
+    finally:
+        os.chdir(origdir)
+        shutil.rmtree(tempdir)

nipype/interfaces/utility.py

@@ -2,6 +2,7 @@
 # vi: set ft=python sts=4 ts=4 sw=4 et:
 import os
 import re
+from cPickle import dumps, loads
 import numpy as np
 import nibabel as nb
 
@@ -11,7 +12,7 @@
     InputMultiPath, BaseInterface, BaseInterfaceInputSpec)
 from nipype.interfaces.io import IOBase, add_traits
 from nipype.testing import assert_equal
-from nipype.utils.misc import getsource, create_function_from_source, dumps
+from nipype.utils.misc import getsource, create_function_from_source
 
 
 class IdentityInterface(IOBase):
@@ -335,7 +336,8 @@ class Function(IOBase):
     input_spec = FunctionInputSpec
     output_spec = DynamicTraitedSpec
 
-    def __init__(self, input_names, output_names, function=None, **inputs):
+    def __init__(self, input_names, output_names, function=None, imports=None,
+                 **inputs):
         """
 
         Parameters
@@ -344,7 +346,15 @@ def __init__(self, input_names, output_names, function=None, **inputs):
         input_names: single str or list
             names corresponding to function inputs
         output_names: single str or list
-            names corresponding to function outputs. has to match the number of outputs
+            names corresponding to function outputs.
+            has to match the number of outputs
+        function : callable
+            callable python object. must be able to execute in an
+            isolated namespace (possibly in concert with the ``imports``
+            parameter)
+        imports : list of strings
+            list of import statements that allow the function to execute
+            in an otherwise empty namespace
         """
 
         super(Function, self).__init__(**inputs)
@@ -354,15 +364,18 @@ def __init__(self, input_names, output_names, function=None, **inputs):
                     self.inputs.function_str = getsource(function)
                 except IOError:
                     raise Exception('Interface Function does not accept ' \
-                                        'function objects defined interactively in a python session')
+                                    'function objects defined interactively ' \
+                                    'in a python session')
             elif isinstance(function, str):
                 self.inputs.function_str = dumps(function)
             else:
                 raise Exception('Unknown type of function')
-        self.inputs.on_trait_change(self._set_function_string, 'function_str')
+        self.inputs.on_trait_change(self._set_function_string,
+                                    'function_str')
         self._input_names = filename_to_list(input_names)
         self._output_names = filename_to_list(output_names)
         add_traits(self.inputs, [name for name in self._input_names])
+        self.imports = imports
         self._out = {}
         for name in self._output_names:
             self._out[name] = None
@@ -373,7 +386,8 @@ def _set_function_string(self, obj, name, old, new):
                 function_source = getsource(new)
             elif isinstance(new, str):
                 function_source = dumps(new)
-            self.inputs.trait_set(trait_change_notify=False, **{'%s' % name: function_source})
+            self.inputs.trait_set(trait_change_notify=False,
+                                  **{'%s' % name: function_source})
 
     def _add_output_traits(self, base):
         undefined_traits = {}
@@ -384,7 +398,8 @@ def _add_output_traits(self, base):
         return base
 
     def _run_interface(self, runtime):
-        function_handle = create_function_from_source(self.inputs.function_str)
+        function_handle = create_function_from_source(self.inputs.function_str,
+                                                      self.imports)
 
         args = {}
         for name in self._input_names:

nipype/utils/misc.py

@@ -46,21 +46,38 @@ def getsource(function):
     src = dumps(dedent(inspect.getsource(function)))
     return src
 
-def create_function_from_source(function_source):
+def create_function_from_source(function_source, imports=None):
     """Return a function object from a function source
+
+    Parameters
+    ----------
+    function_source : pickled string
+        string in pickled form defining a function
+    imports : list of strings
+        list of import statements in string form that allow the function
+        to be executed in an otherwise empty namespace
     """
     ns = {}
+    import_keys = []
     try:
+        if imports is not None:
+            for statement in imports:
+                exec statement in ns
+            import_keys = ns.keys()
+
         exec loads(function_source) in ns
+
     except Exception, msg:
         msg = str(msg) + '\nError executing function:\n %s\n'%function_source
-        msg += '\n'.join( ["Functions in connection strings have to be standalone.",
-                           "They cannot be declared either interactively or inside",
-                           "another function or inline in the connect string. Any",
-                           "imports should be done inside the function"
+        msg += '\n'.join(["Functions in connection strings have to be standalone.",
+                          "They cannot be declared either interactively or inside",
+                          "another function or inline in the connect string. Any",
+                          "imports should be done inside the function"
                            ])
         raise RuntimeError(msg)
-    funcname = [name for name in ns.keys() if name != '__builtins__'][0]
+    ns_funcs = list(set(ns) - set(import_keys + ['__builtins__']))
+    assert len(ns_funcs) == 1, "Function or inputs are ill-defined"
+    funcname = ns_funcs[0]
     func = ns[funcname]
     return func