blob: 32c0556b5bb3584d33244a3915619a705a8957b0 [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 Hettingerd7911a32004-05-01 08:31:36 +000036Whether cast in pure python form or compiled code, tools that use iterators
Raymond Hettinger60eca932003-02-09 06:40:58 +000037are 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):
Raymond Hettinger94a70362008-03-07 20:08:41 +000071 # chain('ABC', 'DEF') --> A B C D E F
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000072 for it in iterables:
73 for element in it:
74 yield element
75 \end{verbatim}
76\end{funcdesc}
77
Raymond Hettinger96ef8112003-02-01 00:10:11 +000078\begin{funcdesc}{count}{\optional{n}}
79 Make an iterator that returns consecutive integers starting with \var{n}.
Raymond Hettingerff294fe2003-12-07 13:00:25 +000080 If not specified \var{n} defaults to zero.
Raymond Hettinger96ef8112003-02-01 00:10:11 +000081 Does not currently support python long integers. Often used as an
82 argument to \function{imap()} to generate consecutive data points.
Raymond Hettingerc7d77662003-08-08 02:40:28 +000083 Also, used with \function{izip()} to add sequence numbers. Equivalent to:
Raymond Hettinger96ef8112003-02-01 00:10:11 +000084
85 \begin{verbatim}
86 def count(n=0):
Raymond Hettinger94a70362008-03-07 20:08:41 +000087 # count(10) --> 10 11 12 13 14 ...
Raymond Hettinger96ef8112003-02-01 00:10:11 +000088 while True:
Raymond Hettinger1b18ba42003-02-21 01:45:34 +000089 yield n
90 n += 1
Raymond Hettinger96ef8112003-02-01 00:10:11 +000091 \end{verbatim}
Raymond Hettinger2012f172003-02-07 05:32:58 +000092
93 Note, \function{count()} does not check for overflow and will return
94 negative numbers after exceeding \code{sys.maxint}. This behavior
95 may change in the future.
Raymond Hettinger96ef8112003-02-01 00:10:11 +000096\end{funcdesc}
97
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000098\begin{funcdesc}{cycle}{iterable}
99 Make an iterator returning elements from the iterable and saving a
100 copy of each. When the iterable is exhausted, return elements from
101 the saved copy. Repeats indefinitely. Equivalent to:
102
103 \begin{verbatim}
104 def cycle(iterable):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000105 # cycle('ABCD') --> A B C D A B C D A B C D ...
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000106 saved = []
107 for element in iterable:
108 yield element
109 saved.append(element)
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000110 while saved:
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000111 for element in saved:
112 yield element
113 \end{verbatim}
114
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000115 Note, this member of the toolkit may require significant
116 auxiliary storage (depending on the length of the iterable).
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000117\end{funcdesc}
118
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000119\begin{funcdesc}{dropwhile}{predicate, iterable}
120 Make an iterator that drops elements from the iterable as long as
121 the predicate is true; afterwards, returns every element. Note,
122 the iterator does not produce \emph{any} output until the predicate
Raymond Hettingerf4d9cae2007-07-14 11:31:35 +0000123 first becomes false, so it may have a lengthy start-up time. Equivalent to:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000124
125 \begin{verbatim}
126 def dropwhile(predicate, iterable):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000127 # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000128 iterable = iter(iterable)
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000129 for x in iterable:
130 if not predicate(x):
131 yield x
132 break
133 for x in iterable:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000134 yield x
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000135 \end{verbatim}
136\end{funcdesc}
137
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000138\begin{funcdesc}{groupby}{iterable\optional{, key}}
139 Make an iterator that returns consecutive keys and groups from the
Raymond Hettinger88e8e342004-07-11 13:20:11 +0000140 \var{iterable}. The \var{key} is a function computing a key value for each
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000141 element. If not specified or is \code{None}, \var{key} defaults to an
Andrew M. Kuchlingdb7dcff2003-12-06 22:29:43 +0000142 identity function and returns the element unchanged. Generally, the
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000143 iterable needs to already be sorted on the same key function.
144
145 The returned group is itself an iterator that shares the underlying
146 iterable with \function{groupby()}. Because the source is shared, when
147 the \function{groupby} object is advanced, the previous group is no
148 longer visible. So, if that data is needed later, it should be stored
149 as a list:
150
151 \begin{verbatim}
152 groups = []
153 uniquekeys = []
154 for k, g in groupby(data, keyfunc):
155 groups.append(list(g)) # Store group iterator as a list
156 uniquekeys.append(k)
157 \end{verbatim}
158
159 \function{groupby()} is equivalent to:
160
161 \begin{verbatim}
162 class groupby(object):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000163 # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
164 # [(list(g)) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000165 def __init__(self, iterable, key=None):
166 if key is None:
167 key = lambda x: x
168 self.keyfunc = key
169 self.it = iter(iterable)
170 self.tgtkey = self.currkey = self.currvalue = xrange(0)
171 def __iter__(self):
172 return self
173 def next(self):
174 while self.currkey == self.tgtkey:
175 self.currvalue = self.it.next() # Exit on StopIteration
176 self.currkey = self.keyfunc(self.currvalue)
177 self.tgtkey = self.currkey
178 return (self.currkey, self._grouper(self.tgtkey))
179 def _grouper(self, tgtkey):
180 while self.currkey == tgtkey:
181 yield self.currvalue
182 self.currvalue = self.it.next() # Exit on StopIteration
183 self.currkey = self.keyfunc(self.currvalue)
184 \end{verbatim}
185 \versionadded{2.4}
186\end{funcdesc}
187
Raymond Hettinger60eca932003-02-09 06:40:58 +0000188\begin{funcdesc}{ifilter}{predicate, iterable}
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000189 Make an iterator that filters elements from iterable returning only
Raymond Hettinger60eca932003-02-09 06:40:58 +0000190 those for which the predicate is \code{True}.
191 If \var{predicate} is \code{None}, return the items that are true.
192 Equivalent to:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000193
194 \begin{verbatim}
Raymond Hettinger60eca932003-02-09 06:40:58 +0000195 def ifilter(predicate, iterable):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000196 # ifilter(lambda x: x%2, range(10)) --> 1 3 5 7 9
Raymond Hettinger60eca932003-02-09 06:40:58 +0000197 if predicate is None:
Guido van Rossum0c9a3182003-10-20 17:01:07 +0000198 predicate = bool
Raymond Hettinger60eca932003-02-09 06:40:58 +0000199 for x in iterable:
200 if predicate(x):
201 yield x
202 \end{verbatim}
203\end{funcdesc}
204
205\begin{funcdesc}{ifilterfalse}{predicate, iterable}
206 Make an iterator that filters elements from iterable returning only
207 those for which the predicate is \code{False}.
208 If \var{predicate} is \code{None}, return the items that are false.
209 Equivalent to:
210
211 \begin{verbatim}
212 def ifilterfalse(predicate, iterable):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000213 # ifilterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
Raymond Hettinger60eca932003-02-09 06:40:58 +0000214 if predicate is None:
Guido van Rossum0c9a3182003-10-20 17:01:07 +0000215 predicate = bool
Raymond Hettinger60eca932003-02-09 06:40:58 +0000216 for x in iterable:
217 if not predicate(x):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000218 yield x
219 \end{verbatim}
220\end{funcdesc}
221
222\begin{funcdesc}{imap}{function, *iterables}
223 Make an iterator that computes the function using arguments from
224 each of the iterables. If \var{function} is set to \code{None}, then
225 \function{imap()} returns the arguments as a tuple. Like
226 \function{map()} but stops when the shortest iterable is exhausted
227 instead of filling in \code{None} for shorter iterables. The reason
228 for the difference is that infinite iterator arguments are typically
229 an error for \function{map()} (because the output is fully evaluated)
230 but represent a common and useful way of supplying arguments to
231 \function{imap()}.
232 Equivalent to:
233
234 \begin{verbatim}
235 def imap(function, *iterables):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000236 # imap(pow, (2,3,10), (5,2,3)) --> 32 9 1000
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000237 iterables = map(iter, iterables)
238 while True:
239 args = [i.next() for i in iterables]
240 if function is None:
241 yield tuple(args)
242 else:
243 yield function(*args)
244 \end{verbatim}
245\end{funcdesc}
246
247\begin{funcdesc}{islice}{iterable, \optional{start,} stop \optional{, step}}
248 Make an iterator that returns selected elements from the iterable.
249 If \var{start} is non-zero, then elements from the iterable are skipped
250 until start is reached. Afterward, elements are returned consecutively
251 unless \var{step} is set higher than one which results in items being
Raymond Hettinger341deb72003-05-02 19:44:20 +0000252 skipped. If \var{stop} is \code{None}, then iteration continues until
253 the iterator is exhausted, if at all; otherwise, it stops at the specified
254 position. Unlike regular slicing,
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000255 \function{islice()} does not support negative values for \var{start},
256 \var{stop}, or \var{step}. Can be used to extract related fields
257 from data where the internal structure has been flattened (for
258 example, a multi-line report may list a name field on every
259 third line). Equivalent to:
260
261 \begin{verbatim}
Raymond Hettinger1b2e0d92005-03-27 20:19:05 +0000262 def islice(iterable, *args):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000263 # islice('ABCDEFG', 2) --> A B
264 # islice('ABCDEFG', 2, 4) --> C D
265 # islice('ABCDEFG', 2, None) --> C D E F G
266 # islice('ABCDEFG', 0, None, 2) --> A C E G
Raymond Hettinger341deb72003-05-02 19:44:20 +0000267 s = slice(*args)
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000268 it = iter(xrange(s.start or 0, s.stop or sys.maxint, s.step or 1))
269 nexti = it.next()
270 for i, element in enumerate(iterable):
271 if i == nexti:
272 yield element
273 nexti = it.next()
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000274 \end{verbatim}
Raymond Hettingerb2594052004-12-05 09:25:51 +0000275
276 If \var{start} is \code{None}, then iteration starts at zero.
277 If \var{step} is \code{None}, then the step defaults to one.
278 \versionchanged[accept \code{None} values for default \var{start} and
279 \var{step}]{2.5}
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000280\end{funcdesc}
281
282\begin{funcdesc}{izip}{*iterables}
283 Make an iterator that aggregates elements from each of the iterables.
284 Like \function{zip()} except that it returns an iterator instead of
285 a list. Used for lock-step iteration over several iterables at a
286 time. Equivalent to:
287
288 \begin{verbatim}
289 def izip(*iterables):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000290 # izip('ABCD', 'xy') --> Ax By
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000291 iterables = map(iter, iterables)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000292 while iterables:
Raymond Hettingera531e5b2006-03-26 01:41:25 +0000293 result = [it.next() for it in iterables]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000294 yield tuple(result)
295 \end{verbatim}
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000296
297 \versionchanged[When no iterables are specified, returns a zero length
Georg Brandldb815ab2006-03-17 16:26:31 +0000298 iterator instead of raising a \exception{TypeError}
299 exception]{2.4}
Raymond Hettingera531e5b2006-03-26 01:41:25 +0000300
301 Note, the left-to-right evaluation order of the iterables is guaranteed.
302 This makes possible an idiom for clustering a data series into n-length
303 groups using \samp{izip(*[iter(s)]*n)}. For data that doesn't fit
304 n-length groups exactly, the last tuple can be pre-padded with fill
305 values using \samp{izip(*[chain(s, [None]*(n-1))]*n)}.
306
307 Note, when \function{izip()} is used with unequal length inputs, subsequent
308 iteration over the longer iterables cannot reliably be continued after
309 \function{izip()} terminates. Potentially, up to one entry will be missing
310 from each of the left-over iterables. This occurs because a value is fetched
311 from each iterator in-turn, but the process ends when one of the iterators
312 terminates. This leaves the last fetched values in limbo (they cannot be
313 returned in a final, incomplete tuple and they are cannot be pushed back
314 into the iterator for retrieval with \code{it.next()}). In general,
315 \function{izip()} should only be used with unequal length inputs when you
316 don't care about trailing, unmatched values from the longer iterables.
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000317\end{funcdesc}
318
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000319\begin{funcdesc}{repeat}{object\optional{, times}}
Raymond Hettinger1b18ba42003-02-21 01:45:34 +0000320 Make an iterator that returns \var{object} over and over again.
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000321 Runs indefinitely unless the \var{times} argument is specified.
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000322 Used as argument to \function{imap()} for invariant parameters
Raymond Hettinger1b18ba42003-02-21 01:45:34 +0000323 to the called function. Also used with \function{izip()} to create
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000324 an invariant part of a tuple record. Equivalent to:
325
326 \begin{verbatim}
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000327 def repeat(object, times=None):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000328 # repeat(10, 3) --> 10 10 10
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000329 if times is None:
330 while True:
331 yield object
332 else:
333 for i in xrange(times):
334 yield object
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000335 \end{verbatim}
336\end{funcdesc}
337
338\begin{funcdesc}{starmap}{function, iterable}
339 Make an iterator that computes the function using arguments tuples
340 obtained from the iterable. Used instead of \function{imap()} when
341 argument parameters are already grouped in tuples from a single iterable
342 (the data has been ``pre-zipped''). The difference between
Raymond Hettinger1b18ba42003-02-21 01:45:34 +0000343 \function{imap()} and \function{starmap()} parallels the distinction
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000344 between \code{function(a,b)} and \code{function(*c)}.
345 Equivalent to:
346
347 \begin{verbatim}
348 def starmap(function, iterable):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000349 # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000350 iterable = iter(iterable)
351 while True:
352 yield function(*iterable.next())
353 \end{verbatim}
354\end{funcdesc}
355
356\begin{funcdesc}{takewhile}{predicate, iterable}
357 Make an iterator that returns elements from the iterable as long as
358 the predicate is true. Equivalent to:
359
360 \begin{verbatim}
361 def takewhile(predicate, iterable):
Raymond Hettinger94a70362008-03-07 20:08:41 +0000362 # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000363 for x in iterable:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000364 if predicate(x):
365 yield x
366 else:
367 break
368 \end{verbatim}
369\end{funcdesc}
370
Raymond Hettingerad983e72003-11-12 14:32:26 +0000371\begin{funcdesc}{tee}{iterable\optional{, n=2}}
372 Return \var{n} independent iterators from a single iterable.
Raymond Hettinger88e8e342004-07-11 13:20:11 +0000373 The case where \code{n==2} is equivalent to:
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000374
375 \begin{verbatim}
376 def tee(iterable):
377 def gen(next, data={}, cnt=[0]):
378 for i in count():
379 if i == cnt[0]:
380 item = data[i] = next()
381 cnt[0] += 1
382 else:
383 item = data.pop(i)
384 yield item
385 it = iter(iterable)
386 return (gen(it.next), gen(it.next))
387 \end{verbatim}
388
Raymond Hettingerad983e72003-11-12 14:32:26 +0000389 Note, once \function{tee()} has made a split, the original \var{iterable}
390 should not be used anywhere else; otherwise, the \var{iterable} could get
391 advanced without the tee objects being informed.
392
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000393 Note, this member of the toolkit may require significant auxiliary
394 storage (depending on how much temporary data needs to be stored).
Andrew M. Kuchling34358202003-12-18 13:28:35 +0000395 In general, if one iterator is going to use most or all of the data before
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000396 the other iterator, it is faster to use \function{list()} instead of
397 \function{tee()}.
398 \versionadded{2.4}
399\end{funcdesc}
400
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000401
402\subsection{Examples \label{itertools-example}}
403
404The following examples show common uses for each tool and
405demonstrate ways they can be combined.
406
407\begin{verbatim}
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000408
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000409# Show a dictionary sorted and grouped by value
410>>> from operator import itemgetter
411>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Raymond Hettinger64958a12003-12-17 20:43:33 +0000412>>> di = sorted(d.iteritems(), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000413>>> for k, g in groupby(di, key=itemgetter(1)):
414... print k, map(itemgetter(0), g)
415...
4161 ['a', 'c', 'e']
4172 ['b', 'd', 'f']
4183 ['g']
419
Raymond Hettinger734fb572004-01-20 20:04:40 +0000420# Find runs of consecutive numbers using groupby. The key to the solution
421# is differencing with a range so that consecutive numbers all appear in
422# same group.
423>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
424>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
425... print map(operator.itemgetter(1), g)
426...
427[1]
428[4, 5, 6]
429[10]
430[15, 16, 17, 18]
431[22]
432[25, 26, 27, 28]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000433
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000434\end{verbatim}
435
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000436
437\subsection{Recipes \label{itertools-recipes}}
438
439This section shows recipes for creating an extended toolset using the
440existing itertools as building blocks.
441
442The extended tools offer the same high performance as the underlying
443toolset. The superior memory performance is kept by processing elements one
444at a time rather than bringing the whole iterable into memory all at once.
445Code volume is kept small by linking the tools together in a functional style
446which helps eliminate temporary variables. High speed is retained by
447preferring ``vectorized'' building blocks over the use of for-loops and
448generators which incur interpreter overhead.
449
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000450
451\begin{verbatim}
Raymond Hettingera098b332003-09-08 23:58:40 +0000452def take(n, seq):
453 return list(islice(seq, n))
454
Raymond Hettinger9e386412003-08-25 05:06:09 +0000455def enumerate(iterable):
456 return izip(count(), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000457
Raymond Hettinger9e386412003-08-25 05:06:09 +0000458def tabulate(function):
459 "Return function(0), function(1), ..."
460 return imap(function, count())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000461
Raymond Hettinger9e386412003-08-25 05:06:09 +0000462def iteritems(mapping):
463 return izip(mapping.iterkeys(), mapping.itervalues())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000464
Raymond Hettinger9e386412003-08-25 05:06:09 +0000465def nth(iterable, n):
Raymond Hettingerf4d9cae2007-07-14 11:31:35 +0000466 "Returns the nth item or raise StopIteration"
467 return islice(iterable, n, None).next()
Raymond Hettinger60eca932003-02-09 06:40:58 +0000468
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000469def all(seq, pred=None):
470 "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000471 for elem in ifilterfalse(pred, seq):
472 return False
473 return True
Raymond Hettinger60eca932003-02-09 06:40:58 +0000474
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000475def any(seq, pred=None):
476 "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000477 for elem in ifilter(pred, seq):
478 return True
479 return False
Raymond Hettinger60eca932003-02-09 06:40:58 +0000480
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000481def no(seq, pred=None):
482 "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000483 for elem in ifilter(pred, seq):
484 return False
485 return True
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000486
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000487def quantify(seq, pred=None):
488 "Count how many times the predicate is true in the sequence"
Raymond Hettinger9e386412003-08-25 05:06:09 +0000489 return sum(imap(pred, seq))
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000490
Raymond Hettinger9e386412003-08-25 05:06:09 +0000491def padnone(seq):
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000492 """Returns the sequence elements and then returns None indefinitely.
493
494 Useful for emulating the behavior of the built-in map() function.
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000495 """
Raymond Hettinger9e386412003-08-25 05:06:09 +0000496 return chain(seq, repeat(None))
Raymond Hettinger863983e2003-04-23 00:09:42 +0000497
Raymond Hettinger9e386412003-08-25 05:06:09 +0000498def ncycles(seq, n):
499 "Returns the sequence elements n times"
500 return chain(*repeat(seq, n))
Raymond Hettinger863983e2003-04-23 00:09:42 +0000501
Raymond Hettinger9e386412003-08-25 05:06:09 +0000502def dotproduct(vec1, vec2):
503 return sum(imap(operator.mul, vec1, vec2))
Raymond Hettinger863983e2003-04-23 00:09:42 +0000504
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000505def flatten(listOfLists):
506 return list(chain(*listOfLists))
507
508def repeatfunc(func, times=None, *args):
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000509 """Repeat calls to func with specified arguments.
510
511 Example: repeatfunc(random.random)
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000512 """
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000513 if times is None:
514 return starmap(func, repeat(args))
515 else:
516 return starmap(func, repeat(args, times))
517
Raymond Hettingerd591f662003-10-26 15:34:50 +0000518def pairwise(iterable):
519 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
520 a, b = tee(iterable)
Raymond Hettinger94a70362008-03-07 20:08:41 +0000521 for elem in b:
522 break
Raymond Hettingerad983e72003-11-12 14:32:26 +0000523 return izip(a, b)
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000524
Raymond Hettingera531e5b2006-03-26 01:41:25 +0000525def grouper(n, iterable, padvalue=None):
526 "grouper(3, 'abcdefg', 'x') --> ('a','b','c'), ('d','e','f'), ('g','x','x')"
527 return izip(*[chain(iterable, repeat(padvalue, n-1))]*n)
528
Raymond Hettinger42f4cfa2007-03-20 21:12:23 +0000529def reverse_map(d):
530 "Return a new dict with swapped keys and values"
531 return dict(izip(d.itervalues(), d))
Raymond Hettingera531e5b2006-03-26 01:41:25 +0000532
Raymond Hettinger94a70362008-03-07 20:08:41 +0000533def roundrobin(*iterables):
534 "roundrobin('abc', 'd', 'ef') --> 'a', 'd', 'e', 'b', 'f', 'c'"
535 # Recipe credited to George Sakkis
536 pending = len(iterables)
537 nexts = cycle(iter(it).next for it in iterables)
538 while pending:
539 try:
540 for next in nexts:
541 yield next()
542 except StopIteration:
543 pending -= 1
544 nexts = cycle(islice(nexts, pending))
545
546def powerset(iterable):
547 "powerset('ab') --> set([]), set(['a']), set(['b']), set(['a', 'b'])"
548 # Recipe credited to Eric Raymond
549 pairs = [(2**i, x) for i, x in enumerate(iterable)]
550 for n in xrange(2**len(pairs)):
551 yield set(x for m, x in pairs if m&n)
552
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000553\end{verbatim}