Improve grouper() recipe to demonstrate all forms of zip() (GH-30837) (GH-30840)

2 files changed, 41 insertions, 9 deletions

diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
index 61d8b86971..34667561c3 100644
--- a/Doc/library/itertools.rst
+++ b/Doc/library/itertools.rst
@@ -813,11 +813,20 @@ which incur interpreter overhead.
            return starmap(func, repeat(args))
        return starmap(func, repeat(args, times))
 
-   def grouper(iterable, n, fillvalue=None):
+   def grouper(iterable, n, *, incomplete='fill', fillvalue=None):
        "Collect data into non-overlapping fixed-length chunks or blocks"
-       # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx
+       # grouper('ABCDEFG', 3, fillvalue='x') --> ABC DEF Gxx
+       # grouper('ABCDEFG', 3, incomplete='strict') --> ABC DEF ValueError
+       # grouper('ABCDEFG', 3, incomplete='ignore') --> ABC DEF
        args = [iter(iterable)] * n
-       return zip_longest(*args, fillvalue=fillvalue)
+       if incomplete == 'fill':
+           return zip_longest(*args, fillvalue=fillvalue)
+       if incomplete == 'strict':
+           return zip(*args, strict=True)
+       if incomplete == 'ignore':
+           return zip(*args)
+       else:
+           raise ValueError('Expected fill, strict, or ignore')
 
    def triplewise(iterable):
        "Return overlapping triplets from an iterable"
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index a12f6f0b97..4c9c597cc4 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -2394,6 +2394,21 @@ Samuele
 ...     else:
 ...         return starmap(func, repeat(args, times))
 
+>>> def grouper(iterable, n, *, incomplete='fill', fillvalue=None):
+...     "Collect data into non-overlapping fixed-length chunks or blocks"
+...     # grouper('ABCDEFG', 3, fillvalue='x') --> ABC DEF Gxx
+...     # grouper('ABCDEFG', 3, incomplete='strict') --> ABC DEF ValueError
+...     # grouper('ABCDEFG', 3, incomplete='ignore') --> ABC DEF
+...     args = [iter(iterable)] * n
+...     if incomplete == 'fill':
+...         return zip_longest(*args, fillvalue=fillvalue)
+...     if incomplete == 'strict':
+...         return zip(*args, strict=True)
+...     if incomplete == 'ignore':
+...         return zip(*args)
+...     else:
+...         raise ValueError('Expected fill, strict, or ignore')
+
 >>> def triplewise(iterable):
 ...     "Return overlapping triplets from an iterable"
 ...     # pairwise('ABCDEFG') -> ABC BCD CDE DEF EFG
@@ -2411,11 +2426,6 @@ Samuele
 ...         window.append(x)
 ...         yield tuple(window)
 
->>> def grouper(n, iterable, fillvalue=None):
-...     "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx"
-...     args = [iter(iterable)] * n
-...     return zip_longest(*args, fillvalue=fillvalue)
-
 >>> def roundrobin(*iterables):
 ...     "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
 ...     # Recipe credited to George Sakkis
@@ -2584,9 +2594,22 @@ True
 >>> dotproduct([1,2,3], [4,5,6])
 32
 
->>> list(grouper(3, 'abcdefg', 'x'))
+>>> list(grouper('abcdefg', 3, fillvalue='x'))
 [('a', 'b', 'c'), ('d', 'e', 'f'), ('g', 'x', 'x')]
 
+>>> it = grouper('abcdefg', 3, incomplete='strict')
+>>> next(it)
+('a', 'b', 'c')
+>>> next(it)
+('d', 'e', 'f')
+>>> next(it)
+Traceback (most recent call last):
+  ...
+ValueError: zip() argument 2 is shorter than argument 1
+
+>>> list(grouper('abcdefg', n=3, incomplete='ignore'))
+[('a', 'b', 'c'), ('d', 'e', 'f')]
+
 >>> list(triplewise('ABCDEFG'))
 [('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E'), ('D', 'E', 'F'), ('E', 'F', 'G')]