[3.15] gh-157242: Add support.inject_memory_error() function

Lib/test/support/__init__.py

@@ -3455,6 +3455,22 @@ def built_with_c_assertions():
         return bool(_testlimitedcapi._py_getbuiltwithassert())
 
 
+def inject_memory_error(start=0, stop=0):
+    """
+    Memory allocation fails after 'start' allocation requests, and until 'stop'
+    allocation requests except when 'stop' is negative or equal to 0 (default)
+    in which case allocation failures never stop.
+
+    Raise SkipTest if the _testcapi extension module is missing
+    """
+    try:
+        import _testcapi
+    except ImportError:
+        raise unittest.SkipTest("_testcapi required")
+
+    _testcapi.set_nomemory(start, stop)
+
+
 @contextlib.contextmanager
 def inject_memory_error_cm(start=0, stop=0):
     """

Lib/test/test_atexit.py

@@ -197,14 +197,14 @@ def test_atexit_with_low_memory(self):
         # callback doesn't cause an infinite loop during finalization.
         code = textwrap.dedent("""
             import atexit
-            import _testcapi
+            from test.support import inject_memory_error
 
             def callback():
                 print("hello")
 
             atexit.register(callback)
             # Simulate low memory condition
-            _testcapi.set_nomemory(0)
+            inject_memory_error()
         """)
 
         with os_helper.temp_dir() as temp_dir:

Lib/test/test_bytes.py

@@ -51,16 +51,10 @@ def __index__(self):
 
 
 @contextlib.contextmanager
-def inject_memory_error(testcase, start):
-    # Raise SkipTest if _testcapi extension module is missing
-    _testcapi = import_helper.import_module('_testcapi')
-
+def inject_memory_error(testcase, start=0):
     with testcase.assertRaises(MemoryError):
-        try:
-            _testcapi.set_nomemory(start)
+        with support.inject_memory_error_cm(start):
             yield
-        finally:
-            _testcapi.remove_mem_hooks()
 
 
 class BaseBytesTest:
@@ -1607,7 +1601,7 @@ def test_resize_error(self):
                     del ba[:offset]
                 else:
                     expected = ba.copy()
-                with inject_memory_error(self, 0):
+                with inject_memory_error(self):
                     ba.resize(1024)
                 self.assertEqual(ba, expected)
 
@@ -1618,7 +1612,7 @@ def test_resize_error(self):
                     del ba[:offset]
                 else:
                     expected = ba.copy()
-                with inject_memory_error(self, 0):
+                with inject_memory_error(self):
                     ba.resize(1)
                 self.assertEqual(ba, expected)
 
@@ -1691,21 +1685,20 @@ def test_take_bytes_error(self):
         # gh-157242: If bytearray.take_bytes() fails (MemoryError),
         # the bytearray must be left unchanged.
 
-        for logical_offset, to_take, mem_errors in (
+        for logical_offset, to_take, start_list in (
             (True, 5, (0, 1)),
             (False, 5, (0, 1)),
             (True, None, (0,)),
         ):
-            for mem_error in mem_errors:
-                with self.subTest(logical_offset=logical_offset,
-                                  to_take=to_take, mem_error=mem_error):
+            for start in start_list:
+                with self.subTest(logical_offset=logical_offset, start=start):
                     ba = bytearray(b'0123456789')
                     if logical_offset:
                         expected = ba[3:]
                         del ba[:3]
                     else:
                         expected = ba.copy()
-                    with inject_memory_error(self, mem_error):
+                    with inject_memory_error(self, start):
                         ba.take_bytes(to_take)
                     self.assertEqual(ba, expected)
 

Lib/test/test_capi/test_bytes.py

@@ -455,6 +455,23 @@ def test_resize(self):
                     writer.resize(_testcapi.PY_SSIZE_T_MAX)
                 self.assertEqual(writer.finish(), b'x' * size)
 
+    @support.nomemtest
+    def test_resize_error(self):
+        # Test PyBytesWriter_Resize() error
+        init = b'x' * self.LARGE_BUFFER
+        writer = self.create_writer(len(init))
+        writer.write(0, init)
+        size = len(init) + 100
+        with self.assertRaises(MemoryError):
+            with support.inject_memory_error_cm():
+                writer.resize(size)
+        suffix = b'still working'
+        writer.write_bytes(suffix, -1)
+        self.assertEqual(writer.finish(), init + suffix)
+
+        # Note: PyBytesWriter_Resize() leaves the buffer unchanged (no resize)
+        # if the new size is smaller than the allocated size
+
     def test_grow(self):
         # Test PyBytesWriter_Grow()
         writer = self.create_writer(0)
@@ -513,24 +530,20 @@ def test_grow(self):
                 self.assertEqual(writer.finish(), b'x' * size)
 
     @support.nomemtest
-    def test_resize_error(self):
-        # Test PyBytesWriter_Resize() error
+    def test_grow_error(self):
+        # Test PyBytesWriter_Grow() error
         init = b'x' * self.LARGE_BUFFER
         writer = self.create_writer(len(init))
         writer.write(0, init)
-        size = len(init) + 100
-        try:
-            with self.assertRaises(MemoryError):
-                _testcapi.set_nomemory(0)
-                writer.resize(size)
-        finally:
-            _testcapi.remove_mem_hooks()
+        with self.assertRaises(MemoryError):
+            with support.inject_memory_error_cm():
+                writer.grow(100)
         suffix = b'still working'
         writer.write_bytes(suffix, -1)
         self.bytes_equal(writer.finish(), init + suffix)
 
-        # Note: PyBytesWriter_Resize() leaves the buffer unchanged (no resize)
-        # if the new size is smaller than the allocated size
+        # Note: PyBytesWriter_Grow() leaves the buffer unchanged (no resize)
+        # if grow is negative.
 
     def test_format_i(self):
         # Test PyBytesWriter_Format()
@@ -554,11 +567,8 @@ def test_format_s_memory_error(self):
         writer = self.create_writer()
         s = b'x' * 500
         with self.assertRaises(MemoryError):
-            try:
-                _testcapi.set_nomemory(0, 0)
+            with support.inject_memory_error_cm():
                 writer.format_s(b's=%s', s)
-            finally:
-                _testcapi.remove_mem_hooks()
         writer.write_bytes(b'after', -1)
 
     @threading_helper.requires_working_threading()

Lib/test/test_class.py

@@ -1016,8 +1016,6 @@ class C:
     @support.nomemtest
     @isolation.runInSubprocess()
     def test_detach_materialized_dict_no_memory(self):
-        import _testcapi
-
         class A:
             def __init__(self):
                 self.a = 1
@@ -1032,11 +1030,8 @@ def __init__(self):
             d = a.__dict__
             try:
                 with support.catch_unraisable_exception() as ex:
-                    _testcapi.set_nomemory(n, n + 1)
-                    try:
+                    with support.inject_memory_error_cm(n, n + 1):
                         del a
-                    finally:
-                        _testcapi.remove_mem_hooks()
                     exc_type = ex.unraisable and ex.unraisable.exc_type
             except MemoryError:
                 # The failing allocation was not in the deallocation code.

Lib/test/test_dict.py

@@ -461,15 +461,11 @@ def test_copy_noncompact(self):
 
     @support.nomemtest
     def test_copy_empty_no_memory(self):
-        import _testcapi
         d = {}
         d_copy = d.copy
         with self.assertRaises(MemoryError):
-            try:
-                _testcapi.set_nomemory(0)
+            with support.inject_memory_error_cm():
                 d_copy()
-            finally:
-                _testcapi.remove_mem_hooks()
 
     def test_get(self):
         d = {}

Lib/test/test_exceptions.py

@@ -1663,14 +1663,14 @@ def test_recursion_normalizing_with_no_memory(self):
         # the size of the list of preallocated MemoryError instances, the
         # Fatal Python error message mentions MemoryError.
         code = """if 1:
-            import _testcapi
+            from test import support
             class C(): pass
             def recurse(cnt):
                 cnt -= 1
                 if cnt:
                     recurse(cnt)
                 else:
-                    _testcapi.set_nomemory(0)
+                    support.inject_memory_error()
                     C()
             recurse(16)
         """
@@ -1846,9 +1846,10 @@ def test_unhandled(self):
     @support.nomemtest
     def test_memory_error_in_PyErr_PrintEx(self):
         code = """if 1:
-            import _testcapi
+            from test import support
+            stop = %d
             class C(): pass
-            _testcapi.set_nomemory(0, %d)
+            support.inject_memory_error(0, stop)
             C()
         """
 
@@ -2013,8 +2014,8 @@ def test_exec_set_nomemory_hang(self):
         warmup_code = "a = list(range(0, 1))\n" * 60
         user_input = warmup_code + dedent("""
             try:
-                import _testcapi
-                _testcapi.set_nomemory(0)
+                from test import support
+                support.inject_memory_error()
                 b = list(range(1000, 2000))
             except Exception as e:
                 import traceback

Lib/test/test_interpreters/test_stress.py

@@ -78,15 +78,9 @@ def run():
 
     @support.nomemtest
     def test_create_interpreter_no_memory(self):
-        import _testcapi
-
-        assertion = self.assertRaises(InterpreterError)
-        try:
-            _testcapi.set_nomemory(0, 1)
-            with assertion:
+        with self.assertRaises(InterpreterError):
+            with support.inject_memory_error_cm(0, 1):
                 _interpreters.create()
-        finally:
-            _testcapi.remove_mem_hooks()
 
 
 if __name__ == '__main__':

Lib/test/test_io/test_memoryio.py

@@ -5,7 +5,6 @@
 
 import unittest
 from test import support
-from test.support import import_helper
 
 import gc
 import io
@@ -654,18 +653,14 @@ def __buffer__(self, flags):
 
     def test_memory_error(self):
         # gh-157242: io.BytesIO() must not close the file on MemoryError
-        _testcapi = import_helper.import_module('_testcapi')
 
         # write()
         stream = self.ioclass()
         stream.write(self.buftype('abc'))
+        data = self.buftype('def')
         with self.assertRaises(MemoryError):
-            try:
-                data = self.buftype('def')
-                _testcapi.set_nomemory(0)
+            with support.inject_memory_error_cm():
                 stream.write(data)
-            finally:
-                _testcapi.remove_mem_hooks()
         stream.write(self.buftype('123'))
         self.assertEqual(stream.getvalue(), self.buftype('abc123'))
 
@@ -674,11 +669,8 @@ def test_memory_error(self):
         stream = self.ioclass()
         stream.write(data)
         with self.assertRaises(MemoryError):
-            try:
-                _testcapi.set_nomemory(0)
+            with support.inject_memory_error_cm():
                 stream.truncate(5)
-            finally:
-                _testcapi.remove_mem_hooks()
         self.assertEqual(stream.getvalue(), data)
 
 

Lib/test/test_list.py

@@ -329,11 +329,12 @@ def test_tier2_invalidates_iterator(self):
     def test_no_memory(self):
         # gh-118331: Make sure we don't crash if list allocation fails
         code = textwrap.dedent("""
-        import _testcapi, sys
+        from test import support
+        import sys
         # Prime the freelist
         l = [None]
         del l
-        _testcapi.set_nomemory(0)
+        support.inject_memory_error()
         l = [None]
         """)
         rc, _, _ = assert_python_failure("-c", code)

Lib/test/test_pyexpat.py

@@ -1028,18 +1028,9 @@ def test_error_path_no_crash(self):
         parser.buffer_text = True
         rc_before = sys.getrefcount(parser)
 
-        # We avoid self.assertRaises(MemoryError) here because the
-        # context manager itself needs memory allocations that fail
-        # while the nomemory hook is active.
-        self.testcapi.set_nomemory(1, 10)
-        raised = False
-        try:
-            parser.ExternalEntityParserCreate(None)
-        except MemoryError:
-            raised = True
-        finally:
-            self.testcapi.remove_mem_hooks()
-        self.assertTrue(raised, "MemoryError not raised")
+        with self.assertRaises(MemoryError):
+            with support.inject_memory_error_cm(1, 10):
+                parser.ExternalEntityParserCreate(None)
 
         rc_after = sys.getrefcount(parser)
         self.assertEqual(rc_after, rc_before)

Lib/test/test_repl.py

@@ -104,10 +104,11 @@ def test_no_memory(self):
         # no memory. Check also that the fix does not break the interactive
         # loop when an exception is raised.
         user_input = """
-            import sys, _testcapi
+            import sys
+            from test import support
             1/0
             print('After the exception.')
-            _testcapi.set_nomemory(0)
+            support.inject_memory_error()
             sys.exit(0)
         """
         user_input = dedent(user_input)

Lib/test/test_str.py

@@ -613,14 +613,9 @@ def test_replace_oom(self):
         s1 = "轘" * 4
         s2 = "&"
         s3 = "&"
-        assertion = self.assertRaises(MemoryError)
-        _testcapi.set_nomemory(0, 0)
-        try:
-            # No allocations made in the test itself:
-            with assertion:
+        with self.assertRaises(MemoryError):
+            with support.inject_memory_error_cm():
                 s1.replace(s2, s3)  # this line used to crash before
-        finally:
-            _testcapi.remove_mem_hooks()
 
     def test_repeat_id_preserving(self):
         a = '123abc1@'

Lib/test/test_weakref.py

@@ -1035,8 +1035,8 @@ def test_no_memory_when_clearing(self):
         # gh-118331: Make sure we do not raise an exception from the destructor
         # when clearing weakrefs if allocating the intermediate tuple fails.
         code = textwrap.dedent("""
-        import _testcapi
         import weakref
+        from test import support
 
         class TestObj:
             pass
@@ -1048,7 +1048,7 @@ def callback(obj):
         # The choice of 50 is arbitrary, but must be large enough to ensure
         # the allocation won't be serviced by the free list.
         wrs = [weakref.ref(obj, callback) for _ in range(50)]
-        _testcapi.set_nomemory(0)
+        support.inject_memory_error()
         del obj
         """).strip()
         res, _ = script_helper.run_python_until_end("-c", code)