diff options
| -rw-r--r-- | Lib/test/test_itertools.py | 10 | ||||
| -rw-r--r-- | Modules/itertoolsmodule.c | 12 |
2 files changed, 20 insertions, 2 deletions
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 81e1711..9e55b2a 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1035,6 +1035,16 @@ class TestBasicOps(unittest.TestCase): for proto in range(pickle.HIGHEST_PROTOCOL + 1): self.pickletest(proto, product(*args)) + def test_product_issue_25021(self): + # test that indices are properly clamped to the length of the tuples + p = product((1, 2),(3,)) + p.__setstate__((0, 0x1000)) # will access tuple element 1 if not clamped + self.assertEqual(next(p), (2, 3)) + # test that empty tuple in the list will result in an immediate StopIteration + p = product((1, 2), (), (3,)) + p.__setstate__((0, 0, 0x1000)) # will access tuple element 1 if not clamped + self.assertRaises(StopIteration, next, p) + def test_repeat(self): self.assertEqual(list(repeat(object='a', times=3)), ['a', 'a', 'a']) self.assertEqual(lzip(range(3),repeat('a')), diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index c7e1919..ff7867d 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -2256,13 +2256,21 @@ product_setstate(productobject *lz, PyObject *state) { PyObject* indexObject = PyTuple_GET_ITEM(state, i); Py_ssize_t index = PyLong_AsSsize_t(indexObject); + PyObject* pool; + Py_ssize_t poolsize; if (index < 0 && PyErr_Occurred()) return NULL; /* not an integer */ + pool = PyTuple_GET_ITEM(lz->pools, i); + poolsize = PyTuple_GET_SIZE(pool); + if (poolsize == 0) { + lz->stopped = 1; + Py_RETURN_NONE; + } /* clamp the index */ if (index < 0) index = 0; - else if (index > n-1) - index = n-1; + else if (index > poolsize-1) + index = poolsize-1; lz->indices[i] = index; } |