diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py index 6e87370c2065ba..316da00bab0283 100644 --- a/Lib/test/pickletester.py +++ b/Lib/test/pickletester.py @@ -2408,6 +2408,22 @@ def test_reduce_calls_base(self): y = self.loads(s) self.assertEqual(y._reduce_called, 1) + def test_reduce_ex_None(self): + c = REX_None() + with self.assertRaises(TypeError): + self.dumps(c) + + def test_reduce_None(self): + c = R_None() + with self.assertRaises(TypeError): + pickle.dumps(c) + + def test_pickle_setstate_None(self): + c = C_None_setstate() + p = self.dumps(c) + with self.assertRaises(TypeError): + self.loads(p) + @no_tracing def test_bad_getattr(self): # Issue #3514: crash when there is an infinite loop in __getattr__ @@ -3348,6 +3364,20 @@ def __setstate__(self, state): def __reduce__(self): return type(self), (), self.state +class REX_None: + """ Setting __reduce_ex__ to None should fail """ + __reduce_ex__ = None + +class R_None: + """ Setting __reduce__ to None should fail """ + __reduce__ = None + +class C_None_setstate: + def __getstate__(self): + return 1 + + __setstate__ = None + # Test classes for newobj @@ -3751,6 +3781,25 @@ def test_unpickling_buffering_readline(self): unpickler = self.unpickler_class(f) self.assertEqual(unpickler.load(), data) + def test_pickle_invalid_reducer_override(self): + # gh-103035 + obj=object() + + f = io.BytesIO() + class MyPickler(self.pickler_class): + pass + pickler = MyPickler(f) + pickler.dump(obj) + + pickler.clear_memo() + pickler.reducer_override = None + with self.assertRaises(TypeError): + pickler.dump(obj) + + pickler.clear_memo() + pickler.reducer_override = 10 + with self.assertRaises(TypeError): + pickler.dump(obj) # Tests for dispatch_table attribute @@ -3913,6 +3962,15 @@ def dumps(obj, protocol=None): self._test_dispatch_table(dumps, dt) + def test_dispatch_table_None_item(self): + # gh-93627 + obj=object() + f = io.BytesIO() + pickler = self.pickler_class(f) + pickler.dispatch_table = {type(obj): None} + with self.assertRaises(TypeError): + pickler.dump(obj) + def _test_dispatch_table(self, dumps, dispatch_table): def custom_load_dump(obj): return pickle.loads(dumps(obj, 0)) diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py index 07cd850ad10850..1a55da39bdc58d 100644 --- a/Lib/test/test_pickle.py +++ b/Lib/test/test_pickle.py @@ -570,87 +570,6 @@ def test_multiprocessing_exceptions(self): ('multiprocessing.context', name)) -class C_with_reduce_ex: - def __reduce_ex__(self, proto): - return object.__reduce_ex__(self, proto) - -class C_None_reduce_ex: - __reduce_ex__ = None - -class C_None_reduce: - __reduce__ = None - -class C_None_setstate: - def __getstate__(self): - return 1 - - __setstate__ = None - -class C_setstate: - def __init__(self): - self.value = 0 - - def __getstate__(self): - return self.value - - def __setstate__(self, state): - print(state) - self.value = state + 1 - -class PickleReductionMethodsTests(unittest.TestCase): - # see gh-93627 - - def test_pickle_invalid_reduction_method(self): - c = C_None_reduce_ex() - with self.assertRaises(TypeError): - pickle.dumps(c) - - c = C_None_reduce() - with self.assertRaises(TypeError): - pickle.dumps(c) - - def test_pickle_invalid_reducer_override(self): - obj=C_with_reduce_ex() - - bio = io.BytesIO() - pickler = pickle._Pickler(bio) - pickler.dump(obj) - - pickler.clear_memo() - pickler.reducer_override = None - with self.assertRaises(TypeError): - pickler.dump(obj) - - pickler.clear_memo() - pickler.reducer_override = 10 - with self.assertRaises(TypeError): - pickler.dump(obj) - - def test_pickle_invalid_dispatch_table(self): - obj=C_with_reduce_ex() - - bio = io.BytesIO() - pickler = pickle._Pickler(bio) - pickler.dispatch_table = {type(obj): None} - with self.assertRaises(TypeError): - pickler.dump(obj) - - -class PickleSetstateTests(unittest.TestCase): - # See gh-103035 - - def test_pickle_setstate(self): - c = C_setstate() - s = pickle.dumps(c) - c2 = pickle.loads(s) - self.assertEqual(c2.value, c.value + 1) - - c = C_None_setstate() - p = pickle.dumps(c) - with self.assertRaises(TypeError): - pickle.loads(p) - - def load_tests(loader, tests, pattern): tests.addTest(doctest.DocTestSuite()) return tests