Permalink
Browse files

ENH: Add ability to pickle ufuncs as requested in ticket #1743

  • Loading branch information...
1 parent 27befc8 commit 4c0576fe9947ef2af8351405e0990cebd83ccbb6 @teoliphant teoliphant committed Feb 6, 2012
Showing with 35 additions and 0 deletions.
  1. +26 −0 numpy/core/__init__.py
  2. +9 −0 numpy/core/tests/test_ufunc.py
@@ -39,3 +39,29 @@
from numpy.testing import Tester
test = Tester().test
bench = Tester().bench
+
+# Make it possible so that ufuncs can be pickled
+# Here are the loading and unloading functions
+# The name numpy.core._ufunc_reconstruct must be
+# available for unpickling to work.
+def _ufunc_reconstruct(module, name):
+ mod = __import__(module)
+ return getattr(mod, name)
+
+def _ufunc_reduce(func):
+ from pickle import whichmodule
+ name = func.__name__
+ return _ufunc_reconstruct, (whichmodule(func,name), name)
+
+
+import sys
+if sys.version_info[0] < 3:
+ import copy_reg as copyreg
+else:
+ import copyreg
+
+copyreg.pickle(ufunc, _ufunc_reduce, _ufunc_reconstruct)
+# Unclutter namespace (must keep _ufunc_reconstruct for unpickling)
+del copyreg
+del sys
+del _ufunc_reduce
@@ -5,6 +5,15 @@
import numpy.core.umath_tests as umt
class TestUfunc(TestCase):
+ def test_pickle(self):
+ import pickle
+ assert pickle.loads(pickle.dumps(np.sin)) is np.sin
+
+ def test_pickle_withstring(self):
+ import pickle
+ astring = "cnumpy.core\n_ufunc_reconstruct\np0\n(S'numpy.core.umath'\np1\nS'cos'\np2\ntp3\nRp4\n."
+ assert pickle.loads(astring) is np.cos
+
def test_reduceat_shifting_sum(self) :
L = 6
x = np.arange(L)

0 comments on commit 4c0576f

Please sign in to comment.