bpo-37210: Fix pure Python pickle when _pickle is unavailable (GH-14016)

miss-islington · vstinner · web-flow · commit cbda40db7b60 · 2019-06-13T05:28:11.000-07:00
Allow pure Python implementation of pickle to work even when the C _pickle module is unavailable. Fix test_pickle when _pickle is missing: declare PyPicklerHookTests outside "if has_c_implementation:" block. (cherry picked from commit 63ab4ba) Co-authored-by: Victor Stinner <vstinner@redhat.com>
diff --git a/Lib/pickle.py b/Lib/pickle.py
@@ -36,10 +36,16 @@
 import codecs
 import _compat_pickle
 
-from _pickle import PickleBuffer
-
 __all__ = ["PickleError", "PicklingError", "UnpicklingError", "Pickler",
-           "Unpickler", "dump", "dumps", "load", "loads", "PickleBuffer"]
+           "Unpickler", "dump", "dumps", "load", "loads"]
+
+try:
+    from _pickle import PickleBuffer
+    __all__.append("PickleBuffer")
+    _HAVE_PICKLE_BUFFER = True
+except ImportError:
+    _HAVE_PICKLE_BUFFER = False
+
 
 # Shortcut for use in isinstance testing
 bytes_types = (bytes, bytearray)
@@ -812,31 +818,32 @@ def save_bytearray(self, obj):
             self.write(BYTEARRAY8 + pack("<Q", n) + obj)
     dispatch[bytearray] = save_bytearray
 
-    def save_picklebuffer(self, obj):
-        if self.proto < 5:
-            raise PicklingError("PickleBuffer can only pickled with "
-                                "protocol >= 5")
-        with obj.raw() as m:
-            if not m.contiguous:
-                raise PicklingError("PickleBuffer can not be pickled when "
-                                    "pointing to a non-contiguous buffer")
-            in_band = True
-            if self._buffer_callback is not None:
-                in_band = bool(self._buffer_callback(obj))
-            if in_band:
-                # Write data in-band
-                # XXX The C implementation avoids a copy here
-                if m.readonly:
-                    self.save_bytes(m.tobytes())
+    if _HAVE_PICKLE_BUFFER:
+        def save_picklebuffer(self, obj):
+            if self.proto < 5:
+                raise PicklingError("PickleBuffer can only pickled with "
+                                    "protocol >= 5")
+            with obj.raw() as m:
+                if not m.contiguous:
+                    raise PicklingError("PickleBuffer can not be pickled when "
+                                        "pointing to a non-contiguous buffer")
+                in_band = True
+                if self._buffer_callback is not None:
+                    in_band = bool(self._buffer_callback(obj))
+                if in_band:
+                    # Write data in-band
+                    # XXX The C implementation avoids a copy here
+                    if m.readonly:
+                        self.save_bytes(m.tobytes())
+                    else:
+                        self.save_bytearray(m.tobytes())
                 else:
-                    self.save_bytearray(m.tobytes())
-            else:
-                # Write data out-of-band
-                self.write(NEXT_BUFFER)
-                if m.readonly:
-                    self.write(READONLY_BUFFER)
+                    # Write data out-of-band
+                    self.write(NEXT_BUFFER)
+                    if m.readonly:
+                        self.write(READONLY_BUFFER)
 
-    dispatch[PickleBuffer] = save_picklebuffer
+        dispatch[PickleBuffer] = save_picklebuffer
 
     def save_str(self, obj):
         if self.bin:
diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py
@@ -203,6 +203,13 @@ def get_dispatch_table(self):
         return collections.ChainMap({}, pickle.dispatch_table)
 
 
+class PyPicklerHookTests(AbstractHookTests):
+    class CustomPyPicklerClass(pickle._Pickler,
+                               AbstractCustomPicklerClass):
+        pass
+    pickler_class = CustomPyPicklerClass
+
+
 if has_c_implementation:
     class CPickleTests(AbstractPickleModuleTests):
         from _pickle import dump, dumps, load, loads, Pickler, Unpickler
@@ -255,12 +262,6 @@ class CChainDispatchTableTests(AbstractDispatchTableTests):
         def get_dispatch_table(self):
             return collections.ChainMap({}, pickle.dispatch_table)
 
-    class PyPicklerHookTests(AbstractHookTests):
-        class CustomPyPicklerClass(pickle._Pickler,
-                                   AbstractCustomPicklerClass):
-            pass
-        pickler_class = CustomPyPicklerClass
-
     class CPicklerHookTests(AbstractHookTests):
         class CustomCPicklerClass(_pickle.Pickler, AbstractCustomPicklerClass):
             pass
diff --git a/Misc/NEWS.d/next/Library/2019-06-12-16-10-50.bpo-37210.r4yMg6.rst b/Misc/NEWS.d/next/Library/2019-06-12-16-10-50.bpo-37210.r4yMg6.rst
@@ -0,0 +1 @@
+Allow pure Python implementation of :mod:`pickle` to work even when the C :mod:`_pickle` module is unavailable.

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Allow pure Python implementation of :mod:`pickle` to work even when the C :mod:`_pickle` module is unavailable.