blob: 826fa454d2e228563953755d07f5f576d01ed878 [file] [log] [blame]
Georg Brandl116aa622007-08-15 14:28:22 +00001
2:mod:`itertools` --- Functions creating iterators for efficient looping
3=======================================================================
4
5.. module:: itertools
6 :synopsis: Functions creating iterators for efficient looping.
7.. moduleauthor:: Raymond Hettinger <python@rcn.com>
8.. sectionauthor:: Raymond Hettinger <python@rcn.com>
9
10
Georg Brandl9afde1c2007-11-01 20:32:30 +000011This module implements a number of :term:`iterator` building blocks inspired by
Georg Brandl116aa622007-08-15 14:28:22 +000012constructs from the Haskell and SML programming languages. Each has been recast
13in a form suitable for Python.
14
15The module standardizes a core set of fast, memory efficient tools that are
16useful by themselves or in combination. Standardization helps avoid the
17readability and reliability problems which arise when many different individuals
18create their own slightly varying implementations, each with their own quirks
19and naming conventions.
20
21The tools are designed to combine readily with one another. This makes it easy
22to construct more specialized tools succinctly and efficiently in pure Python.
23
24For instance, SML provides a tabulation tool: ``tabulate(f)`` which produces a
25sequence ``f(0), f(1), ...``. This toolbox provides :func:`imap` and
26:func:`count` which can be combined to form ``imap(f, count())`` and produce an
27equivalent result.
28
29Likewise, the functional tools are designed to work well with the high-speed
30functions provided by the :mod:`operator` module.
31
32The module author welcomes suggestions for other basic building blocks to be
33added to future versions of the module.
34
35Whether cast in pure python form or compiled code, tools that use iterators are
36more memory efficient (and faster) than their list based counterparts. Adopting
37the principles of just-in-time manufacturing, they create data when and where
38needed instead of consuming memory with the computer equivalent of "inventory".
39
40The performance advantage of iterators becomes more acute as the number of
41elements increases -- at some point, lists grow large enough to severely impact
42memory cache performance and start running slowly.
43
44
45.. seealso::
46
47 The Standard ML Basis Library, `The Standard ML Basis Library
48 <http://www.standardml.org/Basis/>`_.
49
50 Haskell, A Purely Functional Language, `Definition of Haskell and the Standard
51 Libraries <http://www.haskell.org/definition/>`_.
52
53
54.. _itertools-functions:
55
56Itertool functions
57------------------
58
59The following module functions all construct and return iterators. Some provide
60streams of infinite length, so they should only be accessed by functions or
61loops that truncate the stream.
62
63
64.. function:: chain(*iterables)
65
66 Make an iterator that returns elements from the first iterable until it is
67 exhausted, then proceeds to the next iterable, until all of the iterables are
68 exhausted. Used for treating consecutive sequences as a single sequence.
69 Equivalent to::
70
71 def chain(*iterables):
72 for it in iterables:
73 for element in it:
74 yield element
75
76
Christian Heimes836baa52008-02-26 08:18:30 +000077.. function:: combinations(iterable, r)
78
79 Return successive *r* length combinations of elements in the *iterable*.
80
81 Combinations are emitted in a lexicographic sort order. So, if the
82 input *iterable* is sorted, the combination tuples will be produced
83 in sorted order.
84
85 Elements are treated as unique based on their position, not on their
86 value. So if the input elements are unique, there will be no repeat
87 values within a single combination.
88
89 Each result tuple is ordered to match the input order. So, every
90 combination is a subsequence of the input *iterable*.
91
92 Example: ``combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)``
93
94 Equivalent to::
95
96 def combinations(iterable, r):
97 pool = tuple(iterable)
Christian Heimes380f7f22008-02-28 11:19:05 +000098 n = len(pool)
99 assert 0 <= r <= n
100 vec = range(r)
101 yield tuple(pool[i] for i in vec)
102 while 1:
103 for i in reversed(range(r)):
104 if vec[i] != i + n - r:
Christian Heimes836baa52008-02-26 08:18:30 +0000105 break
Christian Heimes380f7f22008-02-28 11:19:05 +0000106 else:
107 return
108 vec[i] += 1
109 for j in range(i+1, r):
110 vec[j] = vec[j-1] + 1
111 yield tuple(pool[i] for i in vec)
Christian Heimes836baa52008-02-26 08:18:30 +0000112
113 .. versionadded:: 2.6
114
Georg Brandl116aa622007-08-15 14:28:22 +0000115.. function:: count([n])
116
117 Make an iterator that returns consecutive integers starting with *n*. If not
Georg Brandl9afde1c2007-11-01 20:32:30 +0000118 specified *n* defaults to zero. Often used as an argument to :func:`imap` to
119 generate consecutive data points. Also, used with :func:`izip` to add sequence
120 numbers. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000121
122 def count(n=0):
123 while True:
124 yield n
125 n += 1
126
Georg Brandl116aa622007-08-15 14:28:22 +0000127
128.. function:: cycle(iterable)
129
130 Make an iterator returning elements from the iterable and saving a copy of each.
131 When the iterable is exhausted, return elements from the saved copy. Repeats
132 indefinitely. Equivalent to::
133
134 def cycle(iterable):
135 saved = []
136 for element in iterable:
137 yield element
138 saved.append(element)
139 while saved:
140 for element in saved:
141 yield element
142
143 Note, this member of the toolkit may require significant auxiliary storage
144 (depending on the length of the iterable).
145
146
147.. function:: dropwhile(predicate, iterable)
148
149 Make an iterator that drops elements from the iterable as long as the predicate
150 is true; afterwards, returns every element. Note, the iterator does not produce
151 *any* output until the predicate first becomes false, so it may have a lengthy
152 start-up time. Equivalent to::
153
154 def dropwhile(predicate, iterable):
155 iterable = iter(iterable)
156 for x in iterable:
157 if not predicate(x):
158 yield x
159 break
160 for x in iterable:
161 yield x
162
163
164.. function:: groupby(iterable[, key])
165
166 Make an iterator that returns consecutive keys and groups from the *iterable*.
167 The *key* is a function computing a key value for each element. If not
168 specified or is ``None``, *key* defaults to an identity function and returns
169 the element unchanged. Generally, the iterable needs to already be sorted on
170 the same key function.
171
172 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
173 generates a break or new group every time the value of the key function changes
174 (which is why it is usually necessary to have sorted the data using the same key
175 function). That behavior differs from SQL's GROUP BY which aggregates common
176 elements regardless of their input order.
177
178 The returned group is itself an iterator that shares the underlying iterable
179 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
180 object is advanced, the previous group is no longer visible. So, if that data
181 is needed later, it should be stored as a list::
182
183 groups = []
184 uniquekeys = []
185 data = sorted(data, key=keyfunc)
186 for k, g in groupby(data, keyfunc):
187 groups.append(list(g)) # Store group iterator as a list
188 uniquekeys.append(k)
189
190 :func:`groupby` is equivalent to::
191
192 class groupby(object):
193 def __init__(self, iterable, key=None):
194 if key is None:
195 key = lambda x: x
196 self.keyfunc = key
197 self.it = iter(iterable)
Christian Heimes5b5e81c2007-12-31 16:14:33 +0000198 self.tgtkey = self.currkey = self.currvalue = object()
Georg Brandl116aa622007-08-15 14:28:22 +0000199 def __iter__(self):
200 return self
201 def __next__(self):
202 while self.currkey == self.tgtkey:
203 self.currvalue = next(self.it) # Exit on StopIteration
204 self.currkey = self.keyfunc(self.currvalue)
205 self.tgtkey = self.currkey
206 return (self.currkey, self._grouper(self.tgtkey))
207 def _grouper(self, tgtkey):
208 while self.currkey == tgtkey:
209 yield self.currvalue
210 self.currvalue = next(self.it) # Exit on StopIteration
211 self.currkey = self.keyfunc(self.currvalue)
212
Georg Brandl116aa622007-08-15 14:28:22 +0000213
214.. function:: ifilter(predicate, iterable)
215
216 Make an iterator that filters elements from iterable returning only those for
217 which the predicate is ``True``. If *predicate* is ``None``, return the items
Georg Brandlf6945182008-02-01 11:56:49 +0000218 that are true. This function is the same as the built-in :func:`filter`
219 function. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000220
221 def ifilter(predicate, iterable):
222 if predicate is None:
223 predicate = bool
224 for x in iterable:
225 if predicate(x):
226 yield x
227
228
229.. function:: ifilterfalse(predicate, iterable)
230
231 Make an iterator that filters elements from iterable returning only those for
232 which the predicate is ``False``. If *predicate* is ``None``, return the items
233 that are false. Equivalent to::
234
235 def ifilterfalse(predicate, iterable):
236 if predicate is None:
237 predicate = bool
238 for x in iterable:
239 if not predicate(x):
240 yield x
241
242
243.. function:: imap(function, *iterables)
244
245 Make an iterator that computes the function using arguments from each of the
Georg Brandlf6945182008-02-01 11:56:49 +0000246 iterables. This function is the same as the built-in :func:`map` function.
247 Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000248
249 def imap(function, *iterables):
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000250 iterables = [iter(it) for it in iterables)
Georg Brandl116aa622007-08-15 14:28:22 +0000251 while True:
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000252 args = [next(it) for it in iterables]
Christian Heimes1af737c2008-01-23 08:24:23 +0000253 if function is None:
254 yield tuple(args)
255 else:
256 yield function(*args)
Georg Brandl116aa622007-08-15 14:28:22 +0000257
258
259.. function:: islice(iterable, [start,] stop [, step])
260
261 Make an iterator that returns selected elements from the iterable. If *start* is
262 non-zero, then elements from the iterable are skipped until start is reached.
263 Afterward, elements are returned consecutively unless *step* is set higher than
264 one which results in items being skipped. If *stop* is ``None``, then iteration
265 continues until the iterator is exhausted, if at all; otherwise, it stops at the
266 specified position. Unlike regular slicing, :func:`islice` does not support
267 negative values for *start*, *stop*, or *step*. Can be used to extract related
268 fields from data where the internal structure has been flattened (for example, a
269 multi-line report may list a name field on every third line). Equivalent to::
270
271 def islice(iterable, *args):
272 s = slice(*args)
Georg Brandlf6945182008-02-01 11:56:49 +0000273 it = range(s.start or 0, s.stop or sys.maxsize, s.step or 1)
Georg Brandl116aa622007-08-15 14:28:22 +0000274 nexti = next(it)
275 for i, element in enumerate(iterable):
276 if i == nexti:
277 yield element
278 nexti = next(it)
279
280 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
281 then the step defaults to one.
282
Georg Brandl116aa622007-08-15 14:28:22 +0000283
284.. function:: izip(*iterables)
285
286 Make an iterator that aggregates elements from each of the iterables. Like
287 :func:`zip` except that it returns an iterator instead of a list. Used for
288 lock-step iteration over several iterables at a time. Equivalent to::
289
290 def izip(*iterables):
291 iterables = map(iter, iterables)
292 while iterables:
293 result = [next(it) for it in iterables]
294 yield tuple(result)
295
Georg Brandl55ac8f02007-09-01 13:51:09 +0000296 When no iterables are specified, return a zero length iterator.
Georg Brandl116aa622007-08-15 14:28:22 +0000297
Christian Heimes1af737c2008-01-23 08:24:23 +0000298 The left-to-right evaluation order of the iterables is guaranteed. This
299 makes possible an idiom for clustering a data series into n-length groups
300 using ``izip(*[iter(s)]*n)``.
Georg Brandl116aa622007-08-15 14:28:22 +0000301
Christian Heimes1af737c2008-01-23 08:24:23 +0000302 :func:`izip` should only be used with unequal length inputs when you don't
303 care about trailing, unmatched values from the longer iterables. If those
304 values are important, use :func:`izip_longest` instead.
Georg Brandl116aa622007-08-15 14:28:22 +0000305
306
307.. function:: izip_longest(*iterables[, fillvalue])
308
309 Make an iterator that aggregates elements from each of the iterables. If the
310 iterables are of uneven length, missing values are filled-in with *fillvalue*.
311 Iteration continues until the longest iterable is exhausted. Equivalent to::
312
313 def izip_longest(*args, **kwds):
314 fillvalue = kwds.get('fillvalue')
315 def sentinel(counter = ([fillvalue]*(len(args)-1)).pop):
316 yield counter() # yields the fillvalue, or raises IndexError
317 fillers = repeat(fillvalue)
318 iters = [chain(it, sentinel(), fillers) for it in args]
319 try:
320 for tup in izip(*iters):
321 yield tup
322 except IndexError:
323 pass
324
325 If one of the iterables is potentially infinite, then the :func:`izip_longest`
326 function should be wrapped with something that limits the number of calls (for
327 example :func:`islice` or :func:`takewhile`).
328
Georg Brandl116aa622007-08-15 14:28:22 +0000329
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000330.. function:: product(*iterables)
331
332 Cartesian product of input iterables.
333
334 Equivalent to nested for-loops in a generator expression. For example,
335 ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``.
336
337 The leftmost iterators are in the outermost for-loop, so the output tuples
338 cycle in a manner similar to an odometer (with the rightmost element
Christian Heimes836baa52008-02-26 08:18:30 +0000339 changing on every iteration). This results in a lexicographic ordering
340 so that if the inputs iterables are sorted, the product tuples are emitted
341 in sorted order.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000342
Christian Heimes836baa52008-02-26 08:18:30 +0000343 Equivalent to the following except that the actual implementation does not
344 build-up intermediate results in memory::
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000345
346 def product(*args):
347 pools = map(tuple, args)
348 if pools:
349 result = [[]]
350 for pool in pools:
351 result = [x+[y] for x in result for y in pool]
352 for prod in result:
353 yield tuple(prod)
354
355
Georg Brandl116aa622007-08-15 14:28:22 +0000356.. function:: repeat(object[, times])
357
358 Make an iterator that returns *object* over and over again. Runs indefinitely
359 unless the *times* argument is specified. Used as argument to :func:`imap` for
360 invariant parameters to the called function. Also used with :func:`izip` to
361 create an invariant part of a tuple record. Equivalent to::
362
363 def repeat(object, times=None):
364 if times is None:
365 while True:
366 yield object
367 else:
368 for i in range(times):
369 yield object
370
371
372.. function:: starmap(function, iterable)
373
Christian Heimes679db4a2008-01-18 09:56:22 +0000374 Make an iterator that computes the function using arguments obtained from
Georg Brandl116aa622007-08-15 14:28:22 +0000375 the iterable. Used instead of :func:`imap` when argument parameters are already
376 grouped in tuples from a single iterable (the data has been "pre-zipped"). The
377 difference between :func:`imap` and :func:`starmap` parallels the distinction
378 between ``function(a,b)`` and ``function(*c)``. Equivalent to::
379
380 def starmap(function, iterable):
Christian Heimes679db4a2008-01-18 09:56:22 +0000381 for args in iterable:
382 yield function(*args)
383
384 .. versionchanged:: 2.6
385 Previously, :func:`starmap` required the function arguments to be tuples.
386 Now, any iterable is allowed.
Georg Brandl116aa622007-08-15 14:28:22 +0000387
388
389.. function:: takewhile(predicate, iterable)
390
391 Make an iterator that returns elements from the iterable as long as the
392 predicate is true. Equivalent to::
393
394 def takewhile(predicate, iterable):
395 for x in iterable:
396 if predicate(x):
397 yield x
398 else:
399 break
400
401
402.. function:: tee(iterable[, n=2])
403
404 Return *n* independent iterators from a single iterable. The case where ``n==2``
405 is equivalent to::
406
407 def tee(iterable):
Christian Heimes5b5e81c2007-12-31 16:14:33 +0000408 def gen(next, data={}):
Georg Brandl116aa622007-08-15 14:28:22 +0000409 for i in count():
Christian Heimes5b5e81c2007-12-31 16:14:33 +0000410 if i in data:
411 yield data.pop(i)
Georg Brandl116aa622007-08-15 14:28:22 +0000412 else:
Christian Heimes5b5e81c2007-12-31 16:14:33 +0000413 data[i] = next()
414 yield data[i]
Georg Brandl116aa622007-08-15 14:28:22 +0000415 it = iter(iterable)
416 return (gen(it.__next__), gen(it.__next__))
417
418 Note, once :func:`tee` has made a split, the original *iterable* should not be
419 used anywhere else; otherwise, the *iterable* could get advanced without the tee
420 objects being informed.
421
422 Note, this member of the toolkit may require significant auxiliary storage
423 (depending on how much temporary data needs to be stored). In general, if one
424 iterator is going to use most or all of the data before the other iterator, it
425 is faster to use :func:`list` instead of :func:`tee`.
426
Georg Brandl116aa622007-08-15 14:28:22 +0000427
428.. _itertools-example:
429
430Examples
431--------
432
433The following examples show common uses for each tool and demonstrate ways they
434can be combined. ::
435
436 >>> amounts = [120.15, 764.05, 823.14]
437 >>> for checknum, amount in izip(count(1200), amounts):
Georg Brandl6911e3c2007-09-04 07:15:32 +0000438 ... print('Check %d is for $%.2f' % (checknum, amount))
Georg Brandl116aa622007-08-15 14:28:22 +0000439 ...
440 Check 1200 is for $120.15
441 Check 1201 is for $764.05
442 Check 1202 is for $823.14
443
444 >>> import operator
445 >>> for cube in imap(operator.pow, range(1,5), repeat(3)):
Georg Brandl6911e3c2007-09-04 07:15:32 +0000446 ... print(cube)
Georg Brandl116aa622007-08-15 14:28:22 +0000447 ...
448 1
449 8
450 27
451 64
452
453 >>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura',
454 ... '', 'martin', '', 'walter', '', 'mark']
455 >>> for name in islice(reportlines, 3, None, 2):
Georg Brandl6911e3c2007-09-04 07:15:32 +0000456 ... print(name.title())
Georg Brandl116aa622007-08-15 14:28:22 +0000457 ...
458 Alex
459 Laura
460 Martin
461 Walter
462 Mark
463
464 # Show a dictionary sorted and grouped by value
465 >>> from operator import itemgetter
466 >>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Fred Drake2e748782007-09-04 17:33:11 +0000467 >>> di = sorted(d.items(), key=itemgetter(1))
Georg Brandl116aa622007-08-15 14:28:22 +0000468 >>> for k, g in groupby(di, key=itemgetter(1)):
Georg Brandl6911e3c2007-09-04 07:15:32 +0000469 ... print(k, map(itemgetter(0), g))
Georg Brandl116aa622007-08-15 14:28:22 +0000470 ...
471 1 ['a', 'c', 'e']
472 2 ['b', 'd', 'f']
473 3 ['g']
474
475 # Find runs of consecutive numbers using groupby. The key to the solution
476 # is differencing with a range so that consecutive numbers all appear in
477 # same group.
478 >>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
479 >>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Georg Brandl6911e3c2007-09-04 07:15:32 +0000480 ... print(map(operator.itemgetter(1), g))
Georg Brandl116aa622007-08-15 14:28:22 +0000481 ...
482 [1]
483 [4, 5, 6]
484 [10]
485 [15, 16, 17, 18]
486 [22]
487 [25, 26, 27, 28]
488
489
490
491.. _itertools-recipes:
492
493Recipes
494-------
495
496This section shows recipes for creating an extended toolset using the existing
497itertools as building blocks.
498
499The extended tools offer the same high performance as the underlying toolset.
500The superior memory performance is kept by processing elements one at a time
501rather than bringing the whole iterable into memory all at once. Code volume is
502kept small by linking the tools together in a functional style which helps
503eliminate temporary variables. High speed is retained by preferring
Georg Brandl9afde1c2007-11-01 20:32:30 +0000504"vectorized" building blocks over the use of for-loops and :term:`generator`\s
505which incur interpreter overhead. ::
Georg Brandl116aa622007-08-15 14:28:22 +0000506
507 def take(n, seq):
508 return list(islice(seq, n))
509
510 def enumerate(iterable):
511 return izip(count(), iterable)
512
513 def tabulate(function):
514 "Return function(0), function(1), ..."
515 return imap(function, count())
516
Georg Brandl116aa622007-08-15 14:28:22 +0000517 def nth(iterable, n):
518 "Returns the nth item or raise StopIteration"
519 return islice(iterable, n, None).next()
520
521 def all(seq, pred=None):
522 "Returns True if pred(x) is true for every element in the iterable"
523 for elem in ifilterfalse(pred, seq):
524 return False
525 return True
526
527 def any(seq, pred=None):
528 "Returns True if pred(x) is true for at least one element in the iterable"
529 for elem in ifilter(pred, seq):
530 return True
531 return False
532
533 def no(seq, pred=None):
534 "Returns True if pred(x) is false for every element in the iterable"
535 for elem in ifilter(pred, seq):
536 return False
537 return True
538
539 def quantify(seq, pred=None):
540 "Count how many times the predicate is true in the sequence"
541 return sum(imap(pred, seq))
542
543 def padnone(seq):
544 """Returns the sequence elements and then returns None indefinitely.
545
546 Useful for emulating the behavior of the built-in map() function.
547 """
548 return chain(seq, repeat(None))
549
550 def ncycles(seq, n):
551 "Returns the sequence elements n times"
552 return chain(*repeat(seq, n))
553
554 def dotproduct(vec1, vec2):
555 return sum(imap(operator.mul, vec1, vec2))
556
557 def flatten(listOfLists):
558 return list(chain(*listOfLists))
559
560 def repeatfunc(func, times=None, *args):
561 """Repeat calls to func with specified arguments.
562
563 Example: repeatfunc(random.random)
564 """
565 if times is None:
566 return starmap(func, repeat(args))
567 else:
568 return starmap(func, repeat(args, times))
569
570 def pairwise(iterable):
571 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
572 a, b = tee(iterable)
573 next(b, None)
574 return izip(a, b)
575
576 def grouper(n, iterable, padvalue=None):
577 "grouper(3, 'abcdefg', 'x') --> ('a','b','c'), ('d','e','f'), ('g','x','x')"
578 return izip(*[chain(iterable, repeat(padvalue, n-1))]*n)
579
Christian Heimes7b3ce6a2008-01-31 14:31:45 +0000580 def roundrobin(*iterables):
581 "roundrobin('abc', 'd', 'ef') --> 'a', 'd', 'e', 'b', 'f', 'c'"
582 # Recipe contributed by George Sakkis
583 pending = len(iterables)
584 nexts = cycle(iter(it).next for it in iterables)
585 while pending:
586 try:
587 for next in nexts:
588 yield next()
589 except StopIteration:
590 pending -= 1
591 nexts = cycle(islice(nexts, pending))
Georg Brandl116aa622007-08-15 14:28:22 +0000592
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000593 def powerset(iterable):
594 "powerset('ab') --> set([]), set(['b']), set(['a']), set(['a', 'b'])"
595 skip = object()
596 for t in product(*izip(repeat(skip), iterable)):
597 yield set(e for e in t if e is not skip)
598