|
17 | 17 | from test import support |
18 | 18 | from test.support import cpython_only, import_helper, os_helper |
19 | 19 | from test.support.import_helper import ensure_lazy_imports |
| 20 | +from test.support.script_helper import assert_python_ok |
20 | 21 |
|
21 | 22 | from test.pickletester import AbstractHookTests |
22 | 23 | from test.pickletester import AbstractUnpickleTests |
@@ -455,6 +456,70 @@ class CPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests, unittest |
455 | 456 | pickler_class = _pickle.Pickler |
456 | 457 | unpickler_class = _pickle.Unpickler |
457 | 458 |
|
| 459 | + @support.requires_subprocess() |
| 460 | + def test_clear_unpickler_memo_memory_error(self): |
| 461 | + import_helper.import_module('_testcapi') |
| 462 | + code = dedent(''' |
| 463 | + import _testcapi |
| 464 | + import io |
| 465 | + import pickle |
| 466 | + from test.support import SuppressCrashReport |
| 467 | +
|
| 468 | + value = [b'repeated'] * 2 |
| 469 | + data = pickle.dumps(value, protocol=4) |
| 470 | + unpickler = pickle.Unpickler(io.BytesIO(data * 2)) |
| 471 | + assert unpickler.load() == value |
| 472 | + memo = unpickler.memo.copy() |
| 473 | + clear = unpickler.memo.clear |
| 474 | + reset = _testcapi.remove_mem_hooks |
| 475 | + failed = False |
| 476 | + with SuppressCrashReport(): |
| 477 | + # Fail the replacement memo allocation, with all callable |
| 478 | + # lookups performed before enabling the allocator hook. |
| 479 | + _testcapi.set_nomemory(0, 1) |
| 480 | + try: |
| 481 | + clear() |
| 482 | + except MemoryError: |
| 483 | + failed = True |
| 484 | + finally: |
| 485 | + reset() |
| 486 | + assert failed |
| 487 | + assert unpickler.memo.copy() == memo |
| 488 | + clear() |
| 489 | + assert unpickler.memo.copy() == {} |
| 490 | + assert unpickler.load() == value |
| 491 | + ''') |
| 492 | + assert_python_ok('-c', code) |
| 493 | + |
| 494 | + @support.requires_subprocess() |
| 495 | + def test_clear_unpickler_memo_reentrant(self): |
| 496 | + code = dedent(r''' |
| 497 | + import io |
| 498 | + import pickle |
| 499 | + from test.support import gc_collect, SuppressCrashReport |
| 500 | +
|
| 501 | + snapshots = [] |
| 502 | + class Entry: |
| 503 | + def __del__(self): |
| 504 | + snapshots.append(unpickler.memo.copy()) |
| 505 | +
|
| 506 | + value = [b'repeated'] * 2 |
| 507 | + # Create the entry through the C unpickler so clearing its |
| 508 | + # sole reference runs the finalizer even without the GIL. |
| 509 | + data = b'\x80\x02c__main__\nEntry\n)\x81q\x00.' |
| 510 | + unpickler = pickle.Unpickler( |
| 511 | + io.BytesIO(data + pickle.dumps(value, protocol=4))) |
| 512 | + unpickler.load() |
| 513 | + gc_collect() |
| 514 | + with SuppressCrashReport(): |
| 515 | + unpickler.memo.clear() |
| 516 | + gc_collect() |
| 517 | + assert snapshots == [{}] |
| 518 | + assert unpickler.memo.copy() == {} |
| 519 | + assert unpickler.load() == value |
| 520 | + ''') |
| 521 | + assert_python_ok('-c', code) |
| 522 | + |
458 | 523 | def test_issue18339(self): |
459 | 524 | unpickler = self.unpickler_class(io.BytesIO()) |
460 | 525 | with self.assertRaises(TypeError): |
|
0 commit comments