@@ -164,11 +164,14 @@ loops that truncate the stream.
164164 Added the optional *initial * parameter.
165165
166166
167- .. function :: batched(iterable, n)
167+ .. function :: batched(iterable, n, *, strict=False )
168168
169169 Batch data from the *iterable * into tuples of length *n *. The last
170170 batch may be shorter than *n *.
171171
172+ If *strict * is true, will raise a :exc: `ValueError ` if the final
173+ batch is shorter than *n *.
174+
172175 Loops over the input iterable and accumulates data into tuples up to
173176 size *n *. The input is consumed lazily, just enough to fill a batch.
174177 The result is yielded as soon as the batch is full or when the input
@@ -190,16 +193,21 @@ loops that truncate the stream.
190193
191194 Roughly equivalent to::
192195
193- def batched(iterable, n):
196+ def batched(iterable, n, *, strict=False ):
194197 # batched('ABCDEFG', 3) --> ABC DEF G
195198 if n < 1:
196199 raise ValueError('n must be at least one')
197200 it = iter(iterable)
198201 while batch := tuple(islice(it, n)):
202+ if strict and len(batch) != n:
203+ raise ValueError('batched(): incomplete batch')
199204 yield batch
200205
201206 .. versionadded :: 3.12
202207
208+ .. versionchanged :: 3.13
209+ Added the *strict * option.
210+
203211
204212.. function :: chain(*iterables)
205213
@@ -1036,10 +1044,15 @@ The following recipes have a more mathematical flavor:
10361044 # sum_of_squares([10, 20, 30]) -> 1400
10371045 return math.sumprod(*tee(it))
10381046
1039- def transpose(it):
1040- "Swap the rows and columns of the input."
1047+ def reshape(matrix, cols):
1048+ "Reshape a 2-D matrix to have a given number of columns."
1049+ # reshape([(0, 1), (2, 3), (4, 5)], 3) --> (0, 1, 2), (3, 4, 5)
1050+ return batched(chain.from_iterable(matrix), cols, strict=True)
1051+
1052+ def transpose(matrix):
1053+ "Swap the rows and columns of a 2-D matrix."
10411054 # transpose([(1, 2, 3), (11, 22, 33)]) --> (1, 11) (2, 22) (3, 33)
1042- return zip(*it , strict=True)
1055+ return zip(*matrix , strict=True)
10431056
10441057 def matmul(m1, m2):
10451058 "Multiply two matrices."
@@ -1254,6 +1267,26 @@ The following recipes have a more mathematical flavor:
12541267 >>> sum_of_squares([10 , 20 , 30 ])
12551268 1400
12561269
1270+ >>> list (reshape([(0 , 1 ), (2 , 3 ), (4 , 5 )], 3 ))
1271+ [(0, 1, 2), (3, 4, 5)]
1272+ >>> M = [(0 , 1 , 2 , 3 ), (4 , 5 , 6 , 7 ), (8 , 9 , 10 , 11 )]
1273+ >>> list (reshape(M, 1 ))
1274+ [(0,), (1,), (2,), (3,), (4,), (5,), (6,), (7,), (8,), (9,), (10,), (11,)]
1275+ >>> list (reshape(M, 2 ))
1276+ [(0, 1), (2, 3), (4, 5), (6, 7), (8, 9), (10, 11)]
1277+ >>> list (reshape(M, 3 ))
1278+ [(0, 1, 2), (3, 4, 5), (6, 7, 8), (9, 10, 11)]
1279+ >>> list (reshape(M, 4 ))
1280+ [(0, 1, 2, 3), (4, 5, 6, 7), (8, 9, 10, 11)]
1281+ >>> list (reshape(M, 5 ))
1282+ Traceback (most recent call last):
1283+ ...
1284+ ValueError: batched(): incomplete batch
1285+ >>> list (reshape(M, 6 ))
1286+ [(0, 1, 2, 3, 4, 5), (6, 7, 8, 9, 10, 11)]
1287+ >>> list (reshape(M, 12 ))
1288+ [(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11)]
1289+
12571290 >>> list (transpose([(1 , 2 , 3 ), (11 , 22 , 33 )]))
12581291 [(1, 11), (2, 22), (3, 33)]
12591292
0 commit comments