Merge pull request #326 from eriknw/curry_reduce

serialize global curried objects.

toolz/_signatures.py

@@ -17,7 +17,7 @@
 import itertools
 import operator
 
-from .compatibility import PY3
+from .compatibility import PY3, import_module
 from .functoolz import (is_partial_args, is_arity, has_varargs,
                         has_keywords, num_required_args)
 
@@ -706,10 +706,7 @@ def expand_sig(sig):
 def create_signature_registry(module_info=module_info, signatures=signatures):
     for module, info in module_info.items():
         if isinstance(module, str):
-            modnames = module.split('.')
-            module = __import__(module)
-            for attr in modnames[1:]:
-                module = getattr(module, attr)
+            module = import_module(module)
         for name, sigs in info.items():
             if hasattr(module, name):
                 new_sigs = tuple(expand_sig(sig) for sig in sigs)

toolz/compatibility.py

@@ -7,7 +7,7 @@
 
 __all__ = ('map', 'filter', 'range', 'zip', 'reduce', 'zip_longest',
            'iteritems', 'iterkeys', 'itervalues', 'filterfalse',
-           'PY3', 'PY34', 'PYPY')
+           'PY3', 'PY34', 'PYPY', 'import_module')
 
 if PY3:
     map = map
@@ -31,3 +31,10 @@
     iteritems = operator.methodcaller('iteritems')
     iterkeys = operator.methodcaller('iterkeys')
     itervalues = operator.methodcaller('itervalues')
+
+try:
+    from importlib import import_module
+except ImportError:
+    def import_module(name):
+        __import__(name)
+        return sys.modules[name]

toolz/functoolz.py

@@ -4,7 +4,7 @@
 from operator import attrgetter
 from textwrap import dedent
 
-from .compatibility import PY3, PY33, PY34, PYPY
+from .compatibility import PY3, PY33, PY34, PYPY, import_module
 from .utils import no_default
 
 
@@ -321,17 +321,33 @@ def __get__(self, instance, owner):
             return self
         return curry(self, instance)
 
-    # pickle protocol because functools.partial objects can't be pickled
-    def __getstate__(self):
-        # dictoolz.keyfilter, I miss you!
+    def __reduce__(self):
+        func = self.func
+        modname = getattr(func, '__module__', None)
+        funcname = getattr(func, '__name__', None)
+        if modname and funcname:
+            module = import_module(modname)
+            obj = getattr(module, funcname, None)
+            if obj is self:
+                return funcname
+            elif isinstance(obj, curry) and obj.func is func:
+                func = '%s.%s' % (modname, funcname)
+
+        # functools.partial objects can't be pickled
         userdict = tuple((k, v) for k, v in self.__dict__.items()
                          if k != '_partial')
-        return self.func, self.args, self.keywords, userdict
-
-    def __setstate__(self, state):
-        func, args, kwargs, userdict = state
-        self.__init__(func, *args, **(kwargs or {}))
-        self.__dict__.update(userdict)
+        state = (type(self), func, self.args, self.keywords, userdict)
+        return (_restore_curry, state)
+
+
+def _restore_curry(cls, func, args, kwargs, userdict):
+    if isinstance(func, str):
+        modname, funcname = func.rsplit('.', 1)
+        module = import_module(modname)
+        func = getattr(module, funcname).func
+    obj = cls(func, *args, **(kwargs or {}))
+    obj.__dict__.update(userdict)
+    return obj
 
 
 @curry

toolz/tests/test_serialization.py

@@ -40,3 +40,18 @@ def test_instanceproperty():
     assert p2.__get__(None) is None
     assert p2.__get__(0) is False
     assert p2.__get__(1) is True
+
+
+def f(x, y):
+    return x, y
+
+
+def test_flip():
+    flip = pickle.loads(pickle.dumps(toolz.functoolz.flip))
+    assert flip is toolz.functoolz.flip
+    g1 = flip(f)
+    g2 = pickle.loads(pickle.dumps(g1))
+    assert g1(1, 2) == g2(1, 2) == f(2, 1)
+    g1 = flip(f)(1)
+    g2 = pickle.loads(pickle.dumps(g1))
+    assert g1(2) == g2(2) == f(2, 1)