Skip to content

GH-98363: Add itertools.batched() #98364

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rhettinger merged 24 commits into from
Oct 17, 2022
Merged

GH-98363: Add itertools.batched() #98364

Show file tree
Hide file tree
Changes from 22 commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
057390d
First draft
rhettinger Oct 17, 2022
047c441
.
rhettinger Oct 17, 2022
9061042
Fix NULL test
rhettinger Oct 17, 2022
72a03c6
Improve the docstring and add the regular docs
rhettinger Oct 17, 2022
21a6f67
.
rhettinger Oct 17, 2022
61b83df
Add tests
rhettinger Oct 17, 2022
93bf013
Add entry to toplevel table in the docs
rhettinger Oct 17, 2022
eb6949c
Remove bogus decref
rhettinger Oct 17, 2022
7fd4016
Flatten before joining
rhettinger Oct 17, 2022
8fb8bdf
Add blurb
rhettinger Oct 17, 2022
2b994a3
Add motivating use case and make equivalent code more accurate
rhettinger Oct 17, 2022
fa9c306
Remove TODO comment
rhettinger Oct 17, 2022
c49fb17
Update Lib/test/test_itertools.py
rhettinger Oct 17, 2022
3fbc9c5
Merge branch 'main' into itertools_batched
rhettinger Oct 17, 2022
4c25341
Neaten-up
rhettinger Oct 17, 2022
41acb45
Make the example more self-explanatory
rhettinger Oct 17, 2022
77de77c
Test pure python equivalent
rhettinger Oct 17, 2022
7390670
Test pure python equivalent
rhettinger Oct 17, 2022
c836bc7
Add GC test
rhettinger Oct 17, 2022
c0ca79d
Add "TestVariousIteratorArgs" tests
rhettinger Oct 17, 2022
8238c28
Fix test case I()
rhettinger Oct 17, 2022
b9cbb46
Remove development note
rhettinger Oct 17, 2022
4799e6a
Remove docstring from pure Python equivalent.
rhettinger Oct 17, 2022
bac2f6c
Fix missing argument in the docs table entry.
rhettinger Oct 17, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 40 additions & 37 deletions Doc/library/itertools.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ Iterator Arguments Results
Iterator Arguments Results Example
============================ ============================ ================================================= =============================================================
:func:`accumulate` p [,func] p0, p0+p1, p0+p1+p2, ... ``accumulate([1,2,3,4,5]) --> 1 3 6 10 15``
:func:`batched` p, n [p0, p1, ..., p_n-1], ... ``batched('ABCDEFG') --> ABC DEF G``
Comment thread
rhettinger marked this conversation as resolved.
Outdated
:func:`chain` p, q, ... p0, p1, ... plast, q0, q1, ... ``chain('ABC', 'DEF') --> A B C D E F``
:func:`chain.from_iterable` iterable p0, p1, ... plast, q0, q1, ... ``chain.from_iterable(['ABC', 'DEF']) --> A B C D E F``
:func:`compress` data, selectors (d[0] if s[0]), (d[1] if s[1]), ... ``compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F``
Expand Down Expand Up @@ -170,6 +171,45 @@ loops that truncate the stream.
.. versionchanged:: 3.8
Added the optional *initial* parameter.


.. function:: batched(iterable, n)

Batch data from the *iterable* into lists of length *n*. The last
batch may be shorter than *n*.

Loops over the input iterable and accumulates data into lists up to
size *n*. The input is consumed lazily, just enough to fill a list.
The result is yielded as soon as the batch is full or when the input
iterable is exhausted:

.. doctest::

>>> flattened_data = ['roses', 'red', 'violets', 'blue', 'sugar', 'sweet']
>>> unflattened = list(batched(flattened_data, 2))
>>> unflattened
[['roses', 'red'], ['violets', 'blue'], ['sugar', 'sweet']]

>>> for batch in batched('ABCDEFG', 3):
... print(batch)
...
['A', 'B', 'C']
['D', 'E', 'F']
['G']

Roughly equivalent to::

def batched(iterable, n):
"Batch data into lists of length n. The last batch may be shorter."
# batched('ABCDEFG', 3) --> ABC DEF G
Copy link
Copy Markdown
Member

@gvanrossum gvanrossum Oct 23, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe tweak the example so it cannot be misunderstood as "return 3 batches"?

Copy link
Copy Markdown

@anderssonjohan anderssonjohan Oct 25, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
# batched('ABCDEFG', 3) --> ABC DEF G
# list(batched('ABCDEFG', 3)) --> [['A', 'B', 'C'], ['D', 'E', 'F'], ['G']]

if n < 1:
raise ValueError('n must be at least one')
it = iter(iterable)
while (batch := list(islice(it, n))):
yield batch

.. versionadded:: 3.12


.. function:: chain(*iterables)

Make an iterator that returns elements from the first iterable until it is
Expand Down Expand Up @@ -858,13 +898,6 @@ which incur interpreter overhead.
else:
raise ValueError('Expected fill, strict, or ignore')

def batched(iterable, n):
"Batch data into lists of length n. The last batch may be shorter."
# batched('ABCDEFG', 3) --> ABC DEF G
it = iter(iterable)
while (batch := list(islice(it, n))):
yield batch

def triplewise(iterable):
"Return overlapping triplets from an iterable"
# triplewise('ABCDEFG') --> ABC BCD CDE DEF EFG
Expand Down Expand Up @@ -1236,36 +1269,6 @@ which incur interpreter overhead.
>>> list(grouper('abcdefg', n=3, incomplete='ignore'))
[('a', 'b', 'c'), ('d', 'e', 'f')]

>>> list(batched('ABCDEFG', 3))
[['A', 'B', 'C'], ['D', 'E', 'F'], ['G']]
>>> list(batched('ABCDEF', 3))
[['A', 'B', 'C'], ['D', 'E', 'F']]
>>> list(batched('ABCDE', 3))
[['A', 'B', 'C'], ['D', 'E']]
>>> list(batched('ABCD', 3))
[['A', 'B', 'C'], ['D']]
>>> list(batched('ABC', 3))
[['A', 'B', 'C']]
>>> list(batched('AB', 3))
[['A', 'B']]
>>> list(batched('A', 3))
[['A']]
>>> list(batched('', 3))
[]
>>> list(batched('ABCDEFG', 2))
[['A', 'B'], ['C', 'D'], ['E', 'F'], ['G']]
>>> list(batched('ABCDEFG', 1))
[['A'], ['B'], ['C'], ['D'], ['E'], ['F'], ['G']]
>>> list(batched('ABCDEFG', 0))
[]
>>> list(batched('ABCDEFG', -1))
Traceback (most recent call last):
...
ValueError: Stop argument for islice() must be None or an integer: 0 <= x <= sys.maxsize.
>>> s = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
>>> all(list(flatten(batched(s[:n], 5))) == list(s[:n]) for n in range(len(s)))
True

>>> list(triplewise('ABCDEFG'))
[('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E'), ('D', 'E', 'F'), ('E', 'F', 'G')]

Expand Down
79 changes: 79 additions & 0 deletions Lib/test/test_itertools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,44 @@ def test_accumulate(self):
with self.assertRaises(TypeError):
list(accumulate([10, 20], 100))

def test_batched(self):
self.assertEqual(list(batched('ABCDEFG', 3)),
[['A', 'B', 'C'], ['D', 'E', 'F'], ['G']])
self.assertEqual(list(batched('ABCDEFG', 2)),
[['A', 'B'], ['C', 'D'], ['E', 'F'], ['G']])
self.assertEqual(list(batched('ABCDEFG', 1)),
[['A'], ['B'], ['C'], ['D'], ['E'], ['F'], ['G']])

with self.assertRaises(TypeError): # Too few arguments
list(batched('ABCDEFG'))
with self.assertRaises(TypeError):
list(batched('ABCDEFG', 3, None)) # Too many arguments
with self.assertRaises(TypeError):
list(batched(None, 3)) # Non-iterable input
with self.assertRaises(TypeError):
list(batched('ABCDEFG', 'hello')) # n is a string
with self.assertRaises(ValueError):
list(batched('ABCDEFG', 0)) # n is zero
with self.assertRaises(ValueError):
list(batched('ABCDEFG', -1)) # n is negative

data = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
for n in range(1, 6):
for i in range(len(data)):
s = data[:i]
batches = list(batched(s, n))
with self.subTest(s=s, n=n, batches=batches):
# Order is preserved and no data is lost
self.assertEqual(''.join(chain(*batches)), s)
# Each batch is an exact list
self.assertTrue(all(type(batch) is list for batch in batches))
# All but the last batch is of size n
if batches:
last_batch = batches.pop()
self.assertTrue(all(len(batch) == n for batch in batches))
self.assertTrue(len(last_batch) <= n)
batches.append(last_batch)

def test_chain(self):

def chain2(*iterables):
Expand Down Expand Up @@ -1737,6 +1775,31 @@ def test_takewhile(self):

class TestPurePythonRoughEquivalents(unittest.TestCase):

def test_batched_recipe(self):
def batched_recipe(iterable, n):
"Batch data into lists of length n. The last batch may be shorter."
# batched('ABCDEFG', 3) --> ABC DEF G
if n < 1:
raise ValueError('n must be at least one')
it = iter(iterable)
while (batch := list(islice(it, n))):
yield batch

for iterable, n in product(
['', 'a', 'ab', 'abc', 'abcd', 'abcde', 'abcdef', 'abcdefg', None],
[-1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, None]):
with self.subTest(iterable=iterable, n=n):
try:
e1, r1 = None, list(batched(iterable, n))
except Exception as e:
e1, r1 = type(e), None
try:
e2, r2 = None, list(batched_recipe(iterable, n))
except Exception as e:
e2, r2 = type(e), None
self.assertEqual(r1, r2)
self.assertEqual(e1, e2)

@staticmethod
def islice(iterable, *args):
s = slice(*args)
Expand Down Expand Up @@ -1788,6 +1851,10 @@ def test_accumulate(self):
a = []
self.makecycle(accumulate([1,2,a,3]), a)

def test_batched(self):
a = []
self.makecycle(batched([1,2,a,3], 2), a)

def test_chain(self):
a = []
self.makecycle(chain(a), a)
Expand Down Expand Up @@ -1972,6 +2039,18 @@ def test_accumulate(self):
self.assertRaises(TypeError, accumulate, N(s))
self.assertRaises(ZeroDivisionError, list, accumulate(E(s)))

def test_batched(self):
s = 'abcde'
r = [['a', 'b'], ['c', 'd'], ['e']]
n = 2
for g in (G, I, Ig, L, R):
with self.subTest(g=g):
self.assertEqual(list(batched(g(s), n)), r)
self.assertEqual(list(batched(S(s), 2)), [])
self.assertRaises(TypeError, batched, X(s), 2)
self.assertRaises(TypeError, batched, N(s), 2)
self.assertRaises(ZeroDivisionError, list, batched(E(s), 2))

def test_chain(self):
for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)):
for g in (G, I, Ig, S, L, R):
Expand Down
2 changes: 2 additions & 0 deletions Misc/NEWS.d/next/Library/2022-10-17-12-49-02.gh-issue-98363.aFmSP-.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Added itertools.batched() to batch data into lists of a given length with
the last list possibly being shorter than the others.
81 changes: 80 additions & 1 deletion Modules/clinic/itertoolsmodule.c.h
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading