bpo-35900: Enable custom reduction callback registration in _pickle (GH-12499)

Enable custom reduction callback registration for functions and classes in
_pickle.c, using the new Pickler's attribute ``reducer_override``.

1 files changed, 71 insertions, 0 deletions

diff --git a/Doc/library/pickle.rst b/Doc/library/pickle.rst
index 3d89536..55005f0 100644
--- a/Doc/library/pickle.rst
+++ b/Doc/library/pickle.rst
@@ -356,6 +356,18 @@ The :mod:`pickle` module exports two classes, :class:`Pickler` and
 
       .. versionadded:: 3.3
 
+   .. method:: reducer_override(self, obj)
+
+      Special reducer that can be defined in :class:`Pickler` subclasses. This
+      method has priority over any reducer in the :attr:`dispatch_table`.  It
+      should conform to the same interface as a :meth:`__reduce__` method, and
+      can optionally return ``NotImplemented`` to fallback on
+      :attr:`dispatch_table`-registered reducers to pickle ``obj``.
+
+      For a detailed example, see :ref:`reducer_override`.
+
+      .. versionadded:: 3.8
+
    .. attribute:: fast
 
       Deprecated. Enable fast mode if set to a true value.  The fast mode
@@ -791,6 +803,65 @@ A sample usage might be something like this::
    >>> new_reader.readline()
    '3: Goodbye!'
 
+.. _reducer_override:
+
+Custom Reduction for Types, Functions, and Other Objects
+--------------------------------------------------------
+
+.. versionadded:: 3.8
+
+Sometimes, :attr:`~Pickler.dispatch_table` may not be flexible enough.
+In particular we may want to customize pickling based on another criterion
+than the object's type, or we may want to customize the pickling of
+functions and classes.
+
+For those cases, it is possible to subclass from the :class:`Pickler` class and
+implement a :meth:`~Pickler.reducer_override` method. This method can return an
+arbitrary reduction tuple (see :meth:`__reduce__`). It can alternatively return
+``NotImplemented`` to fallback to the traditional behavior.
+
+If both the :attr:`~Pickler.dispatch_table` and
+:meth:`~Pickler.reducer_override` are defined, then
+:meth:`~Pickler.reducer_override` method takes priority.
+
+.. Note::
+   For performance reasons, :meth:`~Pickler.reducer_override` may not be
+   called for the following objects: ``None``, ``True``, ``False``, and
+   exact instances of :class:`int`, :class:`float`, :class:`bytes`,
+   :class:`str`, :class:`dict`, :class:`set`, :class:`frozenset`, :class:`list`
+   and :class:`tuple`.
+
+Here is a simple example where we allow pickling and reconstructing
+a given class::
+
+   import io
+   import pickle
+
+   class MyClass:
+       my_attribute = 1
+
+   class MyPickler(pickle.Pickler):
+       def reducer_override(self, obj):
+           """Custom reducer for MyClass."""
+           if getattr(obj, "__name__", None) == "MyClass":
+               return type, (obj.__name__, obj.__bases__,
+                             {'my_attribute': obj.my_attribute})
+           else:
+               # For any other object, fallback to usual reduction
+               return NotImplemented
+
+   f = io.BytesIO()
+   p = MyPickler(f)
+   p.dump(MyClass)
+
+   del MyClass
+
+   unpickled_class = pickle.loads(f.getvalue())
+
+   assert isinstance(unpickled_class, type)
+   assert unpickled_class.__name__ == "MyClass"
+   assert unpickled_class.my_attribute == 1
+
 
 .. _pickle-restrict: