Misc updates to the itertools recipes and tests (GH-98018)

(cherry picked from commit e500cc04517bd65668f2e203c1e37b0cc5b1cf1b)

Co-authored-by: Raymond Hettinger <rhettinger@users.noreply.github.com>

diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
index 2b4d7d535e0933e915206696f938591b5ffb26b3..9de09ed46cfb931040deb4626a2d257691d417fd 100644 (file)
--- a/Doc/library/itertools.rst
+++ b/Doc/library/itertools.rst
@@ -775,10 +775,7 @@ which incur interpreter overhead.
        return sum(map(pred, iterable))
 
    def pad_none(iterable):
-       """Returns the sequence elements and then returns None indefinitely.
-
-       Useful for emulating the behavior of the built-in map() function.
-       """
+       "Returns the sequence elements and then returns None indefinitely."
        return chain(iterable, repeat(None))
 
    def ncycles(iterable, n):
@@ -860,6 +857,13 @@ 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
@@ -1232,6 +1236,36 @@ 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')]