from __future__ import generator_stop

# ======= BEGIN Dynamo patch =======

# Owner(s): ["module: dynamo"]

# ruff: noqa

# flake8: noqa

# Test copied from

# https://raw.githubusercontent.com/python/cpython/refs/tags/v3.13.5/Lib/test/test_generator_stop.py

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()