blob: ac6028b31a747321a35e19ec66b378d4fb926a3c [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):
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}.
Raymond Hettingerff294fe2003-12-07 13:00:25 +000079 If not specified \var{n} defaults to zero.
Raymond Hettinger96ef8112003-02-01 00:10:11 +000080 Does not currently support python long integers. Often used as an
81 argument to \function{imap()} to generate consecutive data points.
Raymond Hettingerc7d77662003-08-08 02:40:28 +000082 Also, used with \function{izip()} to add sequence numbers. Equivalent to:
Raymond Hettinger96ef8112003-02-01 00:10:11 +000083
84 \begin{verbatim}
85 def count(n=0):
Raymond Hettinger96ef8112003-02-01 00:10:11 +000086 while True:
Raymond Hettinger1b18ba42003-02-21 01:45:34 +000087 yield n
88 n += 1
Raymond Hettinger96ef8112003-02-01 00:10:11 +000089 \end{verbatim}
Raymond Hettinger2012f172003-02-07 05:32:58 +000090
91 Note, \function{count()} does not check for overflow and will return
92 negative numbers after exceeding \code{sys.maxint}. This behavior
93 may change in the future.
Raymond Hettinger96ef8112003-02-01 00:10:11 +000094\end{funcdesc}
95
Raymond Hettinger61fe64d2003-02-23 04:40:07 +000096\begin{funcdesc}{cycle}{iterable}
97 Make an iterator returning elements from the iterable and saving a
98 copy of each. When the iterable is exhausted, return elements from
99 the saved copy. Repeats indefinitely. Equivalent to:
100
101 \begin{verbatim}
102 def cycle(iterable):
103 saved = []
104 for element in iterable:
105 yield element
106 saved.append(element)
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000107 while saved:
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000108 for element in saved:
109 yield element
110 \end{verbatim}
111
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000112 Note, this member of the toolkit may require significant
113 auxiliary storage (depending on the length of the iterable).
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000114\end{funcdesc}
115
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000116\begin{funcdesc}{dropwhile}{predicate, iterable}
117 Make an iterator that drops elements from the iterable as long as
118 the predicate is true; afterwards, returns every element. Note,
119 the iterator does not produce \emph{any} output until the predicate
120 is true, so it may have a lengthy start-up time. Equivalent to:
121
122 \begin{verbatim}
123 def dropwhile(predicate, iterable):
124 iterable = iter(iterable)
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000125 for x in iterable:
126 if not predicate(x):
127 yield x
128 break
129 for x in iterable:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000130 yield x
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000131 \end{verbatim}
132\end{funcdesc}
133
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000134\begin{funcdesc}{groupby}{iterable\optional{, key}}
135 Make an iterator that returns consecutive keys and groups from the
Raymond Hettinger88e8e342004-07-11 13:20:11 +0000136 \var{iterable}. The \var{key} is a function computing a key value for each
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000137 element. If not specified or is \code{None}, \var{key} defaults to an
Andrew M. Kuchlingdb7dcff2003-12-06 22:29:43 +0000138 identity function and returns the element unchanged. Generally, the
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000139 iterable needs to already be sorted on the same key function.
140
141 The returned group is itself an iterator that shares the underlying
142 iterable with \function{groupby()}. Because the source is shared, when
143 the \function{groupby} object is advanced, the previous group is no
144 longer visible. So, if that data is needed later, it should be stored
145 as a list:
146
147 \begin{verbatim}
148 groups = []
149 uniquekeys = []
150 for k, g in groupby(data, keyfunc):
151 groups.append(list(g)) # Store group iterator as a list
152 uniquekeys.append(k)
153 \end{verbatim}
154
155 \function{groupby()} is equivalent to:
156
157 \begin{verbatim}
158 class groupby(object):
159 def __init__(self, iterable, key=None):
160 if key is None:
161 key = lambda x: x
162 self.keyfunc = key
163 self.it = iter(iterable)
164 self.tgtkey = self.currkey = self.currvalue = xrange(0)
165 def __iter__(self):
166 return self
167 def next(self):
168 while self.currkey == self.tgtkey:
169 self.currvalue = self.it.next() # Exit on StopIteration
170 self.currkey = self.keyfunc(self.currvalue)
171 self.tgtkey = self.currkey
172 return (self.currkey, self._grouper(self.tgtkey))
173 def _grouper(self, tgtkey):
174 while self.currkey == tgtkey:
175 yield self.currvalue
176 self.currvalue = self.it.next() # Exit on StopIteration
177 self.currkey = self.keyfunc(self.currvalue)
178 \end{verbatim}
179 \versionadded{2.4}
180\end{funcdesc}
181
Raymond Hettinger60eca932003-02-09 06:40:58 +0000182\begin{funcdesc}{ifilter}{predicate, iterable}
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000183 Make an iterator that filters elements from iterable returning only
Raymond Hettinger60eca932003-02-09 06:40:58 +0000184 those for which the predicate is \code{True}.
185 If \var{predicate} is \code{None}, return the items that are true.
186 Equivalent to:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000187
188 \begin{verbatim}
Raymond Hettinger60eca932003-02-09 06:40:58 +0000189 def ifilter(predicate, iterable):
190 if predicate is None:
Guido van Rossum0c9a3182003-10-20 17:01:07 +0000191 predicate = bool
Raymond Hettinger60eca932003-02-09 06:40:58 +0000192 for x in iterable:
193 if predicate(x):
194 yield x
195 \end{verbatim}
196\end{funcdesc}
197
198\begin{funcdesc}{ifilterfalse}{predicate, iterable}
199 Make an iterator that filters elements from iterable returning only
200 those for which the predicate is \code{False}.
201 If \var{predicate} is \code{None}, return the items that are false.
202 Equivalent to:
203
204 \begin{verbatim}
205 def ifilterfalse(predicate, iterable):
206 if predicate is None:
Guido van Rossum0c9a3182003-10-20 17:01:07 +0000207 predicate = bool
Raymond Hettinger60eca932003-02-09 06:40:58 +0000208 for x in iterable:
209 if not predicate(x):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000210 yield x
211 \end{verbatim}
212\end{funcdesc}
213
214\begin{funcdesc}{imap}{function, *iterables}
215 Make an iterator that computes the function using arguments from
216 each of the iterables. If \var{function} is set to \code{None}, then
217 \function{imap()} returns the arguments as a tuple. Like
218 \function{map()} but stops when the shortest iterable is exhausted
219 instead of filling in \code{None} for shorter iterables. The reason
220 for the difference is that infinite iterator arguments are typically
221 an error for \function{map()} (because the output is fully evaluated)
222 but represent a common and useful way of supplying arguments to
223 \function{imap()}.
224 Equivalent to:
225
226 \begin{verbatim}
227 def imap(function, *iterables):
228 iterables = map(iter, iterables)
229 while True:
230 args = [i.next() for i in iterables]
231 if function is None:
232 yield tuple(args)
233 else:
234 yield function(*args)
235 \end{verbatim}
236\end{funcdesc}
237
238\begin{funcdesc}{islice}{iterable, \optional{start,} stop \optional{, step}}
239 Make an iterator that returns selected elements from the iterable.
240 If \var{start} is non-zero, then elements from the iterable are skipped
241 until start is reached. Afterward, elements are returned consecutively
242 unless \var{step} is set higher than one which results in items being
Raymond Hettinger341deb72003-05-02 19:44:20 +0000243 skipped. If \var{stop} is \code{None}, then iteration continues until
244 the iterator is exhausted, if at all; otherwise, it stops at the specified
245 position. Unlike regular slicing,
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000246 \function{islice()} does not support negative values for \var{start},
247 \var{stop}, or \var{step}. Can be used to extract related fields
248 from data where the internal structure has been flattened (for
249 example, a multi-line report may list a name field on every
250 third line). Equivalent to:
251
252 \begin{verbatim}
Raymond Hettinger1b2e0d92005-03-27 20:19:05 +0000253 def islice(iterable, *args):
Raymond Hettinger341deb72003-05-02 19:44:20 +0000254 s = slice(*args)
Raymond Hettingerfdf3bd62005-03-27 20:11:44 +0000255 it = iter(xrange(s.start or 0, s.stop or sys.maxint, s.step or 1))
256 nexti = it.next()
257 for i, element in enumerate(iterable):
258 if i == nexti:
259 yield element
260 nexti = it.next()
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000261 \end{verbatim}
Raymond Hettingerb2594052004-12-05 09:25:51 +0000262
263 If \var{start} is \code{None}, then iteration starts at zero.
264 If \var{step} is \code{None}, then the step defaults to one.
265 \versionchanged[accept \code{None} values for default \var{start} and
266 \var{step}]{2.5}
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000267\end{funcdesc}
268
269\begin{funcdesc}{izip}{*iterables}
270 Make an iterator that aggregates elements from each of the iterables.
271 Like \function{zip()} except that it returns an iterator instead of
272 a list. Used for lock-step iteration over several iterables at a
273 time. Equivalent to:
274
275 \begin{verbatim}
276 def izip(*iterables):
277 iterables = map(iter, iterables)
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000278 while iterables:
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000279 result = [it.next() for it in iterables]
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000280 yield tuple(result)
281 \end{verbatim}
Raymond Hettingerb5a42082003-08-08 05:10:41 +0000282
283 \versionchanged[When no iterables are specified, returns a zero length
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000284 iterator instead of raising a \exception{TypeError}
285 exception]{2.4}
286
287 Note, the left-to-right evaluation order of the iterables is guaranteed.
288 This makes possible an idiom for clustering a data series into n-length
289 groups using \samp{izip(*[iter(s)]*n)}. For data that doesn't fit
290 n-length groups exactly, the last tuple can be pre-padded with fill
291 values using \samp{izip(*[chain(s, [None]*(n-1))]*n)}.
292
293 Note, when \function{izip()} is used with unequal length inputs, subsequent
294 iteration over the longer iterables cannot reliably be continued after
295 \function{izip()} terminates. Potentially, up to one entry will be missing
296 from each of the left-over iterables. This occurs because a value is fetched
297 from each iterator in-turn, but the process ends when one of the iterators
298 terminates. This leaves the last fetched values in limbo (they cannot be
299 returned in a final, incomplete tuple and they are cannot be pushed back
300 into the iterator for retrieval with \code{it.next()}). In general,
301 \function{izip()} should only be used with unequal length inputs when you
302 don't care about trailing, unmatched values from the longer iterables.
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000303\end{funcdesc}
304
Thomas Wouterscf297e42007-02-23 15:07:44 +0000305\begin{funcdesc}{izip_longest}{*iterables\optional{, fillvalue}}
306 Make an iterator that aggregates elements from each of the iterables.
307 If the iterables are of uneven length, missing values are filled-in
308 with \var{fillvalue}. Iteration continues until the longest iterable
309 is exhausted. Equivalent to:
310
311 \begin{verbatim}
312 def izip_longest(*args, **kwds):
313 fillvalue = kwds.get('fillvalue')
314 def sentinel(counter = ([fillvalue]*(len(args)-1)).pop):
315 yield counter() # yields the fillvalue, or raises IndexError
316 fillers = repeat(fillvalue)
317 iters = [chain(it, sentinel(), fillers) for it in args]
318 try:
319 for tup in izip(*iters):
320 yield tup
321 except IndexError:
322 pass
323 \end{verbatim}
324
325 If one of the iterables is potentially infinite, then the
326 \function{izip_longest()} function should be wrapped with something
327 that limits the number of calls (for example \function{islice()} or
328 \function{take()}).
329 \versionadded{2.6}
330\end{funcdesc}
331
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000332\begin{funcdesc}{repeat}{object\optional{, times}}
Raymond Hettinger1b18ba42003-02-21 01:45:34 +0000333 Make an iterator that returns \var{object} over and over again.
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000334 Runs indefinitely unless the \var{times} argument is specified.
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000335 Used as argument to \function{imap()} for invariant parameters
Raymond Hettinger1b18ba42003-02-21 01:45:34 +0000336 to the called function. Also used with \function{izip()} to create
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000337 an invariant part of a tuple record. Equivalent to:
338
339 \begin{verbatim}
Raymond Hettinger61fe64d2003-02-23 04:40:07 +0000340 def repeat(object, times=None):
341 if times is None:
342 while True:
343 yield object
344 else:
345 for i in xrange(times):
346 yield object
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000347 \end{verbatim}
348\end{funcdesc}
349
350\begin{funcdesc}{starmap}{function, iterable}
351 Make an iterator that computes the function using arguments tuples
352 obtained from the iterable. Used instead of \function{imap()} when
353 argument parameters are already grouped in tuples from a single iterable
354 (the data has been ``pre-zipped''). The difference between
Raymond Hettinger1b18ba42003-02-21 01:45:34 +0000355 \function{imap()} and \function{starmap()} parallels the distinction
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000356 between \code{function(a,b)} and \code{function(*c)}.
357 Equivalent to:
358
359 \begin{verbatim}
360 def starmap(function, iterable):
361 iterable = iter(iterable)
362 while True:
363 yield function(*iterable.next())
364 \end{verbatim}
365\end{funcdesc}
366
367\begin{funcdesc}{takewhile}{predicate, iterable}
368 Make an iterator that returns elements from the iterable as long as
369 the predicate is true. Equivalent to:
370
371 \begin{verbatim}
372 def takewhile(predicate, iterable):
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000373 for x in iterable:
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000374 if predicate(x):
375 yield x
376 else:
377 break
378 \end{verbatim}
379\end{funcdesc}
380
Raymond Hettingerad983e72003-11-12 14:32:26 +0000381\begin{funcdesc}{tee}{iterable\optional{, n=2}}
382 Return \var{n} independent iterators from a single iterable.
Raymond Hettinger88e8e342004-07-11 13:20:11 +0000383 The case where \code{n==2} is equivalent to:
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000384
385 \begin{verbatim}
386 def tee(iterable):
387 def gen(next, data={}, cnt=[0]):
388 for i in count():
389 if i == cnt[0]:
390 item = data[i] = next()
391 cnt[0] += 1
392 else:
393 item = data.pop(i)
394 yield item
395 it = iter(iterable)
396 return (gen(it.next), gen(it.next))
397 \end{verbatim}
398
Raymond Hettingerad983e72003-11-12 14:32:26 +0000399 Note, once \function{tee()} has made a split, the original \var{iterable}
400 should not be used anywhere else; otherwise, the \var{iterable} could get
401 advanced without the tee objects being informed.
402
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000403 Note, this member of the toolkit may require significant auxiliary
404 storage (depending on how much temporary data needs to be stored).
Andrew M. Kuchling34358202003-12-18 13:28:35 +0000405 In general, if one iterator is going to use most or all of the data before
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000406 the other iterator, it is faster to use \function{list()} instead of
407 \function{tee()}.
408 \versionadded{2.4}
409\end{funcdesc}
410
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000411
412\subsection{Examples \label{itertools-example}}
413
414The following examples show common uses for each tool and
415demonstrate ways they can be combined.
416
417\begin{verbatim}
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000418
419>>> amounts = [120.15, 764.05, 823.14]
420>>> for checknum, amount in izip(count(1200), amounts):
421... print 'Check %d is for $%.2f' % (checknum, amount)
422...
423Check 1200 is for $120.15
424Check 1201 is for $764.05
425Check 1202 is for $823.14
426
427>>> import operator
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000428>>> for cube in imap(operator.pow, xrange(1,5), repeat(3)):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000429... print cube
430...
4311
4328
43327
Raymond Hettingerd7911a32004-05-01 08:31:36 +000043464
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000435
436>>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura',
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000437 '', 'martin', '', 'walter', '', 'mark']
Raymond Hettinger3567a872003-06-28 05:44:36 +0000438>>> for name in islice(reportlines, 3, None, 2):
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000439... print name.title()
440...
441Alex
442Laura
443Martin
444Walter
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000445Mark
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000446
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000447# Show a dictionary sorted and grouped by value
448>>> from operator import itemgetter
449>>> d = dict(a=1, b=2, c=1, d=2, e=1, f=2, g=3)
Raymond Hettinger64958a12003-12-17 20:43:33 +0000450>>> di = sorted(d.iteritems(), key=itemgetter(1))
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000451>>> for k, g in groupby(di, key=itemgetter(1)):
452... print k, map(itemgetter(0), g)
453...
4541 ['a', 'c', 'e']
4552 ['b', 'd', 'f']
4563 ['g']
457
Raymond Hettinger734fb572004-01-20 20:04:40 +0000458# Find runs of consecutive numbers using groupby. The key to the solution
459# is differencing with a range so that consecutive numbers all appear in
460# same group.
461>>> data = [ 1, 4,5,6, 10, 15,16,17,18, 22, 25,26,27,28]
462>>> for k, g in groupby(enumerate(data), lambda (i,x):i-x):
463... print map(operator.itemgetter(1), g)
464...
465[1]
466[4, 5, 6]
467[10]
468[15, 16, 17, 18]
469[22]
470[25, 26, 27, 28]
Raymond Hettingerd25c1c62003-12-06 16:23:06 +0000471
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000472\end{verbatim}
473
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000474
475\subsection{Recipes \label{itertools-recipes}}
476
477This section shows recipes for creating an extended toolset using the
478existing itertools as building blocks.
479
480The extended tools offer the same high performance as the underlying
481toolset. The superior memory performance is kept by processing elements one
482at a time rather than bringing the whole iterable into memory all at once.
483Code volume is kept small by linking the tools together in a functional style
484which helps eliminate temporary variables. High speed is retained by
485preferring ``vectorized'' building blocks over the use of for-loops and
486generators which incur interpreter overhead.
487
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000488
489\begin{verbatim}
Raymond Hettingera098b332003-09-08 23:58:40 +0000490def take(n, seq):
491 return list(islice(seq, n))
492
Raymond Hettinger9e386412003-08-25 05:06:09 +0000493def enumerate(iterable):
494 return izip(count(), iterable)
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000495
Raymond Hettinger9e386412003-08-25 05:06:09 +0000496def tabulate(function):
497 "Return function(0), function(1), ..."
498 return imap(function, count())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000499
Raymond Hettinger9e386412003-08-25 05:06:09 +0000500def iteritems(mapping):
501 return izip(mapping.iterkeys(), mapping.itervalues())
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000502
Raymond Hettinger9e386412003-08-25 05:06:09 +0000503def nth(iterable, n):
Thomas Wouters89f507f2006-12-13 04:49:30 +0000504 "Returns the nth item or raise IndexError"
505 return list(islice(iterable, n, n+1))[0]
Raymond Hettinger60eca932003-02-09 06:40:58 +0000506
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000507def all(seq, pred=None):
508 "Returns True if pred(x) is true for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000509 for elem in ifilterfalse(pred, seq):
510 return False
511 return True
Raymond Hettinger60eca932003-02-09 06:40:58 +0000512
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000513def any(seq, pred=None):
514 "Returns True if pred(x) is true for at least one element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000515 for elem in ifilter(pred, seq):
516 return True
517 return False
Raymond Hettinger60eca932003-02-09 06:40:58 +0000518
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000519def no(seq, pred=None):
520 "Returns True if pred(x) is false for every element in the iterable"
Raymond Hettinger4533f1f2004-09-23 07:27:39 +0000521 for elem in ifilter(pred, seq):
522 return False
523 return True
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000524
Raymond Hettingerf77d0332005-03-11 22:17:30 +0000525def quantify(seq, pred=None):
526 "Count how many times the predicate is true in the sequence"
Raymond Hettinger9e386412003-08-25 05:06:09 +0000527 return sum(imap(pred, seq))
Raymond Hettingerc7d77662003-08-08 02:40:28 +0000528
Raymond Hettinger9e386412003-08-25 05:06:09 +0000529def padnone(seq):
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000530 """Returns the sequence elements and then returns None indefinitely.
531
532 Useful for emulating the behavior of the built-in map() function.
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000533 """
Raymond Hettinger9e386412003-08-25 05:06:09 +0000534 return chain(seq, repeat(None))
Raymond Hettinger863983e2003-04-23 00:09:42 +0000535
Raymond Hettinger9e386412003-08-25 05:06:09 +0000536def ncycles(seq, n):
537 "Returns the sequence elements n times"
538 return chain(*repeat(seq, n))
Raymond Hettinger863983e2003-04-23 00:09:42 +0000539
Raymond Hettinger9e386412003-08-25 05:06:09 +0000540def dotproduct(vec1, vec2):
541 return sum(imap(operator.mul, vec1, vec2))
Raymond Hettinger863983e2003-04-23 00:09:42 +0000542
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000543def flatten(listOfLists):
544 return list(chain(*listOfLists))
545
546def repeatfunc(func, times=None, *args):
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000547 """Repeat calls to func with specified arguments.
548
549 Example: repeatfunc(random.random)
Raymond Hettingerd7911a32004-05-01 08:31:36 +0000550 """
Raymond Hettinger6a5b0272003-10-24 08:45:23 +0000551 if times is None:
552 return starmap(func, repeat(args))
553 else:
554 return starmap(func, repeat(args, times))
555
Raymond Hettingerd591f662003-10-26 15:34:50 +0000556def pairwise(iterable):
557 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
558 a, b = tee(iterable)
Raymond Hettingerad983e72003-11-12 14:32:26 +0000559 try:
560 b.next()
561 except StopIteration:
562 pass
563 return izip(a, b)
Raymond Hettingerbefa37d2003-06-18 19:25:37 +0000564
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000565def grouper(n, iterable, padvalue=None):
566 "grouper(3, 'abcdefg', 'x') --> ('a','b','c'), ('d','e','f'), ('g','x','x')"
567 return izip(*[chain(iterable, repeat(padvalue, n-1))]*n)
568
569
Raymond Hettinger96ef8112003-02-01 00:10:11 +0000570\end{verbatim}