Add CPython generator/contextlib tests (#150796)

Tests:
* test_generator.py
* test_generator_stop.py
* test_contextlib.py

Minor changes were made to each test to run them inside Dynamo. We
intentionally didn't copy the binary files stored in
`python/Lib/test/archivetestdata` for security reasons. There's a single
test that requires a binary file and it is skipped because of that.

The tests were downloaded from CPython 3.13 and the diff was generated
using `git diff` to apply the changes:

```bash
for f in "test_contextlib" "test_generators" "test_generator_stop"; do
	wget -O "test/dynamo/cpython/3_13/${f}.py" "https://raw.githubusercontent.com/python/cpython/refs/heads/3.13/Lib/test/${f}.py"
	git apply "test/dynamo/cpython/3_13/${f}.diff"
done
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/150796
Approved by: https://github.com/williamwen42

test/dynamo/cpython/3_13/data/README

@@ -0,0 +1,2 @@
+This empty directory serves as destination for temporary files
+created by some tests, in particular, the test_codecmaps_* tests.

test/dynamo/cpython/3_13/test_contextlib.diff

@@ -0,0 +1,195 @@
+diff --git a/test/dynamo/cpython/3_13/test_contextlib.py b/test/dynamo/cpython/3_13/test_contextlib.py
+index cf651959803..6a17bc719eb 100644
+--- a/test/dynamo/cpython/3_13/test_contextlib.py
++++ b/test/dynamo/cpython/3_13/test_contextlib.py
+@@ -1,3 +1,54 @@
++# ======= BEGIN Dynamo patch =======
++# Owner(s): ["module: dynamo"]
++
++# ruff: noqa
++# flake8: noqa
++
++import sys
++import torch
++import torch._dynamo.test_case
++import unittest
++from torch._dynamo.test_case import CPythonTestCase
++from torch.testing._internal.common_utils import run_tests
++
++__TestCase = CPythonTestCase
++
++
++# redirect import statements
++import sys
++import importlib.abc
++
++redirect_imports = (
++    "test.mapping_tests",
++    "test.typinganndata",
++    "test.test_grammar",
++    "test.test_math",
++    "test.test_iter",
++    "test.typinganndata.ann_module",
++)
++
++class RedirectImportFinder(importlib.abc.MetaPathFinder):
++    def find_spec(self, fullname, path, target=None):
++        # Check if the import is the problematic one
++        if fullname in redirect_imports:
++            try:
++                # Attempt to import the standalone module
++                name = fullname.removeprefix("test.")
++                r = importlib.import_module(name)
++                # Redirect the module in sys.modules
++                sys.modules[fullname] = r
++                # Return a module spec from the found module
++                return importlib.util.find_spec(name)
++            except ImportError:
++                return None
++        return None
++
++# Add the custom finder to sys.meta_path
++sys.meta_path.insert(0, RedirectImportFinder())
++
++
++# ======= END DYNAMO PATCH =======
++
+ """Unit tests for contextlib.py, and other context managers."""
+ 
+ import io
+@@ -14,7 +65,7 @@ from test.support.testcase import ExceptionIsLikeMixin
+ import weakref
+ 
+ 
+-class TestAbstractContextManager(unittest.TestCase):
++class TestAbstractContextManager(__TestCase):
+ 
+     def test_enter(self):
+         class DefaultEnter(AbstractContextManager):
+@@ -67,7 +118,7 @@ class TestAbstractContextManager(unittest.TestCase):
+         self.assertFalse(issubclass(NoExit, AbstractContextManager))
+ 
+ 
+-class ContextManagerTestCase(unittest.TestCase):
++class ContextManagerTestCase(__TestCase):
+ 
+     def test_contextmanager_plain(self):
+         state = []
+@@ -396,7 +447,7 @@ def woohoo():
+         self.assertEqual(depth, 0)
+ 
+ 
+-class ClosingTestCase(unittest.TestCase):
++class ClosingTestCase(__TestCase):
+ 
+     @support.requires_docstrings
+     def test_instance_docs(self):
+@@ -430,7 +481,7 @@ class ClosingTestCase(unittest.TestCase):
+         self.assertEqual(state, [1])
+ 
+ 
+-class NullcontextTestCase(unittest.TestCase):
++class NullcontextTestCase(__TestCase):
+     def test_nullcontext(self):
+         class C:
+             pass
+@@ -439,7 +490,7 @@ class NullcontextTestCase(unittest.TestCase):
+             self.assertIs(c_in, c)
+ 
+ 
+-class FileContextTestCase(unittest.TestCase):
++class FileContextTestCase(__TestCase):
+ 
+     def testWithOpen(self):
+         tfn = tempfile.mktemp()
+@@ -457,7 +508,7 @@ class FileContextTestCase(unittest.TestCase):
+         finally:
+             os_helper.unlink(tfn)
+ 
+-class LockContextTestCase(unittest.TestCase):
++class LockContextTestCase(__TestCase):
+ 
+     def boilerPlate(self, lock, locked):
+         self.assertFalse(locked())
+@@ -520,7 +571,7 @@ class mycontext(ContextDecorator):
+         return self.catch
+ 
+ 
+-class TestContextDecorator(unittest.TestCase):
++class TestContextDecorator(__TestCase):
+ 
+     @support.requires_docstrings
+     def test_instance_docs(self):
+@@ -680,7 +731,7 @@ class TestContextDecorator(unittest.TestCase):
+         self.assertEqual(state, [1, 'something else', 999])
+ 
+ 
+-class TestBaseExitStack:
++class _TestBaseExitStack:
+     exit_stack = None
+ 
+     @support.requires_docstrings
+@@ -1141,7 +1192,7 @@ class TestBaseExitStack:
+         self.assertIs(exc.__cause__, exc.__context__)
+ 
+ 
+-class TestExitStack(TestBaseExitStack, unittest.TestCase):
++class TestExitStack(_TestBaseExitStack, __TestCase):
+     exit_stack = ExitStack
+     callback_error_internal_frames = [
+         ('__exit__', 'raise exc'),
+@@ -1149,7 +1200,7 @@ class TestExitStack(TestBaseExitStack, unittest.TestCase):
+     ]
+ 
+ 
+-class TestRedirectStream:
++class _TestRedirectStream:
+ 
+     redirect_stream = None
+     orig_stream = None
+@@ -1206,19 +1257,19 @@ class TestRedirectStream:
+         self.assertEqual(s, "Hello World!\n")
+ 
+ 
+-class TestRedirectStdout(TestRedirectStream, unittest.TestCase):
++class TestRedirectStdout(_TestRedirectStream, __TestCase):
+ 
+     redirect_stream = redirect_stdout
+     orig_stream = "stdout"
+ 
+ 
+-class TestRedirectStderr(TestRedirectStream, unittest.TestCase):
++class TestRedirectStderr(_TestRedirectStream, __TestCase):
+ 
+     redirect_stream = redirect_stderr
+     orig_stream = "stderr"
+ 
+ 
+-class TestSuppress(ExceptionIsLikeMixin, unittest.TestCase):
++class TestSuppress(ExceptionIsLikeMixin, __TestCase):
+ 
+     @support.requires_docstrings
+     def test_instance_docs(self):
+@@ -1315,7 +1366,7 @@ class TestSuppress(ExceptionIsLikeMixin, unittest.TestCase):
+         )
+ 
+ 
+-class TestChdir(unittest.TestCase):
++class TestChdir(__TestCase):
+     def make_relative_path(self, *parts):
+         return os.path.join(
+             os.path.dirname(os.path.realpath(__file__)),
+@@ -1331,6 +1382,7 @@ class TestChdir(unittest.TestCase):
+             self.assertEqual(os.getcwd(), target)
+         self.assertEqual(os.getcwd(), old_cwd)
+ 
++    @unittest.skip("Missing archivetestdata")
+     def test_reentrant(self):
+         old_cwd = os.getcwd()
+         target1 = self.make_relative_path('data')
+@@ -1363,4 +1415,4 @@ class TestChdir(unittest.TestCase):
+ 
+ 
+ if __name__ == "__main__":
+-    unittest.main()
++    run_tests()

test/dynamo/cpython/3_13/test_generator_stop.diff

@@ -0,0 +1,74 @@
+diff --git a/test/dynamo/cpython/3_13/test_generator_stop.py b/test/dynamo/cpython/3_13/test_generator_stop.py
+index bc235ceb00e..cb2a85255cb 100644
+--- a/test/dynamo/cpython/3_13/test_generator_stop.py
++++ b/test/dynamo/cpython/3_13/test_generator_stop.py
+@@ -1,9 +1,60 @@
+ from __future__ import generator_stop
+ 
++# ======= BEGIN Dynamo patch =======
++# Owner(s): ["module: dynamo"]
++
++# ruff: noqa
++# flake8: noqa
++
++import sys
++import torch
++import torch._dynamo.test_case
++import unittest
++from torch._dynamo.test_case import CPythonTestCase
++from torch.testing._internal.common_utils import run_tests
++
++__TestCase = CPythonTestCase
++
++
++# redirect import statements
++import sys
++import importlib.abc
++
++redirect_imports = (
++    "test.mapping_tests",
++    "test.typinganndata",
++    "test.test_grammar",
++    "test.test_math",
++    "test.test_iter",
++    "test.typinganndata.ann_module",
++)
++
++class RedirectImportFinder(importlib.abc.MetaPathFinder):
++    def find_spec(self, fullname, path, target=None):
++        # Check if the import is the problematic one
++        if fullname in redirect_imports:
++            try:
++                # Attempt to import the standalone module
++                name = fullname.removeprefix("test.")
++                r = importlib.import_module(name)
++                # Redirect the module in sys.modules
++                sys.modules[fullname] = r
++                # Return a module spec from the found module
++                return importlib.util.find_spec(name)
++            except ImportError:
++                return None
++        return None
++
++# Add the custom finder to sys.meta_path
++sys.meta_path.insert(0, RedirectImportFinder())
++
++
++# ======= END DYNAMO PATCH =======
++
+ import unittest
+ 
+ 
+-class TestPEP479(unittest.TestCase):
++class TestPEP479(__TestCase):
+     def test_stopiteration_wrapping(self):
+         def f():
+             raise StopIteration
+@@ -30,5 +81,5 @@ class TestPEP479(unittest.TestCase):
+                       'were not properly set')
+ 
+ 
+-if __name__ == '__main__':
+-    unittest.main()
++if __name__ == "__main__":
++    run_tests()

test/dynamo/cpython/3_13/test_generator_stop.py

@@ -0,0 +1,85 @@
+from __future__ import generator_stop
+
+# ======= BEGIN Dynamo patch =======
+# Owner(s): ["module: dynamo"]
+
+# ruff: noqa
+# flake8: noqa
+
+import sys
+import torch
+import torch._dynamo.test_case
+import unittest
+from torch._dynamo.test_case import CPythonTestCase
+from torch.testing._internal.common_utils import run_tests
+
+__TestCase = CPythonTestCase
+
+
+# redirect import statements
+import sys
+import importlib.abc
+
+redirect_imports = (
+    "test.mapping_tests",
+    "test.typinganndata",
+    "test.test_grammar",
+    "test.test_math",
+    "test.test_iter",
+    "test.typinganndata.ann_module",
+)
+
+class RedirectImportFinder(importlib.abc.MetaPathFinder):
+    def find_spec(self, fullname, path, target=None):
+        # Check if the import is the problematic one
+        if fullname in redirect_imports:
+            try:
+                # Attempt to import the standalone module
+                name = fullname.removeprefix("test.")
+                r = importlib.import_module(name)
+                # Redirect the module in sys.modules
+                sys.modules[fullname] = r
+                # Return a module spec from the found module
+                return importlib.util.find_spec(name)
+            except ImportError:
+                return None
+        return None
+
+# Add the custom finder to sys.meta_path
+sys.meta_path.insert(0, RedirectImportFinder())
+
+
+# ======= END DYNAMO PATCH =======
+
+import unittest
+
+
+class TestPEP479(__TestCase):
+    def test_stopiteration_wrapping(self):
+        def f():
+            raise StopIteration
+        def g():
+            yield f()
+        with self.assertRaisesRegex(RuntimeError,
+                                    "generator raised StopIteration"):
+            next(g())
+
+    def test_stopiteration_wrapping_context(self):
+        def f():
+            raise StopIteration
+        def g():
+            yield f()
+
+        try:
+            next(g())
+        except RuntimeError as exc:
+            self.assertIs(type(exc.__cause__), StopIteration)
+            self.assertIs(type(exc.__context__), StopIteration)
+            self.assertTrue(exc.__suppress_context__)
+        else:
+            self.fail('__cause__, __context__, or __suppress_context__ '
+                      'were not properly set')
+
+
+if __name__ == "__main__":
+    run_tests()

test/dynamo/cpython/3_13/test_generators.diff

@@ -0,0 +1,289 @@
+diff --git a/test/dynamo/cpython/3_13/test_generators.py b/test/dynamo/cpython/3_13/test_generators.py
+index e48d79d34f4..40a02d644a9 100644
+--- a/test/dynamo/cpython/3_13/test_generators.py
++++ b/test/dynamo/cpython/3_13/test_generators.py
+@@ -1,3 +1,53 @@
++# ======= BEGIN Dynamo patch =======
++# Owner(s): ["module: dynamo"]
++
++# ruff: noqa
++# flake8: noqa
++
++import sys
++import torch
++import torch._dynamo.test_case
++import unittest
++from torch._dynamo.test_case import CPythonTestCase
++from torch.testing._internal.common_utils import run_tests
++
++__TestCase = CPythonTestCase
++
++# redirect import statements
++import sys
++import importlib.abc
++
++redirect_imports = (
++    "test.mapping_tests",
++    "test.typinganndata",
++    "test.test_grammar",
++    "test.test_math",
++    "test.test_iter",
++    "test.typinganndata.ann_module",
++)
++
++class RedirectImportFinder(importlib.abc.MetaPathFinder):
++    def find_spec(self, fullname, path, target=None):
++        # Check if the import is the problematic one
++        if fullname in redirect_imports:
++            try:
++                # Attempt to import the standalone module
++                name = fullname.removeprefix("test.")
++                r = importlib.import_module(name)
++                # Redirect the module in sys.modules
++                sys.modules[fullname] = r
++                # Return a module spec from the found module
++                return importlib.util.find_spec(name)
++            except ImportError:
++                return None
++        return None
++
++# Add the custom finder to sys.meta_path
++sys.meta_path.insert(0, RedirectImportFinder())
++
++
++# ======= END DYNAMO PATCH =======
++
+ import copy
+ import gc
+ import pickle
+@@ -22,7 +72,7 @@ except ImportError:
+ @unittest.skipUnless(_testcapi is not None and
+                      hasattr(_testcapi, "raise_SIGINT_then_send_None"),
+                      "needs _testcapi.raise_SIGINT_then_send_None")
+-class SignalAndYieldFromTest(unittest.TestCase):
++class SignalAndYieldFromTest(__TestCase):
+ 
+     def generator1(self):
+         return (yield from self.generator2())
+@@ -46,7 +96,7 @@ class SignalAndYieldFromTest(unittest.TestCase):
+         self.assertEqual(exc.value, "PASSED")
+ 
+ 
+-class FinalizationTest(unittest.TestCase):
++class FinalizationTest(__TestCase):
+ 
+     def test_frame_resurrect(self):
+         # A generator frame can be resurrected by a generator's finalization.
+@@ -113,7 +163,7 @@ class FinalizationTest(unittest.TestCase):
+             self.assertEqual(cm.exception.value, 2)
+ 
+ 
+-class GeneratorTest(unittest.TestCase):
++class GeneratorTest(__TestCase):
+ 
+     def test_name(self):
+         def func():
+@@ -246,8 +296,31 @@ class GeneratorTest(unittest.TestCase):
+         #This should not raise
+         loop()
+ 
++    @unittest.expectedFailure
++    def test_genexpr_only_calls_dunder_iter_once(self):
++
++        class Iterator:
++
++            def __init__(self):
++                self.val = 0
++
++            def __next__(self):
++                if self.val == 2:
++                    raise StopIteration
++                self.val += 1
++                return self.val
++
++            # No __iter__ method
++
++        class C:
++
++            def __iter__(self):
++                return Iterator()
++
++        self.assertEqual([1,2], list(i for i in C()))
++
+ 
+-class ModifyUnderlyingIterableTest(unittest.TestCase):
++class ModifyUnderlyingIterableTest(__TestCase):
+     iterables = [
+         range(0),
+         range(20),
+@@ -319,7 +392,7 @@ class ModifyUnderlyingIterableTest(unittest.TestCase):
+         self.process_tests(get_generator_genfunc)
+ 
+ 
+-class ExceptionTest(unittest.TestCase):
++class ExceptionTest(__TestCase):
+     # Tests for the issue #23353: check that the currently handled exception
+     # is correctly saved/restored in PyEval_EvalFrameEx().
+ 
+@@ -528,7 +601,7 @@ class ExceptionTest(unittest.TestCase):
+         self.assertEqual(cm.exception.value.value, 2)
+ 
+ 
+-class GeneratorCloseTest(unittest.TestCase):
++class GeneratorCloseTest(__TestCase):
+ 
+     def test_close_no_return_value(self):
+         def f():
+@@ -630,90 +703,7 @@ class GeneratorCloseTest(unittest.TestCase):
+         self.assertIsNone(f_wr())
+ 
+ 
+-# See https://github.com/python/cpython/issues/125723
+-class GeneratorDeallocTest(unittest.TestCase):
+-    def test_frame_outlives_generator(self):
+-        def g1():
+-            a = 42
+-            yield sys._getframe()
+-
+-        def g2():
+-            a = 42
+-            yield
+-
+-        def g3(obj):
+-            a = 42
+-            obj.frame = sys._getframe()
+-            yield
+-
+-        class ObjectWithFrame():
+-            def __init__(self):
+-                self.frame = None
+-
+-        def get_frame(index):
+-            if index == 1:
+-                return next(g1())
+-            elif index == 2:
+-                gen = g2()
+-                next(gen)
+-                return gen.gi_frame
+-            elif index == 3:
+-                obj = ObjectWithFrame()
+-                next(g3(obj))
+-                return obj.frame
+-            else:
+-                return None
+-
+-        for index in (1, 2, 3):
+-            with self.subTest(index=index):
+-                frame = get_frame(index)
+-                frame_locals = frame.f_locals
+-                self.assertIn('a', frame_locals)
+-                self.assertEqual(frame_locals['a'], 42)
+-
+-    def test_frame_locals_outlive_generator(self):
+-        frame_locals1 = None
+-
+-        def g1():
+-            nonlocal frame_locals1
+-            frame_locals1 = sys._getframe().f_locals
+-            a = 42
+-            yield
+-
+-        def g2():
+-            a = 42
+-            yield sys._getframe().f_locals
+-
+-        def get_frame_locals(index):
+-            if index == 1:
+-                nonlocal frame_locals1
+-                next(g1())
+-                return frame_locals1
+-            if index == 2:
+-                return next(g2())
+-            else:
+-                return None
+-
+-        for index in (1, 2):
+-            with self.subTest(index=index):
+-                frame_locals = get_frame_locals(index)
+-                self.assertIn('a', frame_locals)
+-                self.assertEqual(frame_locals['a'], 42)
+-
+-    def test_frame_locals_outlive_generator_with_exec(self):
+-        def g():
+-            a = 42
+-            yield locals(), sys._getframe().f_locals
+-
+-        locals_ = {'g': g}
+-        for i in range(10):
+-            exec("snapshot, live_locals = next(g())", locals=locals_)
+-            for l in (locals_['snapshot'], locals_['live_locals']):
+-                self.assertIn('a', l)
+-                self.assertEqual(l['a'], 42)
+-
+-
+-class GeneratorThrowTest(unittest.TestCase):
++class GeneratorThrowTest(__TestCase):
+ 
+     def test_exception_context_with_yield(self):
+         def f():
+@@ -812,7 +802,7 @@ class GeneratorThrowTest(unittest.TestCase):
+             gen.throw(ValueError)
+ 
+ 
+-class GeneratorStackTraceTest(unittest.TestCase):
++class GeneratorStackTraceTest(__TestCase):
+ 
+     def check_stack_names(self, frame, expected):
+         names = []
+@@ -861,7 +851,7 @@ class GeneratorStackTraceTest(unittest.TestCase):
+         self.check_yield_from_example(call_throw)
+ 
+ 
+-class YieldFromTests(unittest.TestCase):
++class YieldFromTests(__TestCase):
+     def test_generator_gi_yieldfrom(self):
+         def a():
+             self.assertEqual(inspect.getgeneratorstate(gen_b), inspect.GEN_RUNNING)
+@@ -2752,21 +2742,27 @@ test_generators just happened to be the test that drew these out.
+ 
+ """
+ 
+-__test__ = {"tut":      tutorial_tests,
+-            "pep":      pep_tests,
+-            "email":    email_tests,
+-            "fun":      fun_tests,
+-            "syntax":   syntax_tests,
+-            "conjoin":  conjoin_tests,
+-            "weakref":  weakref_tests,
+-            "coroutine":  coroutine_tests,
+-            "refleaks": refleaks_tests,
+-            }
+-
+-def load_tests(loader, tests, pattern):
+-    tests.addTest(doctest.DocTestSuite())
+-    return tests
++# __test__ = {"tut":      tutorial_tests,
++#             "pep":      pep_tests,
++#             "email":    email_tests,
++#             "fun":      fun_tests,
++#             "syntax":   syntax_tests,
++#             "conjoin":  conjoin_tests,
++#             "weakref":  weakref_tests,
++#             "coroutine":  coroutine_tests,
++#             "refleaks": refleaks_tests,
++#             }
++
++# def load_tests(loader, tests, pattern):
++#     # ======= BEGIN Dynamo patch =======
++#     suite = doctest.DocTestSuite()
++#     for test in suite:
++#         # Dynamically change base class
++#         test.__class__ = type(test.__class__.__name__, (__TestCase, test.__class__), {})
++#     tests.addTests(suite)
++#     # ======= END DYNAMO PATCH =======
++#     return tests
+ 
+ 
+ if __name__ == "__main__":
+-    unittest.main()
++    run_tests()