Add CPython test test_itertools (#156981)

Test the itertools module

Pull Request resolved: https://github.com/pytorch/pytorch/pull/156981
Approved by: https://github.com/zou3519
ghstack dependencies: #157799, #157800, #157801, #157802

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

@@ -0,0 +1,264 @@
+diff --git a/test/dynamo/cpython/3_13/test_itertools.py b/test/dynamo/cpython/3_13/test_itertools.py
+index 7d5ba727389..637fbe545cd 100644
+--- a/test/dynamo/cpython/3_13/test_itertools.py
++++ b/test/dynamo/cpython/3_13/test_itertools.py
+@@ -1,3 +1,25 @@
++# ======= 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_itertools.py
++
++import torch
++import torch._dynamo.test_case
++from torch._dynamo.test_case import CPythonTestCase
++from torch.testing._internal.common_utils import (
++    run_tests,
++    skipIfTorchDynamo,
++    slowTest,
++)
++
++__TestCase = CPythonTestCase
++
++# ======= END DYNAMO PATCH =======
++
+ import doctest
+ import unittest
+ import itertools
+@@ -90,10 +112,10 @@ def fact(n):
+     return prod(range(1, n+1))
+ 
+ # root level methods for pickling ability
+-def testR(r):
++def _testR(r):
+     return r[0]
+ 
+-def testR2(r):
++def _testR2(r):
+     return r[2]
+ 
+ def underten(x):
+@@ -102,7 +124,7 @@ def underten(x):
+ picklecopiers = [lambda s, proto=proto: pickle.loads(pickle.dumps(s, proto))
+                  for proto in range(pickle.HIGHEST_PROTOCOL + 1)]
+ 
+-class TestBasicOps(unittest.TestCase):
++class TestBasicOps(__TestCase):
+ 
+     def pickletest(self, protocol, it, stop=4, take=1, compare=None):
+         """Test that an iterator is the same after pickling, also when part-consumed"""
+@@ -888,7 +910,7 @@ class TestBasicOps(unittest.TestCase):
+         # Check normal pickled
+         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+             dup = []
+-            for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
++            for k, g in pickle.loads(pickle.dumps(groupby(s, _testR), proto)):
+                 for elem in g:
+                     self.assertEqual(k, elem[0])
+                     dup.append(elem)
+@@ -896,8 +918,8 @@ class TestBasicOps(unittest.TestCase):
+ 
+         # Check nested case
+         dup = []
+-        for k, g in groupby(s, testR):
+-            for ik, ig in groupby(g, testR2):
++        for k, g in groupby(s, _testR):
++            for ik, ig in groupby(g, _testR2):
+                 for elem in ig:
+                     self.assertEqual(k, elem[0])
+                     self.assertEqual(ik, elem[2])
+@@ -907,8 +929,8 @@ class TestBasicOps(unittest.TestCase):
+         # Check nested and pickled
+         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+             dup = []
+-            for k, g in pickle.loads(pickle.dumps(groupby(s, testR), proto)):
+-                for ik, ig in pickle.loads(pickle.dumps(groupby(g, testR2), proto)):
++            for k, g in pickle.loads(pickle.dumps(groupby(s, _testR), proto)):
++                for ik, ig in pickle.loads(pickle.dumps(groupby(g, _testR2), proto)):
+                     for elem in ig:
+                         self.assertEqual(k, elem[0])
+                         self.assertEqual(ik, elem[2])
+@@ -917,7 +939,7 @@ class TestBasicOps(unittest.TestCase):
+ 
+ 
+         # Check case where inner iterator is not used
+-        keys = [k for k, g in groupby(s, testR)]
++        keys = [k for k, g in groupby(s, _testR)]
+         expectedkeys = set([r[0] for r in s])
+         self.assertEqual(set(keys), expectedkeys)
+         self.assertEqual(len(keys), len(expectedkeys))
+@@ -925,7 +947,7 @@ class TestBasicOps(unittest.TestCase):
+         # Check case where inner iterator is used after advancing the groupby
+         # iterator
+         s = list(zip('AABBBAAAA', range(9)))
+-        it = groupby(s, testR)
++        it = groupby(s, _testR)
+         _, g1 = next(it)
+         _, g2 = next(it)
+         _, g3 = next(it)
+@@ -936,7 +958,7 @@ class TestBasicOps(unittest.TestCase):
+         self.assertEqual(list(g3), [])
+ 
+         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+-            it = groupby(s, testR)
++            it = groupby(s, _testR)
+             _, g = next(it)
+             next(it)
+             next(it)
+@@ -1038,6 +1060,7 @@ class TestBasicOps(unittest.TestCase):
+         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+             self.pickletest(proto, filterfalse(isEven, range(6)))
+ 
++    @skipIfTorchDynamo("infinite loop in torch dynamo")
+     def test_zip(self):
+         # XXX This is rather silly now that builtin zip() calls zip()...
+         ans = [(x,y) for x, y in zip('abc',count())]
+@@ -1082,6 +1105,7 @@ class TestBasicOps(unittest.TestCase):
+         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+             self.pickletest(proto, zip('abc', count()))
+ 
++    @skipIfTorchDynamo("infinite loop in torch dynamo")
+     def test_ziplongest(self):
+         for args in [
+                 ['abc', range(6)],
+@@ -1767,6 +1791,7 @@ class TestBasicOps(unittest.TestCase):
+         script_helper.assert_python_ok("-c", script)
+ 
+     # Issue 13454: Crash when deleting backward iterator from tee()
++    @skipIfTorchDynamo("infinite loop in torch dynamo")
+     def test_tee_del_backward(self):
+         forward, backward = tee(repeat(None, 20000000))
+         try:
+@@ -1920,7 +1945,7 @@ class TestBasicOps(unittest.TestCase):
+                     tp.foobar = 1
+ 
+ 
+-class TestExamples(unittest.TestCase):
++class TestExamples(__TestCase):
+ 
+     def test_accumulate(self):
+         self.assertEqual(list(accumulate([1,2,3,4,5])), [1, 3, 6, 10, 15])
+@@ -2032,7 +2057,7 @@ class TestExamples(unittest.TestCase):
+         self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
+ 
+ 
+-class TestPurePythonRoughEquivalents(unittest.TestCase):
++class TestPurePythonRoughEquivalents(__TestCase):
+ 
+     def test_batched_recipe(self):
+         def batched_recipe(iterable, n):
+@@ -2081,6 +2106,7 @@ class TestPurePythonRoughEquivalents(unittest.TestCase):
+             for i, element in zip(range(i + 1, stop), iterable):
+                 pass
+ 
++    @skipIfTorchDynamo("infinite loop in torch dynamo")
+     def test_islice_recipe(self):
+         self.assertEqual(list(self.islice('ABCDEFG', 2)), list('AB'))
+         self.assertEqual(list(self.islice('ABCDEFG', 2, 4)), list('CD'))
+@@ -2265,7 +2291,7 @@ class TestPurePythonRoughEquivalents(unittest.TestCase):
+             raise
+ 
+ 
+-class TestGC(unittest.TestCase):
++class TestGC(__TestCase):
+ 
+     def makecycle(self, iterator, container):
+         container.append(iterator)
+@@ -2465,7 +2491,7 @@ def L(seqn):
+     return chain(map(lambda x:x, R(Ig(G(seqn)))))
+ 
+ 
+-class TestVariousIteratorArgs(unittest.TestCase):
++class TestVariousIteratorArgs(__TestCase):
+ 
+     def test_accumulate(self):
+         s = [1,2,3,4,5]
+@@ -2644,7 +2670,7 @@ class TestVariousIteratorArgs(unittest.TestCase):
+             self.assertRaises(TypeError, tee, N(s))
+             self.assertRaises(ZeroDivisionError, list, tee(E(s))[0])
+ 
+-class LengthTransparency(unittest.TestCase):
++class LengthTransparency(__TestCase):
+ 
+     def test_repeat(self):
+         self.assertEqual(operator.length_hint(repeat(None, 50)), 50)
+@@ -2657,7 +2683,7 @@ class LengthTransparency(unittest.TestCase):
+         self.assertEqual(operator.length_hint(repeat(None, times=-1)), 0)
+         self.assertEqual(operator.length_hint(repeat(None, times=-2)), 0)
+ 
+-class RegressionTests(unittest.TestCase):
++class RegressionTests(__TestCase):
+ 
+     def test_sf_793826(self):
+         # Fix Armin Rigo's successful efforts to wreak havoc
+@@ -2718,6 +2744,7 @@ class RegressionTests(unittest.TestCase):
+ 
+     @support.skip_if_pgo_task
+     @support.requires_resource('cpu')
++    @slowTest
+     def test_long_chain_of_empty_iterables(self):
+         # Make sure itertools.chain doesn't run into recursion limits when
+         # dealing with long chains of empty iterables. Even with a high
+@@ -2750,7 +2777,7 @@ class RegressionTests(unittest.TestCase):
+             next(g, None)  # shouldn't crash
+ 
+ 
+-class SubclassWithKwargsTest(unittest.TestCase):
++class SubclassWithKwargsTest(__TestCase):
+     def test_keywords_in_subclass(self):
+         # count is not subclassable...
+         testcases = [
+@@ -2805,49 +2832,5 @@ class SubclassWithKwargsTest(unittest.TestCase):
+                 self.assertEqual(u.newarg, 3)
+ 
+ 
+-@support.cpython_only
+-class SizeofTest(unittest.TestCase):
+-    def setUp(self):
+-        self.ssize_t = struct.calcsize('n')
+-
+-    check_sizeof = support.check_sizeof
+-
+-    def test_product_sizeof(self):
+-        basesize = support.calcobjsize('3Pi')
+-        check = self.check_sizeof
+-        check(product('ab', '12'), basesize + 2 * self.ssize_t)
+-        check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t)
+-
+-    def test_combinations_sizeof(self):
+-        basesize = support.calcobjsize('3Pni')
+-        check = self.check_sizeof
+-        check(combinations('abcd', 3), basesize + 3 * self.ssize_t)
+-        check(combinations(range(10), 4), basesize + 4 * self.ssize_t)
+-
+-    def test_combinations_with_replacement_sizeof(self):
+-        cwr = combinations_with_replacement
+-        basesize = support.calcobjsize('3Pni')
+-        check = self.check_sizeof
+-        check(cwr('abcd', 3), basesize + 3 * self.ssize_t)
+-        check(cwr(range(10), 4), basesize + 4 * self.ssize_t)
+-
+-    def test_permutations_sizeof(self):
+-        basesize = support.calcobjsize('4Pni')
+-        check = self.check_sizeof
+-        check(permutations('abcd'),
+-              basesize + 4 * self.ssize_t + 4 * self.ssize_t)
+-        check(permutations('abcd', 3),
+-              basesize + 4 * self.ssize_t + 3 * self.ssize_t)
+-        check(permutations('abcde', 3),
+-              basesize + 5 * self.ssize_t + 3 * self.ssize_t)
+-        check(permutations(range(10), 4),
+-              basesize + 10 * self.ssize_t + 4 * self.ssize_t)
+-
+-
+-def load_tests(loader, tests, pattern):
+-    tests.addTest(doctest.DocTestSuite(itertools))
+-    return tests
+-
+-
+ if __name__ == "__main__":
+-    unittest.main()
++    run_tests()