blob: cd867b00ac9602e8a6ed239698563a84c64f9660 [file] [log] [blame]
Raymond Hettinger96ef8112003-02-01 00:10:11 +00001\section{\module{itertools} ---
2 Functions creating iterators for efficient looping}
3
4\declaremodule{standard}{itertools}
5\modulesynopsis{Functions creating iterators for efficient looping.}
6\moduleauthor{Raymond Hettinger}{python@rcn.com}
7\sectionauthor{Raymond Hettinger}{python@rcn.com}
8\versionadded{2.3}
9
10
11This module implements a number of iterator building blocks inspired
12by constructs from the Haskell and SML programming languages. Each
13has been recast in a form suitable for Python.
14
Raymond Hettinger60eca932003-02-09 06:40:58 +000015The module standardizes a core set of fast, memory efficient tools
16that are useful by themselves or in combination. Standardization helps
17avoid the readability and reliability problems which arise when many
18different individuals create their own slightly varying implementations,
19each with their own quirks and naming conventions.
Raymond Hettinger96ef8112003-02-01 00:10:11 +000020
Raymond Hettinger1b18ba42003-02-21 01:45:34 +000021The tools are designed to combine readily with one another. This makes
Raymond Hettinger60eca932003-02-09 06:40:58 +000022it easy to construct more specialized tools succinctly and efficiently
23in pure Python.
Raymond Hettinger96ef8112003-02-01 00:10:11 +000024
Raymond Hettinger1b18ba42003-02-21 01:45:34 +000025For instance, SML provides a tabulation tool: \code{tabulate(f)}
Raymond Hettinger60eca932003-02-09 06:40:58 +000026which produces a sequence \code{f(0), f(1), ...}. This toolbox
27provides \function{imap()} and \function{count()} which can be combined
Raymond Hettinger1b18ba42003-02-21 01:45:34 +000028to form \code{imap(f, count())} and produce an equivalent result.
Raymond Hettinger96ef8112003-02-01 00:10:11 +000029
Raymond Hettinger863983e2003-04-23 00:09:42 +000030Likewise, the functional tools are designed to work well with the
31high-speed functions provided by the \refmodule{operator} module.
32
33The module author welcomes suggestions for other basic building blocks
34to be added to future versions of the module.
35
Raymond Hettinger60eca932003-02-09 06:40:58 +000036Whether cast in pure python form or C code, tools that use iterators
37are more memory efficient (and faster) than their list based counterparts.
38Adopting the principles of just-in-time manufacturing, they create
39data when and where needed instead of consuming memory with the
40computer equivalent of ``inventory''.
Raymond Hettinger96ef8112003-02-01 00:10:11 +000041
Raymond Hettinger863983e2003-04-23 00:09:42 +000042The performance advantage of iterators becomes more acute as the number
43of elements increases -- at some point, lists grow large enough to
Raymond Hettinger7e431102003-09-22 15:00:55 +000044severely impact memory cache performance and start running slowly.
Raymond Hettinger96ef8112003-02-01 00:10:11 +000045
46\begin{seealso}
47 \seetext{The Standard ML Basis Library,
48 \citetitle[http://www.standardml.org/Basis/]
49 {The Standard ML Basis Library}.}
50
51 \seetext{Haskell, A Purely Functional Language,
52 \citetitle[http://www.haskell.org/definition/]
53 {Definition of Haskell and the Standard Libraries}.}
54\end{seealso}
55
56
57\subsection{Itertool functions \label{itertools-functions}}
58
59The following module functions all construct and return iterators.
60Some provide streams of infinite length, so they should only be accessed
61by functions or loops that truncate the stream.
62
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000063\begin{funcdesc}{chain}{*iterables}
64 Make an iterator that returns elements from the first iterable until
65 it is exhausted, then proceeds to the next iterable, until all of the
66 iterables are exhausted. Used for treating consecutive sequences as
67 a single sequence. Equivalent to:
68
69 \begin{verbatim}
70 def chain(*iterables):
71 for it in iterables:
72 for element in it:
73 yield element
74 \end{verbatim}
75\end{funcdesc}
76
Raymond Hettinger96ef8112003-02-01 00:10:11 +000077\begin{funcdesc}{count}{\optional{n}}
78 Make an iterator that returns consecutive integers starting with \var{n}.
79 Does not currently support python long integers. Often used as an
80 argument to \function{imap()} to generate consecutive data points.
Raymond Hettingerc7d77662003-08-08 02:40:28 +000081 Also, used with \function{izip()} to add sequence numbers. Equivalent to:
Raymond Hettinger96ef8112003-02-01 00:10:11 +000082
83 \begin{verbatim}
84 def count(n=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +000085 while True:
Raymond Hettinger1b18ba42003-02-21 01:45:34 +000086 yield n
87 n += 1
Raymond Hettinger96ef8112003-02-01 00:10:11 +000088 \end{verbatim}
Raymond Hettinger2012f172003-02-07 05:32:58 +000089
90 Note, \function{count()} does not check for overflow and will return
91 negative numbers after exceeding \code{sys.maxint}. This behavior
92 may change in the future.
Raymond Hettinger96ef8112003-02-01 00:10:11 +000093\end{funcdesc}
94
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000095\begin{funcdesc}{cycle}{iterable}
96 Make an iterator returning elements from the iterable and saving a
97 copy of each. When the iterable is exhausted, return elements from
98 the saved copy. Repeats indefinitely. Equivalent to:
99
100 \begin{verbatim}
101 def cycle(iterable):
102 saved = []
103 for element in iterable:
104 yield element
105 saved.append(element)
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000106 while saved:
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000107 for element in saved:
108 yield element
109 \end{verbatim}
110
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000111 Note, this member of the toolkit may require significant
112 auxiliary storage (depending on the length of the iterable).
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000113\end{funcdesc}
114
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000115\begin{funcdesc}{dropwhile}{predicate, iterable}
116 Make an iterator that drops elements from the iterable as long as
117 the predicate is true; afterwards, returns every element. Note,
118 the iterator does not produce \emph{any} output until the predicate
119 is true, so it may have a lengthy start-up time. Equivalent to:
120
121 \begin{verbatim}
122 def dropwhile(predicate, iterable):
123 iterable = iter(iterable)
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000124 for x in iterable:
125 if not predicate(x):
126 yield x
127 break
128 for x in iterable:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000129 yield x
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000130 \end{verbatim}
131\end{funcdesc}
132
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000133\begin{funcdesc}{groupby}{iterable\optional{, key}}
134 Make an iterator that returns consecutive keys and groups from the
Andrew M. Kuchlingdb7dcff2003-12-06 22:29:43 +0000135 \var{iterable}. \var{key} is a function computing a key value for each
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000136 element. If not specified or is \code{None}, \var{key} defaults to an
Andrew M. Kuchlingdb7dcff2003-12-06 22:29:43 +0000137 identity function and returns the element unchanged. Generally, the
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000138 iterable needs to already be sorted on the same key function.
139
140 The returned group is itself an iterator that shares the underlying
141 iterable with \function{groupby()}. Because the source is shared, when
142 the \function{groupby} object is advanced, the previous group is no
143 longer visible. So, if that data is needed later, it should be stored
144 as a list:
145
146 \begin{verbatim}
147 groups = []
148 uniquekeys = []
149 for k, g in groupby(data, keyfunc):
150 groups.append(list(g)) # Store group iterator as a list
151 uniquekeys.append(k)
152 \end{verbatim}
153
154 \function{groupby()} is equivalent to:
155
156 \begin{verbatim}
157 class groupby(object):
158 def __init__(self, iterable, key=None):
159 if key is None:
160 key = lambda x: x
161 self.keyfunc = key
162 self.it = iter(iterable)
163 self.tgtkey = self.currkey = self.currvalue = xrange(0)
164 def __iter__(self):
165 return self
166 def next(self):
167 while self.currkey == self.tgtkey:
168 self.currvalue = self.it.next() # Exit on StopIteration
169 self.currkey = self.keyfunc(self.currvalue)
170 self.tgtkey = self.currkey
171 return (self.currkey, self._grouper(self.tgtkey))
172 def _grouper(self, tgtkey):
173 while self.currkey == tgtkey:
174 yield self.currvalue
175 self.currvalue = self.it.next() # Exit on StopIteration
176 self.currkey = self.keyfunc(self.currvalue)
177 \end{verbatim}
178 \versionadded{2.4}
179\end{funcdesc}
180
Raymond Hettinger60eca932003-02-09 06:40:58 +0000181\begin{funcdesc}{ifilter}{predicate, iterable}
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000182 Make an iterator that filters elements from iterable returning only
Raymond Hettinger60eca932003-02-09 06:40:58 +0000183 those for which the predicate is \code{True}.
184 If \var{predicate} is \code{None}, return the items that are true.
185 Equivalent to:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000186
187 \begin{verbatim}
Raymond Hettinger60eca932003-02-09 06:40:58 +0000188 def ifilter(predicate, iterable):
189 if predicate is None:
Guido van Rossum0c9a3182003-10-20 17:01:07 +0000190 predicate = bool
Raymond Hettinger60eca932003-02-09 06:40:58 +0000191 for x in iterable:
192 if predicate(x):
193 yield x
194 \end{verbatim}
195\end{funcdesc}
196
197\begin{funcdesc}{ifilterfalse}{predicate, iterable}
198 Make an iterator that filters elements from iterable returning only
199 those for which the predicate is \code{False}.
200 If \var{predicate} is \code{None}, return the items that are false.
201 Equivalent to:
202
203 \begin{verbatim}
204 def ifilterfalse(predicate, iterable):
205 if predicate is None:
Guido van Rossum0c9a3182003-10-20 17:01:07 +0000206 predicate = bool
Raymond Hettinger60eca932003-02-09 06:40:58 +0000207 for x in iterable:
208 if not predicate(x):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000209 yield x
210 \end{verbatim}
211\end{funcdesc}
212
213\begin{funcdesc}{imap}{function, *iterables}
214 Make an iterator that computes the function using arguments from
215 each of the iterables. If \var{function} is set to \code{None}, then
216 \function{imap()} returns the arguments as a tuple. Like
217 \function{map()} but stops when the shortest iterable is exhausted
218 instead of filling in \code{None} for shorter iterables. The reason
219 for the difference is that infinite iterator arguments are typically
220 an error for \function{map()} (because the output is fully evaluated)
221 but represent a common and useful way of supplying arguments to
222 \function{imap()}.
223 Equivalent to:
224
225 \begin{verbatim}
226 def imap(function, *iterables):
227 iterables = map(iter, iterables)
228 while True:
229 args = [i.next() for i in iterables]
230 if function is None:
231 yield tuple(args)
232 else:
233 yield function(*args)
234 \end{verbatim}
235\end{funcdesc}
236
237\begin{funcdesc}{islice}{iterable, \optional{start,} stop \optional{, step}}
238 Make an iterator that returns selected elements from the iterable.
239 If \var{start} is non-zero, then elements from the iterable are skipped
240 until start is reached. Afterward, elements are returned consecutively
241 unless \var{step} is set higher than one which results in items being
Raymond Hettinger341deb72003-05-02 19:44:20 +0000242 skipped. If \var{stop} is \code{None}, then iteration continues until
243 the iterator is exhausted, if at all; otherwise, it stops at the specified
244 position. Unlike regular slicing,
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000245 \function{islice()} does not support negative values for \var{start},
246 \var{stop}, or \var{step}. Can be used to extract related fields
247 from data where the internal structure has been flattened (for
248 example, a multi-line report may list a name field on every
249 third line). Equivalent to:
250
251 \begin{verbatim}
252 def islice(iterable, *args):
Raymond Hettinger341deb72003-05-02 19:44:20 +0000253 s = slice(*args)
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000254 next, stop, step = s.start or 0, s.stop, s.step or 1
Raymond Hettinger60eca932003-02-09 06:40:58 +0000255 for cnt, element in enumerate(iterable):
256 if cnt < next:
257 continue
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000258 if stop is not None and cnt >= stop:
Raymond Hettinger60eca932003-02-09 06:40:58 +0000259 break
260 yield element
Raymond Hettinger14ef54c2003-05-02 19:04:37 +0000261 next += step
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000262 \end{verbatim}
263\end{funcdesc}
264
265\begin{funcdesc}{izip}{*iterables}
266 Make an iterator that aggregates elements from each of the iterables.
267 Like \function{zip()} except that it returns an iterator instead of
268 a list. Used for lock-step iteration over several iterables at a
269 time. Equivalent to:
270
271 \begin{verbatim}
272 def izip(*iterables):
273 iterables = map(iter, iterables)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000274 while iterables:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000275 result = [i.next() for i in iterables]
276 yield tuple(result)
277 \end{verbatim}
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000278
279 \versionchanged[When no iterables are specified, returns a zero length
280 iterator instead of raising a TypeError exception]{2.4}
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000281\end{funcdesc}
282
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000283\begin{funcdesc}{repeat}{object\optional{, times}}
Raymond Hettinger1b18ba42003-02-21 01:45:34 +0000284 Make an iterator that returns \var{object} over and over again.
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000285 Runs indefinitely unless the \var{times} argument is specified.
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000286 Used as argument to \function{imap()} for invariant parameters
Raymond Hettinger1b18ba42003-02-21 01:45:34 +0000287 to the called function. Also used with \function{izip()} to create
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000288 an invariant part of a tuple record. Equivalent to:
289
290 \begin{verbatim}
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000291 def repeat(object, times=None):
292 if times is None:
293 while True:
294 yield object
295 else:
296 for i in xrange(times):
297 yield object
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000298 \end{verbatim}
299\end{funcdesc}
300
301\begin{funcdesc}{starmap}{function, iterable}
302 Make an iterator that computes the function using arguments tuples
303 obtained from the iterable. Used instead of \function{imap()} when
304 argument parameters are already grouped in tuples from a single iterable
305 (the data has been ``pre-zipped''). The difference between
Raymond Hettinger1b18ba42003-02-21 01:45:34 +0000306 \function{imap()} and \function{starmap()} parallels the distinction
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000307 between \code{function(a,b)} and \code{function(*c)}.
308 Equivalent to:
309
310 \begin{verbatim}
311 def starmap(function, iterable):
312 iterable = iter(iterable)
313 while True:
314 yield function(*iterable.next())
315 \end{verbatim}
316\end{funcdesc}
317
318\begin{funcdesc}{takewhile}{predicate, iterable}
319 Make an iterator that returns elements from the iterable as long as
320 the predicate is true. Equivalent to:
321
322 \begin{verbatim}
323 def takewhile(predicate, iterable):
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000324 for x in iterable:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000325 if predicate(x):
326 yield x
327 else:
328 break
329 \end{verbatim}
330\end{funcdesc}
331
Raymond Hettingerad983e72003-11-12 14:32:26 +0000332\begin{funcdesc}{tee}{iterable\optional{, n=2}}
333 Return \var{n} independent iterators from a single iterable.
334 The case where \var{n} is two is equivalent to:
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000335
336 \begin{verbatim}
337 def tee(iterable):
338 def gen(next, data={}, cnt=[0]):
339 for i in count():
340 if i == cnt[0]:
341 item = data[i] = next()
342 cnt[0] += 1
343 else:
344 item = data.pop(i)
345 yield item
346 it = iter(iterable)
347 return (gen(it.next), gen(it.next))
348 \end{verbatim}
349
Raymond Hettingerad983e72003-11-12 14:32:26 +0000350 Note, once \function{tee()} has made a split, the original \var{iterable}
351 should not be used anywhere else; otherwise, the \var{iterable} could get
352 advanced without the tee objects being informed.
353
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000354 Note, this member of the toolkit may require significant auxiliary
355 storage (depending on how much temporary data needs to be stored).
356 In general, if one iterator is going use most or all of the data before
357 the other iterator, it is faster to use \function{list()} instead of
358 \function{tee()}.
359 \versionadded{2.4}
360\end{funcdesc}
361
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000362
363\subsection{Examples \label{itertools-example}}
364
365The following examples show common uses for each tool and
366demonstrate ways they can be combined.
367
368\begin{verbatim}
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000369
370>>> amounts = [120.15, 764.05, 823.14]
371>>> for checknum, amount in izip(count(1200), amounts):
372... print 'Check %d is for $%.2f' % (checknum, amount)
373...
374Check 1200 is for $120.15
375Check 1201 is for $764.05
376Check 1202 is for $823.14
377
378>>> import operator
379>>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
380... print cube
381...
3821
3838
38427
385
386>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura',
387 '', 'martin', '', 'walter', '', 'samuele']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000388>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000389... print name.title()
390...
391Alex
392Laura
393Martin
394Walter
395Samuele
396
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000397# Show a dictionary sorted and grouped by value
398>>> from operator import itemgetter
399>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
400>>> di = list.sorted(d.iteritems(), key=itemgetter(1))
401>>> for k, g in groupby(di, key=itemgetter(1)):
402... print k, map(itemgetter(0), g)
403...
4041 ['a', 'c', 'e']
4052 ['b', 'd', 'f']
4063 ['g']
407
408
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000409\end{verbatim}
410
Raymond Hettingera098b332003-09-08 23:58:40 +0000411This section shows how itertools can be combined to create other more
412powerful itertools. Note that \function{enumerate()} and \method{iteritems()}
413already have efficient implementations in Python. They are only included here
414to illustrate how higher level tools can be created from building blocks.
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000415
416\begin{verbatim}
Raymond Hettingera098b332003-09-08 23:58:40 +0000417def take(n, seq):
418 return list(islice(seq, n))
419
Raymond Hettinger9e386412003-08-25 05:06:09 +0000420def enumerate(iterable):
421 return izip(count(), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000422
Raymond Hettinger9e386412003-08-25 05:06:09 +0000423def tabulate(function):
424 "Return function(0), function(1), ..."
425 return imap(function, count())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000426
Raymond Hettinger9e386412003-08-25 05:06:09 +0000427def iteritems(mapping):
428 return izip(mapping.iterkeys(), mapping.itervalues())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000429
Raymond Hettinger9e386412003-08-25 05:06:09 +0000430def nth(iterable, n):
431 "Returns the nth item"
432 return list(islice(iterable, n, n+1))
Raymond Hettinger60eca932003-02-09 06:40:58 +0000433
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000434def all(seq, pred=bool):
Raymond Hettinger9e386412003-08-25 05:06:09 +0000435 "Returns True if pred(x) is True for every element in the iterable"
436 return False not in imap(pred, seq)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000437
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000438def any(seq, pred=bool):
Raymond Hettinger9e386412003-08-25 05:06:09 +0000439 "Returns True if pred(x) is True at least one element in the iterable"
440 return True in imap(pred, seq)
Raymond Hettinger60eca932003-02-09 06:40:58 +0000441
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000442def no(seq, pred=bool):
Raymond Hettinger9e386412003-08-25 05:06:09 +0000443 "Returns True if pred(x) is False for every element in the iterable"
444 return True not in imap(pred, seq)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000445
Raymond Hettingerdbe3d282003-10-05 16:47:36 +0000446def quantify(seq, pred=bool):
Raymond Hettinger9e386412003-08-25 05:06:09 +0000447 "Count how many times the predicate is True in the sequence"
448 return sum(imap(pred, seq))
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000449
Raymond Hettinger9e386412003-08-25 05:06:09 +0000450def padnone(seq):
451 "Returns the sequence elements and then returns None indefinitely"
452 return chain(seq, repeat(None))
Raymond Hettinger863983e2003-04-23 00:09:42 +0000453
Raymond Hettinger9e386412003-08-25 05:06:09 +0000454def ncycles(seq, n):
455 "Returns the sequence elements n times"
456 return chain(*repeat(seq, n))
Raymond Hettinger863983e2003-04-23 00:09:42 +0000457
Raymond Hettinger9e386412003-08-25 05:06:09 +0000458def dotproduct(vec1, vec2):
459 return sum(imap(operator.mul, vec1, vec2))
Raymond Hettinger863983e2003-04-23 00:09:42 +0000460
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000461def flatten(listOfLists):
462 return list(chain(*listOfLists))
463
464def repeatfunc(func, times=None, *args):
465 "Repeat calls to func with specified arguments."
466 "Example: repeatfunc(random.random)"
467 if times is None:
468 return starmap(func, repeat(args))
469 else:
470 return starmap(func, repeat(args, times))
471
Raymond Hettingerd591f662003-10-26 15:34:50 +0000472def pairwise(iterable):
473 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
474 a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000475 try:
476 b.next()
477 except StopIteration:
478 pass
479 return izip(a, b)
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000480
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000481\end{verbatim}