blob: 80fe34dc8eb111df9d83b5b85e745a3508b7ad19 [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
77.. function:: count([n])
78
79 Make an iterator that returns consecutive integers starting with *n*. If not
Georg Brandl9afde1c2007-11-01 20:32:30 +000080 specified *n* defaults to zero. Often used as an argument to :func:`imap` to
81 generate consecutive data points. Also, used with :func:`izip` to add sequence
82 numbers. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +000083
84 def count(n=0):
85 while True:
86 yield n
87 n += 1
88
Georg Brandl116aa622007-08-15 14:28:22 +000089
90.. function:: cycle(iterable)
91
92 Make an iterator returning elements from the iterable and saving a copy of each.
93 When the iterable is exhausted, return elements from the saved copy. Repeats
94 indefinitely. Equivalent to::
95
96 def cycle(iterable):
97 saved = []
98 for element in iterable:
99 yield element
100 saved.append(element)
101 while saved:
102 for element in saved:
103 yield element
104
105 Note, this member of the toolkit may require significant auxiliary storage
106 (depending on the length of the iterable).
107
108
109.. function:: dropwhile(predicate, iterable)
110
111 Make an iterator that drops elements from the iterable as long as the predicate
112 is true; afterwards, returns every element. Note, the iterator does not produce
113 *any* output until the predicate first becomes false, so it may have a lengthy
114 start-up time. Equivalent to::
115
116 def dropwhile(predicate, iterable):
117 iterable = iter(iterable)
118 for x in iterable:
119 if not predicate(x):
120 yield x
121 break
122 for x in iterable:
123 yield x
124
125
126.. function:: groupby(iterable[, key])
127
128 Make an iterator that returns consecutive keys and groups from the *iterable*.
129 The *key* is a function computing a key value for each element. If not
130 specified or is ``None``, *key* defaults to an identity function and returns
131 the element unchanged. Generally, the iterable needs to already be sorted on
132 the same key function.
133
134 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
135 generates a break or new group every time the value of the key function changes
136 (which is why it is usually necessary to have sorted the data using the same key
137 function). That behavior differs from SQL's GROUP BY which aggregates common
138 elements regardless of their input order.
139
140 The returned group is itself an iterator that shares the underlying iterable
141 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
142 object is advanced, the previous group is no longer visible. So, if that data
143 is needed later, it should be stored as a list::
144
145 groups = []
146 uniquekeys = []
147 data = sorted(data, key=keyfunc)
148 for k, g in groupby(data, keyfunc):
149 groups.append(list(g)) # Store group iterator as a list
150 uniquekeys.append(k)
151
152 :func:`groupby` is equivalent to::
153
154 class groupby(object):
155 def __init__(self, iterable, key=None):
156 if key is None:
157 key = lambda x: x
158 self.keyfunc = key
159 self.it = iter(iterable)
Christian Heimes5b5e81c2007-12-31 16:14:33 +0000160 self.tgtkey = self.currkey = self.currvalue = object()
Georg Brandl116aa622007-08-15 14:28:22 +0000161 def __iter__(self):
162 return self
163 def __next__(self):
164 while self.currkey == self.tgtkey:
165 self.currvalue = next(self.it) # Exit on StopIteration
166 self.currkey = self.keyfunc(self.currvalue)
167 self.tgtkey = self.currkey
168 return (self.currkey, self._grouper(self.tgtkey))
169 def _grouper(self, tgtkey):
170 while self.currkey == tgtkey:
171 yield self.currvalue
172 self.currvalue = next(self.it) # Exit on StopIteration
173 self.currkey = self.keyfunc(self.currvalue)
174
Georg Brandl116aa622007-08-15 14:28:22 +0000175
176.. function:: ifilter(predicate, iterable)
177
178 Make an iterator that filters elements from iterable returning only those for
179 which the predicate is ``True``. If *predicate* is ``None``, return the items
180 that are true. Equivalent to::
181
182 def ifilter(predicate, iterable):
183 if predicate is None:
184 predicate = bool
185 for x in iterable:
186 if predicate(x):
187 yield x
188
189
190.. function:: ifilterfalse(predicate, iterable)
191
192 Make an iterator that filters elements from iterable returning only those for
193 which the predicate is ``False``. If *predicate* is ``None``, return the items
194 that are false. Equivalent to::
195
196 def ifilterfalse(predicate, iterable):
197 if predicate is None:
198 predicate = bool
199 for x in iterable:
200 if not predicate(x):
201 yield x
202
203
204.. function:: imap(function, *iterables)
205
206 Make an iterator that computes the function using arguments from each of the
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000207 iterables. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000208
209 def imap(function, *iterables):
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000210 iterables = [iter(it) for it in iterables)
Georg Brandl116aa622007-08-15 14:28:22 +0000211 while True:
Raymond Hettinger1dfde1d2008-01-22 23:25:35 +0000212 args = [next(it) for it in iterables]
Christian Heimes1af737c2008-01-23 08:24:23 +0000213 if function is None:
214 yield tuple(args)
215 else:
216 yield function(*args)
Georg Brandl116aa622007-08-15 14:28:22 +0000217
218
219.. function:: islice(iterable, [start,] stop [, step])
220
221 Make an iterator that returns selected elements from the iterable. If *start* is
222 non-zero, then elements from the iterable are skipped until start is reached.
223 Afterward, elements are returned consecutively unless *step* is set higher than
224 one which results in items being skipped. If *stop* is ``None``, then iteration
225 continues until the iterator is exhausted, if at all; otherwise, it stops at the
226 specified position. Unlike regular slicing, :func:`islice` does not support
227 negative values for *start*, *stop*, or *step*. Can be used to extract related
228 fields from data where the internal structure has been flattened (for example, a
229 multi-line report may list a name field on every third line). Equivalent to::
230
231 def islice(iterable, *args):
232 s = slice(*args)
Christian Heimesa37d4c62007-12-04 23:02:19 +0000233 it = iter(range(s.start or 0, s.stop or sys.maxsize, s.step or 1))
Georg Brandl116aa622007-08-15 14:28:22 +0000234 nexti = next(it)
235 for i, element in enumerate(iterable):
236 if i == nexti:
237 yield element
238 nexti = next(it)
239
240 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
241 then the step defaults to one.
242
Georg Brandl116aa622007-08-15 14:28:22 +0000243
244.. function:: izip(*iterables)
245
246 Make an iterator that aggregates elements from each of the iterables. Like
247 :func:`zip` except that it returns an iterator instead of a list. Used for
248 lock-step iteration over several iterables at a time. Equivalent to::
249
250 def izip(*iterables):
251 iterables = map(iter, iterables)
252 while iterables:
253 result = [next(it) for it in iterables]
254 yield tuple(result)
255
Georg Brandl55ac8f02007-09-01 13:51:09 +0000256 When no iterables are specified, return a zero length iterator.
Georg Brandl116aa622007-08-15 14:28:22 +0000257
Christian Heimes1af737c2008-01-23 08:24:23 +0000258 The left-to-right evaluation order of the iterables is guaranteed. This
259 makes possible an idiom for clustering a data series into n-length groups
260 using ``izip(*[iter(s)]*n)``.
Georg Brandl116aa622007-08-15 14:28:22 +0000261
Christian Heimes1af737c2008-01-23 08:24:23 +0000262<<<<<<< .working
Georg Brandl116aa622007-08-15 14:28:22 +0000263 Note, when :func:`izip` is used with unequal length inputs, subsequent
264 iteration over the longer iterables cannot reliably be continued after
265 :func:`izip` terminates. Potentially, up to one entry will be missing from
266 each of the left-over iterables. This occurs because a value is fetched from
267 each iterator in- turn, but the process ends when one of the iterators
268 terminates. This leaves the last fetched values in limbo (they cannot be
269 returned in a final, incomplete tuple and they are cannot be pushed back into
270 the iterator for retrieval with ``next(it)``). In general, :func:`izip`
271 should only be used with unequal length inputs when you don't care about
272 trailing, unmatched values from the longer iterables.
Christian Heimes1af737c2008-01-23 08:24:23 +0000273=======
274 :func:`izip` should only be used with unequal length inputs when you don't
275 care about trailing, unmatched values from the longer iterables. If those
276 values are important, use :func:`izip_longest` instead.
277>>>>>>> .merge-right.r60208
Georg Brandl116aa622007-08-15 14:28:22 +0000278
279
280.. function:: izip_longest(*iterables[, fillvalue])
281
282 Make an iterator that aggregates elements from each of the iterables. If the
283 iterables are of uneven length, missing values are filled-in with *fillvalue*.
284 Iteration continues until the longest iterable is exhausted. Equivalent to::
285
286 def izip_longest(*args, **kwds):
287 fillvalue = kwds.get('fillvalue')
288 def sentinel(counter = ([fillvalue]*(len(args)-1)).pop):
289 yield counter() # yields the fillvalue, or raises IndexError
290 fillers = repeat(fillvalue)
291 iters = [chain(it, sentinel(), fillers) for it in args]
292 try:
293 for tup in izip(*iters):
294 yield tup
295 except IndexError:
296 pass
297
298 If one of the iterables is potentially infinite, then the :func:`izip_longest`
299 function should be wrapped with something that limits the number of calls (for
300 example :func:`islice` or :func:`takewhile`).
301
Georg Brandl116aa622007-08-15 14:28:22 +0000302
303.. function:: repeat(object[, times])
304
305 Make an iterator that returns *object* over and over again. Runs indefinitely
306 unless the *times* argument is specified. Used as argument to :func:`imap` for
307 invariant parameters to the called function. Also used with :func:`izip` to
308 create an invariant part of a tuple record. Equivalent to::
309
310 def repeat(object, times=None):
311 if times is None:
312 while True:
313 yield object
314 else:
315 for i in range(times):
316 yield object
317
318
319.. function:: starmap(function, iterable)
320
Christian Heimes679db4a2008-01-18 09:56:22 +0000321 Make an iterator that computes the function using arguments obtained from
Georg Brandl116aa622007-08-15 14:28:22 +0000322 the iterable. Used instead of :func:`imap` when argument parameters are already
323 grouped in tuples from a single iterable (the data has been "pre-zipped"). The
324 difference between :func:`imap` and :func:`starmap` parallels the distinction
325 between ``function(a,b)`` and ``function(*c)``. Equivalent to::
326
327 def starmap(function, iterable):
Christian Heimes679db4a2008-01-18 09:56:22 +0000328 for args in iterable:
329 yield function(*args)
330
331 .. versionchanged:: 2.6
332 Previously, :func:`starmap` required the function arguments to be tuples.
333 Now, any iterable is allowed.
Georg Brandl116aa622007-08-15 14:28:22 +0000334
335
336.. function:: takewhile(predicate, iterable)
337
338 Make an iterator that returns elements from the iterable as long as the
339 predicate is true. Equivalent to::
340
341 def takewhile(predicate, iterable):
342 for x in iterable:
343 if predicate(x):
344 yield x
345 else:
346 break
347
348
349.. function:: tee(iterable[, n=2])
350
351 Return *n* independent iterators from a single iterable. The case where ``n==2``
352 is equivalent to::
353
354 def tee(iterable):
Christian Heimes5b5e81c2007-12-31 16:14:33 +0000355 def gen(next, data={}):
Georg Brandl116aa622007-08-15 14:28:22 +0000356 for i in count():
Christian Heimes5b5e81c2007-12-31 16:14:33 +0000357 if i in data:
358 yield data.pop(i)
Georg Brandl116aa622007-08-15 14:28:22 +0000359 else:
Christian Heimes5b5e81c2007-12-31 16:14:33 +0000360 data[i] = next()
361 yield data[i]
Georg Brandl116aa622007-08-15 14:28:22 +0000362 it = iter(iterable)
363 return (gen(it.__next__), gen(it.__next__))
364
365 Note, once :func:`tee` has made a split, the original *iterable* should not be
366 used anywhere else; otherwise, the *iterable* could get advanced without the tee
367 objects being informed.
368
369 Note, this member of the toolkit may require significant auxiliary storage
370 (depending on how much temporary data needs to be stored). In general, if one
371 iterator is going to use most or all of the data before the other iterator, it
372 is faster to use :func:`list` instead of :func:`tee`.
373
Georg Brandl116aa622007-08-15 14:28:22 +0000374
375.. _itertools-example:
376
377Examples
378--------
379
380The following examples show common uses for each tool and demonstrate ways they
381can be combined. ::
382
383 >>> amounts = [120.15, 764.05, 823.14]
384 >>> for checknum, amount in izip(count(1200), amounts):
Georg Brandl6911e3c2007-09-04 07:15:32 +0000385 ... print('Check %d is for $%.2f' % (checknum, amount))
Georg Brandl116aa622007-08-15 14:28:22 +0000386 ...
387 Check 1200 is for $120.15
388 Check 1201 is for $764.05
389 Check 1202 is for $823.14
390
391 >>> import operator
392 >>> for cube in imap(operator.pow, range(1,5), repeat(3)):
Georg Brandl6911e3c2007-09-04 07:15:32 +0000393 ... print(cube)
Georg Brandl116aa622007-08-15 14:28:22 +0000394 ...
395 1
396 8
397 27
398 64
399
400 >>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura',
401 ... '', 'martin', '', 'walter', '', 'mark']
402 >>> for name in islice(reportlines, 3, None, 2):
Georg Brandl6911e3c2007-09-04 07:15:32 +0000403 ... print(name.title())
Georg Brandl116aa622007-08-15 14:28:22 +0000404 ...
405 Alex
406 Laura
407 Martin
408 Walter
409 Mark
410
411 # Show a dictionary sorted and grouped by value
412 >>> from operator import itemgetter
413 >>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Fred Drake2e748782007-09-04 17:33:11 +0000414 >>> di = sorted(d.items(), key=itemgetter(1))
Georg Brandl116aa622007-08-15 14:28:22 +0000415 >>> for k, g in groupby(di, key=itemgetter(1)):
Georg Brandl6911e3c2007-09-04 07:15:32 +0000416 ... print(k, map(itemgetter(0), g))
Georg Brandl116aa622007-08-15 14:28:22 +0000417 ...
418 1 ['a', 'c', 'e']
419 2 ['b', 'd', 'f']
420 3 ['g']
421
422 # Find runs of consecutive numbers using groupby. The key to the solution
423 # is differencing with a range so that consecutive numbers all appear in
424 # same group.
425 >>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
426 >>> for k, g in groupby(enumerate(data), lambda t:t[0]-t[1]):
Georg Brandl6911e3c2007-09-04 07:15:32 +0000427 ... print(map(operator.itemgetter(1), g))
Georg Brandl116aa622007-08-15 14:28:22 +0000428 ...
429 [1]
430 [4, 5, 6]
431 [10]
432 [15, 16, 17, 18]
433 [22]
434 [25, 26, 27, 28]
435
436
437
438.. _itertools-recipes:
439
440Recipes
441-------
442
443This section shows recipes for creating an extended toolset using the existing
444itertools as building blocks.
445
446The extended tools offer the same high performance as the underlying toolset.
447The superior memory performance is kept by processing elements one at a time
448rather than bringing the whole iterable into memory all at once. Code volume is
449kept small by linking the tools together in a functional style which helps
450eliminate temporary variables. High speed is retained by preferring
Georg Brandl9afde1c2007-11-01 20:32:30 +0000451"vectorized" building blocks over the use of for-loops and :term:`generator`\s
452which incur interpreter overhead. ::
Georg Brandl116aa622007-08-15 14:28:22 +0000453
454 def take(n, seq):
455 return list(islice(seq, n))
456
457 def enumerate(iterable):
458 return izip(count(), iterable)
459
460 def tabulate(function):
461 "Return function(0), function(1), ..."
462 return imap(function, count())
463
Georg Brandl116aa622007-08-15 14:28:22 +0000464 def nth(iterable, n):
465 "Returns the nth item or raise StopIteration"
466 return islice(iterable, n, None).next()
467
468 def all(seq, pred=None):
469 "Returns True if pred(x) is true for every element in the iterable"
470 for elem in ifilterfalse(pred, seq):
471 return False
472 return True
473
474 def any(seq, pred=None):
475 "Returns True if pred(x) is true for at least one element in the iterable"
476 for elem in ifilter(pred, seq):
477 return True
478 return False
479
480 def no(seq, pred=None):
481 "Returns True if pred(x) is false for every element in the iterable"
482 for elem in ifilter(pred, seq):
483 return False
484 return True
485
486 def quantify(seq, pred=None):
487 "Count how many times the predicate is true in the sequence"
488 return sum(imap(pred, seq))
489
490 def padnone(seq):
491 """Returns the sequence elements and then returns None indefinitely.
492
493 Useful for emulating the behavior of the built-in map() function.
494 """
495 return chain(seq, repeat(None))
496
497 def ncycles(seq, n):
498 "Returns the sequence elements n times"
499 return chain(*repeat(seq, n))
500
501 def dotproduct(vec1, vec2):
502 return sum(imap(operator.mul, vec1, vec2))
503
504 def flatten(listOfLists):
505 return list(chain(*listOfLists))
506
507 def repeatfunc(func, times=None, *args):
508 """Repeat calls to func with specified arguments.
509
510 Example: repeatfunc(random.random)
511 """
512 if times is None:
513 return starmap(func, repeat(args))
514 else:
515 return starmap(func, repeat(args, times))
516
517 def pairwise(iterable):
518 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
519 a, b = tee(iterable)
520 next(b, None)
521 return izip(a, b)
522
523 def grouper(n, iterable, padvalue=None):
524 "grouper(3, 'abcdefg', 'x') --> ('a','b','c'), ('d','e','f'), ('g','x','x')"
525 return izip(*[chain(iterable, repeat(padvalue, n-1))]*n)
526
527
528