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, 68 insertions, 0 deletions

diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py
index 19e8823a731..4f8c2942df9 100644
--- a/Lib/test/pickletester.py
+++ b/Lib/test/pickletester.py
@@ -4,6 +4,7 @@ import dbm
 import io
 import functools
 import os
+import math
 import pickle
 import pickletools
 import shutil
@@ -3013,6 +3014,73 @@ def setstate_bbb(obj, state):
     obj.a = "custom state_setter"
 
 
+
+class AbstractCustomPicklerClass:
+    """Pickler implementing a reducing hook using reducer_override."""
+    def reducer_override(self, obj):
+        obj_name = getattr(obj, "__name__", None)
+
+        if obj_name == 'f':
+            # asking the pickler to save f as 5
+            return int, (5, )
+
+        if obj_name == 'MyClass':
+            return str, ('some str',)
+
+        elif obj_name == 'g':
+            # in this case, the callback returns an invalid result (not a 2-5
+            # tuple or a string), the pickler should raise a proper error.
+            return False
+
+        elif obj_name == 'h':
+            # Simulate a case when the reducer fails. The error should
+            # be propagated to the original ``dump`` call.
+            raise ValueError('The reducer just failed')
+
+        return NotImplemented
+
+class AbstractHookTests(unittest.TestCase):
+    def test_pickler_hook(self):
+        # test the ability of a custom, user-defined CPickler subclass to
+        # override the default reducing routines of any type using the method
+        # reducer_override
+
+        def f():
+            pass
+
+        def g():
+            pass
+
+        def h():
+            pass
+
+        class MyClass:
+            pass
+
+        for proto in range(0, pickle.HIGHEST_PROTOCOL + 1):
+            with self.subTest(proto=proto):
+                bio = io.BytesIO()
+                p = self.pickler_class(bio, proto)
+
+                p.dump([f, MyClass, math.log])
+                new_f, some_str, math_log = pickle.loads(bio.getvalue())
+
+                self.assertEqual(new_f, 5)
+                self.assertEqual(some_str, 'some str')
+                # math.log does not have its usual reducer overriden, so the
+                # custom reduction callback should silently direct the pickler
+                # to the default pickling by attribute, by returning
+                # NotImplemented
+                self.assertIs(math_log, math.log)
+
+                with self.assertRaises(pickle.PicklingError):
+                    p.dump(g)
+
+                with self.assertRaisesRegex(
+                        ValueError, 'The reducer just failed'):
+                    p.dump(h)
+
+
 class AbstractDispatchTableTests(unittest.TestCase):
 
     def test_default_dispatch_table(self):