ENH customizable serialization process for mp.Queue

Author: tomMoral

Lib/multiprocessing/queues.py

@@ -27,6 +27,22 @@
 
 from .util import debug, info, Finalize, register_after_fork, is_exiting
 
+#
+#  Sendable Object, with a serialization protocol
+#
+
+
+class _SendableObject(object):
+    def __init__(self, obj, serialization=None):
+        self.obj = obj
+        self.serialization = serialization
+
+    def serialize(self):
+        if self.serialization:
+            return self.serialization(self.obj)
+        return self.obj
+
+
 #
 # Queue type using a pipe, buffer and thread
 #
@@ -78,17 +94,26 @@ def _after_fork(self):
         self._poll = self._reader.poll
 
     def put(self, obj, block=True, timeout=None):
+        self._put_bytes(obj, block=block, timeout=timeout,
+                        serialization=_ForkingPickler.dumps)
+
+    def _put_bytes(self, obj, block=True, timeout=None, serialization=None):
         assert not self._closed, "Queue {0!r} has been closed".format(self)
         if not self._sem.acquire(block, timeout):
             raise Full
 
         with self._notempty:
             if self._thread is None:
                 self._start_thread()
-            self._buffer.append(obj)
+            self._buffer.append(_SendableObject(
+                obj, serialization=serialization))
             self._notempty.notify()
 
     def get(self, block=True, timeout=None):
+        return self._get_bytes(block=block, timeout=timeout,
+                               deserialization=_ForkingPickler.loads)
+
+    def _get_bytes(self, block=True, timeout=None, deserialization=None):
         if block and timeout is None:
             with self._rlock:
                 res = self._recv_bytes()
@@ -109,8 +134,10 @@ def get(self, block=True, timeout=None):
                 self._sem.release()
             finally:
                 self._rlock.release()
-        # unserialize the data after having released the lock
-        return _ForkingPickler.loads(res)
+        # un-serialize the data after having released the lock
+        if deserialization:
+            return deserialization(res)
+        return res
 
     def qsize(self):
         # Raises NotImplementedError on Mac OSX because of broken sem_getvalue()
@@ -233,7 +260,7 @@ def _feed(buffer, notempty, send_bytes, writelock, close, ignore_epipe,
                             return
 
                         # serialize the data before acquiring the lock
-                        obj = _ForkingPickler.dumps(obj)
+                        obj = obj.serialize()
                         if wacquire is None:
                             send_bytes(obj)
                         else:
@@ -255,7 +282,7 @@ def _feed(buffer, notempty, send_bytes, writelock, close, ignore_epipe,
                     info('error in queue thread: %s', e)
                     return
                 else:
-                    onerror(e, obj)
+                    onerror(e, obj.obj)
 
     @staticmethod
     def _on_queue_feeder_error(e, obj):
@@ -299,7 +326,8 @@ def put(self, obj, block=True, timeout=None):
         with self._notempty, self._cond:
             if self._thread is None:
                 self._start_thread()
-            self._buffer.append(obj)
+            self._buffer.append(_SendableObject(
+                obj, serialization=_ForkingPickler.dumps))
             self._unfinished_tasks.release()
             self._notempty.notify()
 
@@ -342,14 +370,25 @@ def __setstate__(self, state):
         self._poll = self._reader.poll
 
     def get(self):
+        # Get the object and deserialize it with the _ForkingPickler
+        return self._get_bytes(deserialization=_ForkingPickler.loads)
+
+    def _get_bytes(self, deserialization=None):
         with self._rlock:
             res = self._reader.recv_bytes()
         # unserialize the data after having released the lock
-        return _ForkingPickler.loads(res)
+        if deserialization:
+            return deserialization(res)
+        return res
 
     def put(self, obj):
+        # Get the object and deserialize it with the _ForkingPickler
+        self._put_bytes(obj, serialization=_ForkingPickler.dumps)
+
+    def _put_bytes(self, obj, serialization=None):
         # serialize the data before acquiring the lock
-        obj = _ForkingPickler.dumps(obj)
+        if serialization:
+            obj = serialization(obj)
         if self._wlock is None:
             # writes to a message oriented win32 pipe are atomic
             self._writer.send_bytes(obj)

Lib/test/_test_multiprocessing.py

@@ -18,6 +18,7 @@
 import random
 import logging
 import struct
+import pickle
 import operator
 import weakref
 import test.support
@@ -1066,6 +1067,24 @@ def _on_queue_feeder_error(e, obj):
         # Assert that the serialization and the hook have been called correctly
         self.assertTrue(not_serializable_obj.reduce_was_called)
         self.assertTrue(not_serializable_obj.on_queue_feeder_error_was_called)
+
+    def test_queue_serialization(self):
+        # bpo-30006: verify feeder handles exceptions using the
+        # _on_queue_feeder_error hook.
+        if self.TYPE != 'processes':
+            self.skipTest('test not appropriate for {}'.format(self.TYPE))
+
+        q = self.Queue()
+
+        # Custom serialization
+        def serialization(x):
+            return pickle.dumps(x * 2)
+        q._put_bytes(21, serialization=serialization)
+        self.assertEqual(q.get(), 42)
+
+        # Custom bytes channels
+        q._put_bytes(bytes(42), serialization=None)
+        self.assertEqual(q._get_bytes(), bytes(42))
 #
 #
 #