Merge pull request #3041 from minrk/requirepush

support non-modules in @require

IPython/parallel/controller/dependency.py

@@ -5,7 +5,7 b' Authors:'
 * Min RK
 """
 #-----------------------------------------------------------------------------
-#  Copyright (C) 2010-2011  The IPython Development Team
+#  Copyright (C) 2013  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
@@ -17,6 +17,7 b' from IPython.parallel.client.asyncresult import AsyncResult'
 from IPython.parallel.error import UnmetDependency
 from IPython.parallel.util import interactive
 from IPython.utils import py3compat
+from IPython.utils.pickleutil import can, uncan
 
 class depend(object):
     """Dependency decorator, for use with tasks.
@@ -58,12 +59,12 b' class dependent(object):'
         self.df = df
         self.dargs = dargs
         self.dkwargs = dkwargs
-
-    def __call__(self, *args, **kwargs):
-        # if hasattr(self.f, 'func_globals') and hasattr(self.df, 'func_globals'):
-        #     self.df.func_globals = self.f.func_globals
+    
+    def check_dependency(self):
         if self.df(*self.dargs, **self.dkwargs) is False:
             raise UnmetDependency()
+    
+    def __call__(self, *args, **kwargs):
         return self.f(*args, **kwargs)
     
     if not py3compat.PY3:
@@ -72,41 +73,62 b' class dependent(object):'
             return self.func_name
 
 @interactive
-def _require(*names):
+def _require(*modules, **mapping):
     """Helper for @require decorator."""
     from IPython.parallel.error import UnmetDependency
+    from IPython.utils.pickleutil import uncan
     user_ns = globals()
-    for name in names:
-        if name in user_ns:
-            continue
+    for name in modules:
         try:
-            exec 'import %s'%name in user_ns
+            exec 'import %s' % name in user_ns
         except ImportError:
             raise UnmetDependency(name)
+            
+    for name, cobj in mapping.items():
+        user_ns[name] = uncan(cobj, user_ns)
     return True
 
-def require(*mods):
-    """Simple decorator for requiring names to be importable.
+def require(*objects, **mapping):
+    """Simple decorator for requiring local objects and modules to be available
+    when the decorated function is called on the engine.
+    
+    Modules specified by name or passed directly will be imported
+    prior to calling the decorated function.
+    
+    Objects other than modules will be pushed as a part of the task.
+    Functions can be passed positionally,
+    and will be pushed to the engine with their __name__.
+    Other objects can be passed by keyword arg.
     
     Examples
     --------
     
     In [1]: @require('numpy')
        ...: def norm(a):
-       ...:     import numpy
        ...:     return numpy.linalg.norm(a,2)
+
+    In [2]: foo = lambda x: x*x
+    In [3]: @require(foo)
+       ...: def bar(a):
+       ...:     return foo(1-a)
     """
     names = []
-    for mod in mods:
-        if isinstance(mod, ModuleType):
-            mod = mod.__name__
+    for obj in objects:
+        if isinstance(obj, ModuleType):
+            obj = obj.__name__
         
-        if isinstance(mod, basestring):
-            names.append(mod)
+        if isinstance(obj, basestring):
+            names.append(obj)
+        elif hasattr(obj, '__name__'):
+            mapping[obj.__name__] = obj
         else:
-            raise TypeError("names must be modules or module names, not %s"%type(mod))
+            raise TypeError("Objects other than modules and functions "
+                "must be passed by kwarg, but got: %s" % type(obj)
+            )
     
-    return depend(_require, *names)
+    for name, obj in mapping.items():
+        mapping[name] = can(obj)
+    return depend(_require, *names, **mapping)
 
 class Dependency(set):
     """An object for representing a set of msg_id dependencies.

IPython/parallel/tests/test_dependency.py

@@ -37,6 +37,10 b' def wait(n):'
     time.sleep(n)
     return n
 
[e-mail]
+def func(x):
+    return x*x
+
 mixed = map(str, range(10))
 completed = map(str, range(0,10,2))
 failed = map(str, range(1,10,2))
@@ -104,3 +108,29 b' class DependencyTest(ClusterTestCase):'
         dep.all=False
         self.assertUnmet(dep)
         self.assertUnreachable(dep)
+    
+    def test_require_function(self):
+        
+        @pmod.interactive
+        def bar(a):
+            return func(a)
+
+        @pmod.require(func)
+        @pmod.interactive
+        def bar2(a):
+            return func(a)
+        
+        self.client[:].clear()
+        self.assertRaisesRemote(NameError, self.view.apply_sync, bar, 5)
+        ar = self.view.apply_async(bar2, 5)
+        self.assertEqual(ar.get(5), func(5))
+
+    def test_require_object(self):
+        
+        @pmod.require(foo=func)
+        @pmod.interactive
+        def bar(a):
+            return foo(a)
+
+        ar = self.view.apply_async(bar, 5)
+        self.assertEqual(ar.get(5), func(5))

IPython/utils/pickleutil.py

@@ -49,9 +49,26 b' else:'
 
 
 class CannedObject(object):
-    def __init__(self, obj, keys=[]):
+    def __init__(self, obj, keys=[], hook=None):
+        """can an object for safe pickling
+        
+        Parameters
+        ==========
+        
+        obj:
+            The object to be canned
+        keys: list (optional)
+            list of attribute names that will be explicitly canned / uncanned
+        hook: callable (optional)
+            An optional extra callable,
+            which can do additional processing of the uncanned object.
+        
+        large data may be offloaded into the buffers list,
+        used for zero-copy transfers.
+        """
         self.keys = keys
         self.obj = copy.copy(obj)
+        self.hook = can(hook)
         for key in keys:
             setattr(self.obj, key, can(getattr(obj, key)))
         
@@ -60,8 +77,13 b' class CannedObject(object):'
     def get_object(self, g=None):
         if g is None:
             g = {}
+        obj = self.obj
         for key in self.keys:
-            setattr(self.obj, key, uncan(getattr(self.obj, key), g))
+            setattr(obj, key, uncan(getattr(obj, key), g))
+        
+        if self.hook:
+            self.hook = uncan(self.hook, g)
+            self.hook(obj, g)
         return self.obj
     
 
@@ -302,6 +324,11 b' def uncan_sequence(obj, g=None):'
     else:
         return obj
 
+def _uncan_dependent_hook(dep, g=None):
+    dep.check_dependency()
+    
+def can_dependent(obj):
+    return CannedObject(obj, keys=('f', 'df'), hook=_uncan_dependent_hook)
 
 #-------------------------------------------------------------------------------
 # API dictionaries
@@ -310,7 +337,7 b' def uncan_sequence(obj, g=None):'
 # These dicts can be extended for custom serialization of new objects
 
 can_map = {
-    'IPython.parallel.dependent' : lambda obj: CannedObject(obj, keys=('f','df')),
+    'IPython.parallel.dependent' : can_dependent,
     'numpy.ndarray' : CannedArray,
     FunctionType : CannedFunction,
     bytes : CannedBytes,