Issue #14166: Pickler objects now have an optional `dispatch_table` attribute...

Issue #14166: Pickler objects now have an optional `dispatch_table` attribute which allows to set custom per-pickler reduction functions.
Patch by sbt.

Doc/library/copyreg.rst

@@ -32,6 +32,8 @@ Such constructors may be factory functions or class instances.
    returned by *function* at pickling time.  :exc:`TypeError` will be raised if
    *object* is a class or *constructor* is not callable.
 
-   See the :mod:`pickle` module for more details on the interface expected of
-   *function* and *constructor*.
-
+   See the :mod:`pickle` module for more details on the interface
+   expected of *function* and *constructor*.  Note that the
+   :attr:`~pickle.Pickler.dispatch_table` attribute of a pickler
+   object or subclass of :class:`pickle.Pickler` can also be used for
+   declaring reduction functions.

Doc/library/pickle.rst

@@ -285,6 +285,29 @@ The :mod:`pickle` module exports two classes, :class:`Pickler` and
 
       See :ref:`pickle-persistent` for details and examples of uses.
 
+   .. attribute:: dispatch_table
+
+      A pickler object's dispatch table is a registry of *reduction
+      functions* of the kind which can be declared using
+      :func:`copyreg.pickle`.  It is a mapping whose keys are classes
+      and whose values are reduction functions.  A reduction function
+      takes a single argument of the associated class and should
+      conform to the same interface as a :meth:`~object.__reduce__`
+      method.
+
+      By default, a pickler object will not have a
+      :attr:`dispatch_table` attribute, and it will instead use the
+      global dispatch table managed by the :mod:`copyreg` module.
+      However, to customize the pickling for a specific pickler object
+      one can set the :attr:`dispatch_table` attribute to a dict-like
+      object.  Alternatively, if a subclass of :class:`Pickler` has a
+      :attr:`dispatch_table` attribute then this will be used as the
+      default dispatch table for instances of that class.
+
+      See :ref:`pickle-dispatch` for usage examples.
+
+      .. versionadded:: 3.3
+
    .. attribute:: fast
 
       Deprecated. Enable fast mode if set to a true value.  The fast mode
@@ -575,6 +598,44 @@ pickle external objects by reference.
 
 .. literalinclude:: ../includes/dbpickle.py
 
+.. _pickle-dispatch:
+
+Dispatch Tables
+^^^^^^^^^^^^^^^
+
+If one wants to customize pickling of some classes without disturbing
+any other code which depends on pickling, then one can create a
+pickler with a private dispatch table.
+
+The global dispatch table managed by the :mod:`copyreg` module is
+available as :data:`copyreg.dispatch_table`.  Therefore, one may
+choose to use a modified copy of :data:`copyreg.dispatch_table` as a
+private dispatch table.
+
+For example ::
+
+   f = io.BytesIO()
+   p = pickle.Pickler(f)
+   p.dispatch_table = copyreg.dispatch_table.copy()
+   p.dispatch_table[SomeClass] = reduce_SomeClass
+
+creates an instance of :class:`pickle.Pickler` with a private dispatch
+table which handles the ``SomeClass`` class specially.  Alternatively,
+the code ::
+
+   class MyPickler(pickle.Pickler):
+       dispatch_table = copyreg.dispatch_table.copy()
+       dispatch_table[SomeClass] = reduce_SomeClass
+   f = io.BytesIO()
+   p = MyPickler(f)
+
+does the same, but all instances of ``MyPickler`` will by default
+share the same dispatch table.  The equivalent code using the
+:mod:`copyreg` module is ::
+
+   copyreg.pickle(SomeClass, reduce_SomeClass)
+   f = io.BytesIO()
+   p = pickle.Pickler(f)
 
 .. _pickle-state:
 

Lib/pickle.py

@@ -297,8 +297,8 @@ class _Pickler:
             f(self, obj) # Call unbound method with explicit self
             return
 
-        # Check copyreg.dispatch_table
-        reduce = dispatch_table.get(t)
+        # Check private dispatch table if any, or else copyreg.dispatch_table
+        reduce = getattr(self, 'dispatch_table', dispatch_table).get(t)
         if reduce:
             rv = reduce(obj)
         else:

Lib/test/pickletester.py

@@ -1605,6 +1605,105 @@ class AbstractPicklerUnpicklerObjectTests(unittest.TestCase):
                 self.assertEqual(unpickler.load(), data)
 
 
+# Tests for dispatch_table attribute
+
+REDUCE_A = 'reduce_A'
+
+class AAA(object):
+    def __reduce__(self):
+        return str, (REDUCE_A,)
+
+class BBB(object):
+    pass
+
+class AbstractDispatchTableTests(unittest.TestCase):
+
+    def test_default_dispatch_table(self):
+        # No dispatch_table attribute by default
+        f = io.BytesIO()
+        p = self.pickler_class(f, 0)
+        with self.assertRaises(AttributeError):
+            p.dispatch_table
+        self.assertFalse(hasattr(p, 'dispatch_table'))
+
+    def test_class_dispatch_table(self):
+        # A dispatch_table attribute can be specified class-wide
+        dt = self.get_dispatch_table()
+
+        class MyPickler(self.pickler_class):
+            dispatch_table = dt
+
+        def dumps(obj, protocol=None):
+            f = io.BytesIO()
+            p = MyPickler(f, protocol)
+            self.assertEqual(p.dispatch_table, dt)
+            p.dump(obj)
+            return f.getvalue()
+
+        self._test_dispatch_table(dumps, dt)
+
+    def test_instance_dispatch_table(self):
+        # A dispatch_table attribute can also be specified instance-wide
+        dt = self.get_dispatch_table()
+
+        def dumps(obj, protocol=None):
+            f = io.BytesIO()
+            p = self.pickler_class(f, protocol)
+            p.dispatch_table = dt
+            self.assertEqual(p.dispatch_table, dt)
+            p.dump(obj)
+            return f.getvalue()
+
+        self._test_dispatch_table(dumps, dt)
+
+    def _test_dispatch_table(self, dumps, dispatch_table):
+        def custom_load_dump(obj):
+            return pickle.loads(dumps(obj, 0))
+
+        def default_load_dump(obj):
+            return pickle.loads(pickle.dumps(obj, 0))
+
+        # pickling complex numbers using protocol 0 relies on copyreg
+        # so check pickling a complex number still works
+        z = 1 + 2j
+        self.assertEqual(custom_load_dump(z), z)
+        self.assertEqual(default_load_dump(z), z)
+
+        # modify pickling of complex
+        REDUCE_1 = 'reduce_1'
+        def reduce_1(obj):
+            return str, (REDUCE_1,)
+        dispatch_table[complex] = reduce_1
+        self.assertEqual(custom_load_dump(z), REDUCE_1)
+        self.assertEqual(default_load_dump(z), z)
+
+        # check picklability of AAA and BBB
+        a = AAA()
+        b = BBB()
+        self.assertEqual(custom_load_dump(a), REDUCE_A)
+        self.assertIsInstance(custom_load_dump(b), BBB)
+        self.assertEqual(default_load_dump(a), REDUCE_A)
+        self.assertIsInstance(default_load_dump(b), BBB)
+
+        # modify pickling of BBB
+        dispatch_table[BBB] = reduce_1
+        self.assertEqual(custom_load_dump(a), REDUCE_A)
+        self.assertEqual(custom_load_dump(b), REDUCE_1)
+        self.assertEqual(default_load_dump(a), REDUCE_A)
+        self.assertIsInstance(default_load_dump(b), BBB)
+
+        # revert pickling of BBB and modify pickling of AAA
+        REDUCE_2 = 'reduce_2'
+        def reduce_2(obj):
+            return str, (REDUCE_2,)
+        dispatch_table[AAA] = reduce_2
+        del dispatch_table[BBB]
+        self.assertEqual(custom_load_dump(a), REDUCE_2)
+        self.assertIsInstance(custom_load_dump(b), BBB)
+        self.assertEqual(default_load_dump(a), REDUCE_A)
+        self.assertIsInstance(default_load_dump(b), BBB)
+
+
 if __name__ == "__main__":
     # Print some stuff that can be used to rewrite DATA{0,1,2}
     from pickletools import dis

Lib/test/test_pickle.py

 import pickle
 import io
+import collections
 
 from test import support
 
@@ -7,6 +8,7 @@ from test.pickletester import AbstractPickleTests
 from test.pickletester import AbstractPickleModuleTests
 from test.pickletester import AbstractPersistentPicklerTests
 from test.pickletester import AbstractPicklerUnpicklerObjectTests
+from test.pickletester import AbstractDispatchTableTests
 from test.pickletester import BigmemPickleTests
 
 try:
@@ -80,6 +82,18 @@ class PyPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
     unpickler_class = pickle._Unpickler
 
 
+class PyDispatchTableTests(AbstractDispatchTableTests):
+    pickler_class = pickle._Pickler
+    def get_dispatch_table(self):
+        return pickle.dispatch_table.copy()
+
+
+class PyChainDispatchTableTests(AbstractDispatchTableTests):
+    pickler_class = pickle._Pickler
+    def get_dispatch_table(self):
+        return collections.ChainMap({}, pickle.dispatch_table)
+
+
 if has_c_implementation:
     class CPicklerTests(PyPicklerTests):
         pickler = _pickle.Pickler
@@ -101,14 +115,26 @@ if has_c_implementation:
         pickler_class = _pickle.Pickler
         unpickler_class = _pickle.Unpickler
 
+    class CDispatchTableTests(AbstractDispatchTableTests):
+        pickler_class = pickle.Pickler
+        def get_dispatch_table(self):
+            return pickle.dispatch_table.copy()
+
+    class CChainDispatchTableTests(AbstractDispatchTableTests):
+        pickler_class = pickle.Pickler
+        def get_dispatch_table(self):
+            return collections.ChainMap({}, pickle.dispatch_table)
+
 
 def test_main():
-    tests = [PickleTests, PyPicklerTests, PyPersPicklerTests]
+    tests = [PickleTests, PyPicklerTests, PyPersPicklerTests,
+             PyDispatchTableTests, PyChainDispatchTableTests]
     if has_c_implementation:
         tests.extend([CPicklerTests, CPersPicklerTests,
                       CDumpPickle_LoadPickle, DumpPickle_CLoadPickle,
                       PyPicklerUnpicklerObjectTests,
                       CPicklerUnpicklerObjectTests,
+                      CDispatchTableTests, CChainDispatchTableTests,
                       InMemoryPickleTests])
     support.run_unittest(*tests)
     support.run_doctest(pickle)

Misc/NEWS

@@ -511,6 +511,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #14166: Pickler objects now have an optional ``dispatch_table``
+  attribute which allows to set custom per-pickler reduction functions.
+  Patch by sbt.
+
 - Issue #14177: marshal.loads() now raises TypeError when given an unicode
   string.  Patch by Guilherme Gonçalves.
 

Modules/_pickle.c

@@ -319,6 +319,7 @@ typedef struct PicklerObject {
                                    objects to support self-referential objects
                                    pickling. */
     PyObject *pers_func;        /* persistent_id() method, can be NULL */
+    PyObject *dispatch_table;   /* private dispatch_table, can be NULL */
     PyObject *arg;
 
     PyObject *write;            /* write() method of the output stream. */
@@ -764,6 +765,7 @@ _Pickler_New(void)
         return NULL;
 
     self->pers_func = NULL;
+    self->dispatch_table = NULL;
     self->arg = NULL;
     self->write = NULL;
     self->proto = 0;
@@ -3176,17 +3178,24 @@ save(PicklerObject *self, PyObject *obj, int pers_save)
     /* XXX: This part needs some unit tests. */
 
     /* Get a reduction callable, and call it.  This may come from
-     * copyreg.dispatch_table, the object's __reduce_ex__ method,
-     * or the object's __reduce__ method.
+     * self.dispatch_table, copyreg.dispatch_table, the object's
+     * __reduce_ex__ method, or the object's __reduce__ method.
      */
-    reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
+    if (self->dispatch_table == NULL) {
+        reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
+        /* PyDict_GetItem() unlike PyObject_GetItem() and
+           PyObject_GetAttr() returns a borrowed ref */
+        Py_XINCREF(reduce_func);
+    } else {
+        reduce_func = PyObject_GetItem(self->dispatch_table, (PyObject *)type);
+        if (reduce_func == NULL) {
+            if (PyErr_ExceptionMatches(PyExc_KeyError))
+                PyErr_Clear();
+            else
+                goto error;
+        }
+    }
     if (reduce_func != NULL) {
-        /* Here, the reference count of the reduce_func object returned by
-           PyDict_GetItem needs to be increased to be consistent with the one
-           returned by PyObject_GetAttr. This is allow us to blindly DECREF
-           reduce_func at the end of the save() routine.
-        */
-        Py_INCREF(reduce_func);
         Py_INCREF(obj);
         reduce_value = _Pickler_FastCall(self, reduce_func, obj);
     }
@@ -3359,6 +3368,7 @@ Pickler_dealloc(PicklerObject *self)
     Py_XDECREF(self->output_buffer);
     Py_XDECREF(self->write);
     Py_XDECREF(self->pers_func);
+    Py_XDECREF(self->dispatch_table);
     Py_XDECREF(self->arg);
     Py_XDECREF(self->fast_memo);
 
@@ -3372,6 +3382,7 @@ Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
 {
     Py_VISIT(self->write);
     Py_VISIT(self->pers_func);
+    Py_VISIT(self->dispatch_table);
     Py_VISIT(self->arg);
     Py_VISIT(self->fast_memo);
     return 0;
@@ -3383,6 +3394,7 @@ Pickler_clear(PicklerObject *self)
     Py_CLEAR(self->output_buffer);
     Py_CLEAR(self->write);
     Py_CLEAR(self->pers_func);
+    Py_CLEAR(self->dispatch_table);
     Py_CLEAR(self->arg);
     Py_CLEAR(self->fast_memo);
 
@@ -3427,6 +3439,7 @@ Pickler_init(PicklerObject *self, PyObject *args, PyObject *kwds)
     PyObject *proto_obj = NULL;
     PyObject *fix_imports = Py_True;
     _Py_IDENTIFIER(persistent_id);
+    _Py_IDENTIFIER(dispatch_table);
 
     if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:Pickler",
                                      kwlist, &file, &proto_obj, &fix_imports))
@@ -3468,6 +3481,13 @@ Pickler_init(PicklerObject *self, PyObject *args, PyObject *kwds)
         if (self->pers_func == NULL)
             return -1;
     }
+    self->dispatch_table = NULL;
+    if (_PyObject_HasAttrId((PyObject *)self, &PyId_dispatch_table)) {
+        self->dispatch_table = _PyObject_GetAttrId((PyObject *)self,
+                                                   &PyId_dispatch_table);
+        if (self->dispatch_table == NULL)
+            return -1;
+    }
     return 0;
 }
 
@@ -3749,6 +3769,7 @@ Pickler_set_persid(PicklerObject *self, PyObject *value)
 static PyMemberDef Pickler_members[] = {
     {"bin", T_INT, offsetof(PicklerObject, bin)},
     {"fast", T_INT, offsetof(PicklerObject, fast)},
+    {"dispatch_table", T_OBJECT_EX, offsetof(PicklerObject, dispatch_table)},
     {NULL}
 };