[pypy-commit] pypy default: Issue #1853

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r91606:0b207c82e5fd
Date: 2017-06-15 11:25 +0200
http://bitbucket.org/pypy/pypy/changeset/0b207c82e5fd/

Log:	Issue #1853

	Add cPickle.Unpickler.find_global

diff --git a/lib_pypy/cPickle.py b/lib_pypy/cPickle.py
--- a/lib_pypy/cPickle.py
+++ b/lib_pypy/cPickle.py
@@ -431,7 +431,14 @@
         self.append(obj)
 
     def find_class(self, module, name):
-        # Subclasses may override this
+        if self.find_global is None:
+            raise UnpicklingError(
+                "Global and instance pickles are not supported.")
+        return self.find_global(module, name)
+
+    def find_global(self, module, name):
+        # This can officially be patched directly in the Unpickler
+        # instance, according to the docs
         __import__(module)
         mod = sys.modules[module]
         klass = getattr(mod, name)
diff --git a/pypy/module/test_lib_pypy/test_cPickle.py b/pypy/module/test_lib_pypy/test_cPickle.py
--- a/pypy/module/test_lib_pypy/test_cPickle.py
+++ b/pypy/module/test_lib_pypy/test_cPickle.py
@@ -9,3 +9,25 @@
 def test_bad_key():
     e = py.test.raises(cPickle.UnpicklingError, cPickle.loads, "v")
     assert str(e.value) == "invalid load key, 'v'."
+
+def test_find_global():
+    import time, cStringIO
+    entry = time.strptime('Fri Mar 27 22:20:42 2017')
+    f = cStringIO.StringIO()
+    cPickle.Pickler(f).dump(entry)
+
+    f = cStringIO.StringIO(f.getvalue())
+    e = cPickle.Unpickler(f).load()
+    assert e == entry
+
+    f = cStringIO.StringIO(f.getvalue())
+    up = cPickle.Unpickler(f)
+    up.find_global = None
+    e = py.test.raises(cPickle.UnpicklingError, up.load)
+    assert str(e.value) == "Global and instance pickles are not supported."
+
+    f = cStringIO.StringIO(f.getvalue())
+    up = cPickle.Unpickler(f)
+    up.find_global = lambda module, name: lambda a, b: (a, b)
+    e = up.load()
+    assert e == ((2017, 3, 27, 22, 20, 42, 4, 86, -1), {})