Issue #26494: Fixed crash on iterating exhausting iterators.

Affected classes are generic sequence iterators, iterators of str, bytes,
bytearray, list, tuple, set, frozenset, dict, OrderedDict, corresponding
views and os.scandir() iterator.

9 files changed, 55 insertions, 0 deletions

diff --git a/Lib/test/seq_tests.py b/Lib/test/seq_tests.py
index 2416249..72f4845 100644
--- a/Lib/test/seq_tests.py
+++ b/Lib/test/seq_tests.py
@@ -5,6 +5,7 @@ Tests common to tuple, list and UserList.UserList
 import unittest
 import sys
 import pickle
+from test import support
 
 # Various iterables
 # This is used for checking the constructor (here and in test_deque.py)
@@ -408,3 +409,7 @@ class CommonTest(unittest.TestCase):
             lst2 = pickle.loads(pickle.dumps(lst, proto))
             self.assertEqual(lst2, lst)
             self.assertNotEqual(id(lst2), id(lst))
+
+    def test_free_after_iterating(self):
+        support.check_free_after_iterating(self, iter, self.type2test)
+        support.check_free_after_iterating(self, reversed, self.type2test)
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
index 7914943..04e8629 100644
--- a/Lib/test/support/__init__.py
+++ b/Lib/test/support/__init__.py
@@ -2432,3 +2432,22 @@ def run_in_subinterp(code):
                                      "memory allocations")
     import _testcapi
     return _testcapi.run_in_subinterp(code)
+
+
+def check_free_after_iterating(test, iter, cls, args=()):
+    class A(cls):
+        def __del__(self):
+            nonlocal done
+            done = True
+            try:
+                next(it)
+            except StopIteration:
+                pass
+
+    done = False
+    it = iter(A(*args))
+    # Issue 26494: Shouldn't crash
+    test.assertRaises(StopIteration, next, it)
+    # The sequence should be deallocated just after the end of iterating
+    gc_collect()
+    test.assertTrue(done)
diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py
index 5d0351b..90bcc42 100644
--- a/Lib/test/test_bytes.py
+++ b/Lib/test/test_bytes.py
@@ -761,6 +761,10 @@ class BaseBytesTest:
         self.assertRaisesRegex(TypeError, r'\bendswith\b', b.endswith,
                                 x, None, None, None)
 
+    def test_free_after_iterating(self):
+        test.support.check_free_after_iterating(self, iter, self.type2test)
+        test.support.check_free_after_iterating(self, reversed, self.type2test)
+
 
 class BytesTest(BaseBytesTest, unittest.TestCase):
     type2test = bytes
diff --git a/Lib/test/test_deque.py b/Lib/test/test_deque.py
index 29b44f6..6dbea49 100644
--- a/Lib/test/test_deque.py
+++ b/Lib/test/test_deque.py
@@ -918,6 +918,10 @@ class TestSequence(seq_tests.CommonTest):
         # For now, bypass tests that require slicing
         pass
 
+    def test_free_after_iterating(self):
+        # For now, bypass tests that require slicing
+        self.skipTest("Exhausted deque iterator doesn't free a deque")
+
 #==============================================================================
 
 libreftest = """
diff --git a/Lib/test/test_dict.py b/Lib/test/test_dict.py
index 1049be5..7dd44b9 100644
--- a/Lib/test/test_dict.py
+++ b/Lib/test/test_dict.py
@@ -954,6 +954,12 @@ class DictTest(unittest.TestCase):
         d = {X(): 0, 1: 1}
         self.assertRaises(RuntimeError, d.update, other)
 
+    def test_free_after_iterating(self):
+        support.check_free_after_iterating(self, iter, dict)
+        support.check_free_after_iterating(self, lambda d: iter(d.keys()), dict)
+        support.check_free_after_iterating(self, lambda d: iter(d.values()), dict)
+        support.check_free_after_iterating(self, lambda d: iter(d.items()), dict)
+
 from test import mapping_tests
 
 class GeneralMappingTests(mapping_tests.BasicTestMappingProtocol):
diff --git a/Lib/test/test_iter.py b/Lib/test/test_iter.py
index 56e21f8..54ddbaa 100644
--- a/Lib/test/test_iter.py
+++ b/Lib/test/test_iter.py
@@ -3,6 +3,7 @@
 import sys
 import unittest
 from test.support import run_unittest, TESTFN, unlink, cpython_only
+from test.support import check_free_after_iterating
 import pickle
 import collections.abc
 
@@ -980,6 +981,9 @@ class TestCase(unittest.TestCase):
         self.assertEqual(next(it), 0)
         self.assertEqual(next(it), 1)
 
+    def test_free_after_iterating(self):
+        check_free_after_iterating(self, iter, SequenceClass, (0,))
+
 
 def test_main():
     run_unittest(TestCase)
diff --git a/Lib/test/test_ordered_dict.py b/Lib/test/test_ordered_dict.py
index e1564ba..6fbc1b4 100644
--- a/Lib/test/test_ordered_dict.py
+++ b/Lib/test/test_ordered_dict.py
@@ -608,6 +608,12 @@ class OrderedDictTests:
         gc.collect()
         self.assertIsNone(r())
 
+    def test_free_after_iterating(self):
+        support.check_free_after_iterating(self, iter, self.OrderedDict)
+        support.check_free_after_iterating(self, lambda d: iter(d.keys()), self.OrderedDict)
+        support.check_free_after_iterating(self, lambda d: iter(d.values()), self.OrderedDict)
+        support.check_free_after_iterating(self, lambda d: iter(d.items()), self.OrderedDict)
+
 
 class PurePythonOrderedDictTests(OrderedDictTests, unittest.TestCase):
 
diff --git a/Lib/test/test_set.py b/Lib/test/test_set.py
index ade39fb..15ae42c 100644
--- a/Lib/test/test_set.py
+++ b/Lib/test/test_set.py
@@ -364,6 +364,9 @@ class TestJointOps:
         gc.collect()
         self.assertTrue(ref() is None, "Cycle was not collected")
 
+    def test_free_after_iterating(self):
+        support.check_free_after_iterating(self, iter, self.thetype)
+
 class TestSet(TestJointOps, unittest.TestCase):
     thetype = set
     basetype = set
diff --git a/Lib/test/test_unicode.py b/Lib/test/test_unicode.py
index c30310e..c281146 100644
--- a/Lib/test/test_unicode.py
+++ b/Lib/test/test_unicode.py
@@ -2729,6 +2729,10 @@ class UnicodeTest(string_tests.CommonTest,
                 # Check that the second call returns the same result
                 self.assertEqual(getargs_s_hash(s), chr(k).encode() * (i + 1))
 
+    def test_free_after_iterating(self):
+        support.check_free_after_iterating(self, iter, str)
+        support.check_free_after_iterating(self, reversed, str)
+
 
 class StringModuleTest(unittest.TestCase):
     def test_formatter_parser(self):