8000 gh-123884 Tee of tee was not producing n independent iterators (gh-12… · python/cpython@909c6f7 · GitHub
[go: up one dir, main page]
Skip to content

Commit 909c6f7

Browse files
rhettingerrhettinger
authored
gh-123884 Tee of tee was not producing n independent iterators (gh-124490)
1 parent fb6bd31 commit 909c6f7

8 files changed

+91
-89
lines changed

‎Doc/library/itertools.rst

Lines changed: 30 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -691,25 +691,36 @@ loops that truncate the stream.
691691

692692
def tee(iterable, n=2):
693693
if n < 0:
694-
raise ValueError('n must be >= 0')
695-
iterator = iter(iterable)
696-
shared_link = [None, None]
697-
return tuple(_tee(iterator, shared_link) for _ in range(n))
698-
699-
def _tee(iterator, link):
700-
try:
701-
while True:
702-
if link[1] is None:
703-
link[0] = next(iterator)
704-
link[1] = [None, None]
705-
value, link = link
706-
yield value
707-
except StopIteration:
708-
return
709-
710-
Once a :func:`tee` has been created, the original *iterable* should not be
711-
used anywhere else; otherwise, the *iterable* could get advanced without
712-
the tee objects being informed.
694+
raise ValueError
695+
if n == 0:
696+
return ()
697+
iterator = _tee(iterable)
698+
result = [iterator]
699+
for _ in range(n - 1):
700+
result.append(_tee(iterator))
701+
return tuple(result)
702+
703+
class _tee:
704+
705+
def __init__(self, iterable):
706+
it = iter(iterable)
707+
if isinstance(it, _tee):
708+
self.iterator = it.iterator
709+
self.link = it.link
710+
else:
711+
self.iterator = it
712+
self.link = [None, None]
713+
714+
def __iter__(self):
715+
return self
716+
717+
def __next__(self):
718+
link = self.link
719+
if link[1] is None:
720+
link[0] = next(self.iterator)
721+
link[1] = [None, None]
722+
value, self.link = link
723+
return value
713724

714725
When the input *iterable* is already a tee iterator object, all
715726
members of the return tuple are constructed as if they had been

‎Include/internal/pycore_global_objects_fini_generated.h

Lines changed: 0 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Include/internal/pycore_global_strings.h

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,6 @@ struct _Py_global_strings {
9393
STRUCT_FOR_ID(__classdictcell__)
9494
STRUCT_FOR_ID(__complex__)
9595
STRUCT_FOR_ID(__contains__)
96-
STRUCT_FOR_ID(__copy__)
9796
STRUCT_FOR_ID(__ctypes_from_outparam__)
9897
STRUCT_FOR_ID(__del__)
9998
STRUCT_FOR_ID(__delattr__)

‎Include/internal/pycore_runtime_init_generated.h

Lines changed: 0 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Include/internal/pycore_unicodeobject_generated.h

Lines changed: 0 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Lib/test/test_itertools.py

Lines changed: 48 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1249,10 +1249,11 @@ def test_tee(self):
12491249
self.assertEqual(len(result), n)
12501250
self.assertEqual([list(x) for x in result], [list('abc')]*n)
12511251

1252-
# tee pass-through to copyable iterator
1252+
# tee objects are independent (see bug gh-123884)
12531253
a, b = tee('abc')
12541254
c, d = tee(a)
1255-
self.assertTrue(a is c)
1255+
e, f = tee(c)
1256+
self.assertTrue(len({a, b, c, d, e, f}) == 6)
12561257

12571258
# test tee_new
12581259
t1, t2 = tee('abc')
@@ -1759,21 +1760,36 @@ def test_tee_recipe(self):
17591760

17601761
def tee(iterable, n=2):
17611762
if n < 0:
1762-
raise ValueError('n must be >= 0')
1763-
iterator = iter(iterable)
1764-
shared_link = [None, None]
1765-
return tuple(_tee(iterator, shared_link) for _ in range(n))
1763+
raise ValueError
1764+
if n == 0:
1765+
return ()
1766+
iterator = _tee(iterable)
1767+
result = [iterator]
1768+
for _ in range(n - 1):
1769+
result.append(_tee(iterator))
1770+
return tuple(result)
1771+
1772+
class _tee:
1773+
1774+
def __init__(self, iterable):
1775+
it = iter(iterable)
1776+
if isinstance(it, _tee):
1777+
self.iterator = it.iterator
1778+
self.link = it.link
1779+
else:
1780+
self.iterator = it
1781+
self.link = [None, None]
17661782

1767-
def _tee(iterator, link):
1768-
try:
1769-
while True:
1770-
if link[1] is None:
1771-
link[0] = next(iterator)
1772-
link[1] = [None, None]
1773-
value, link = link
1774-
yield value
1775-
except StopIteration:
1776-
return
1783+
def __iter__(self):
1784+
return self
1785+
1786+
def __next__(self):
1787+
link = self.link
1788+
if link[1] is None:
1789+
link[0] = next(self.iterator)
1790+
link[1] = [None, None]
1791+
value, self.link = link
1792+
return value
17771793

17781794
# End tee() recipe #############################################
17791795

@@ -1819,12 +1835,10 @@ def _tee(iterator, link):
18191835
self.assertRaises(TypeError, tee, [1,2], 'x')
18201836
self.assertRaises(TypeError, tee, [1,2], 3, 'x')
18211837

1822-
# Tests not applicable to the tee() recipe
1823-
if False:
1824-
# tee object should be instantiable
1825-
a, b = tee('abc')
1826-
c = type(a)('def')
1827-
self.assertEqual(list(c), list('def'))
1838+
# tee object should be instantiable
1839+
a, b = tee('abc')
1840+
c = type(a)('def')
1841+
self.assertEqual(list(c), list('def'))
18281842

18291843
# test long-lagged and multi-way split
18301844
a, b, c = tee(range(2000), 3)
@@ -1845,21 +1859,19 @@ def _tee(iterator, link):
18451859
self.assertEqual(len(result), n)
18461860
self.assertEqual([list(x) for x in result], [list('abc')]*n)
18471861

1862+
# tee objects are independent (see bug gh-123884)
1863+
a, b = tee('abc')
1864+
c, d = tee(a)
1865+
e, f = tee(c)
1866+
self.assertTrue(len({a, b, c, d, e, f}) == 6)
18481867

1849-
# Tests not applicable to the tee() recipe
1850-
if False:
1851-
# tee pass-through to copyable iterator
1852-
a, b = tee('abc')
1853-
c, d = tee(a)
1854-
self.assertTrue(a is c)
1855-
1856-
# test tee_new
1857-
t1, t2 = tee('abc')
1858-
tnew = type(t1)
1859-
self.assertRaises(TypeError, tnew)
1860-
self.assertRaises(TypeError, tnew, 10)
1861-
t3 = tnew(t1)
1862-
self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
1868+
# test tee_new
1869+
t1, t2 = tee('abc')
1870+
tnew = type(t1)
1871+
self.assertRaises(TypeError, tnew)
1872+
self.assertRaises(TypeError, tnew, 10)
1873+
t3 = tnew(t1)
1874+
self.assertTrue(list(t1) == list(t2) == list(t3) == list('abc'))
18631875

18641876
# test that tee objects are weak referencable
18651877
a, b = tee(range(10))

‎Misc/NEWS.d/next/Library/2024-09-24-22-38-51.gh-issue-123884.iEPTK4.rst

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
Fixed bug in itertools.tee() handling of other tee inputs (a tee in a tee).
2+
The output now has the promised *n* independent new iterators. Formerly,
3+
the first iterator was identical (not independent) to the input iterator.
4+
This would sometimes give surprising results.

‎Modules/itertoolsmodule.c

Lines changed: 9 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1036,7 +1036,7 @@ itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
10361036
/*[clinic end generated code: output=1c64519cd859c2f0 input=c99a1472c425d66d]*/
10371037
{
10381038
Py_ssize_t i;
1039-
PyObject *it, *copyable, *copyfunc, *result;
1039+
PyObject *it, *to, *result;
10401040

10411041
if (n < 0) {
10421042
PyErr_SetString(PyExc_ValueError, "n must be >= 0");
@@ -1053,41 +1053,23 @@ itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
10531053
return NULL;
10541054
}
10551055

1056-
if (PyObject_GetOptionalAttr(it, &_Py_ID(__copy__), &copyfunc) < 0) {
1057-
Py_DECREF(it);
1056+
itertools_state *state = get_module_state(module);
1057+
to = tee_fromiterable(state, it);
1058+
Py_DECREF(it);
1059+
if (to == NULL) {
10581060
Py_DECREF(result);
10591061
return NULL;
10601062
10000 }
1061-
if (copyfunc != NULL) {
1062-
copyable = it;
1063-
}
1064-
else {
1065-
itertools_state *state = get_module_state(module);
1066-
copyable = tee_fromiterable(state, it);
1067-
Py_DECREF(it);
1068-
if (copyable == NULL) {
1069-
Py_DECREF(result);
1070-
return NULL;
1071-
}
1072-
copyfunc = PyObject_GetAttr(copyable, &_Py_ID(__copy__));
1073-
if (copyfunc == NULL) {
1074-
Py_DECREF(copyable);
1075-
Py_DECREF(result);
1076-
return NULL;
1077-
}
1078-
}
10791063

1080-
PyTuple_SET_ITEM(result, 0, copyable);
1064+
PyTuple_SET_ITEM(result, 0, to);
10811065
for (i = 1; i < n; i++) {
1082-
copyable = _PyObject_CallNoArgs(copyfunc);
1083-
if (copyable == NULL) {
1084-
Py_DECREF(copyfunc);
1066+
to = tee_copy((teeobject *)to, NULL);
1067+
if (to == NULL) {
10851068
Py_DECREF(result);
10861069
return NULL;
10871070
}
1088-
PyTuple_SET_ITEM(result, i, copyable);
1071+
PyTuple_SET_ITEM(result, i, to);
10891072
}
1090-
Py_DECREF(copyfunc);
10911073
return result;
10921074
}
10931075

0 commit comments

Comments
 (0)
0