diff options
Diffstat (limited to 'Lib/test/pickletester.py')
-rw-r--r-- | Lib/test/pickletester.py | 68 |
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): |