blob: f489535c5361808123314e90002a546a0b8c4d24 [file] [log] [blame]
Georg Brandl116aa622007-08-15 14:28:22 +00001:mod:`itertools` --- Functions creating iterators for efficient looping
2=======================================================================
3
4.. module:: itertools
Raymond Hettinger30c73622010-12-01 23:45:20 +00005 :synopsis: Functions creating iterators for efficient looping.
Georg Brandl116aa622007-08-15 14:28:22 +00006.. moduleauthor:: Raymond Hettinger <python@rcn.com>
7.. sectionauthor:: Raymond Hettinger <python@rcn.com>
8
9
Christian Heimesfe337bf2008-03-23 21:54:12 +000010.. testsetup::
11
Raymond Hettinger30c73622010-12-01 23:45:20 +000012 from itertools import *
Christian Heimesfe337bf2008-03-23 21:54:12 +000013
14
Raymond Hettingerf76b9202009-02-17 20:00:59 +000015This module implements a number of :term:`iterator` building blocks inspired
16by constructs from APL, Haskell, and SML. Each has been recast in a form
17suitable for Python.
Georg Brandl116aa622007-08-15 14:28:22 +000018
19The module standardizes a core set of fast, memory efficient tools that are
Raymond Hettingerf76b9202009-02-17 20:00:59 +000020useful by themselves or in combination. Together, they form an "iterator
21algebra" making it possible to construct specialized tools succinctly and
22efficiently in pure Python.
Georg Brandl116aa622007-08-15 14:28:22 +000023
24For instance, SML provides a tabulation tool: ``tabulate(f)`` which produces a
Ezio Melottib6605992010-01-21 20:57:24 +000025sequence ``f(0), f(1), ...``. The same effect can be achieved in Python
Raymond Hettingera6c60372008-03-13 01:26:19 +000026by combining :func:`map` and :func:`count` to form ``map(f, count())``.
Georg Brandl116aa622007-08-15 14:28:22 +000027
Raymond Hettinger2c109ab2009-03-12 00:29:44 +000028These tools and their built-in counterparts also work well with the high-speed
29functions in the :mod:`operator` module. For example, the multiplication
30operator can be mapped across two vectors to form an efficient dot-product:
31``sum(map(operator.mul, vector1, vector2))``.
Georg Brandl116aa622007-08-15 14:28:22 +000032
Georg Brandl116aa622007-08-15 14:28:22 +000033
Raymond Hettingerf76b9202009-02-17 20:00:59 +000034**Infinite Iterators:**
Georg Brandl116aa622007-08-15 14:28:22 +000035
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000036================== ================= ================================================= =========================================
37Iterator Arguments Results Example
38================== ================= ================================================= =========================================
39:func:`count` start, [step] start, start+step, start+2*step, ... ``count(10) --> 10 11 12 13 14 ...``
40:func:`cycle` p p0, p1, ... plast, p0, p1, ... ``cycle('ABCD') --> A B C D A B C D ...``
41:func:`repeat` elem [,n] elem, elem, elem, ... endlessly or up to n times ``repeat(10, 3) --> 10 10 10``
42================== ================= ================================================= =========================================
Georg Brandl116aa622007-08-15 14:28:22 +000043
Raymond Hettingerf76b9202009-02-17 20:00:59 +000044**Iterators terminating on the shortest input sequence:**
45
Benjamin Peterson2989f582014-01-16 10:10:13 -050046============================ ============================ ================================================= =============================================================
47Iterator Arguments Results Example
48============================ ============================ ================================================= =============================================================
49:func:`accumulate` p [,func] p0, p0+p1, p0+p1+p2, ... ``accumulate([1,2,3,4,5]) --> 1 3 6 10 15``
50:func:`chain` p, q, ... p0, p1, ... plast, q0, q1, ... ``chain('ABC', 'DEF') --> A B C D E F``
51:func:`chain.from_iterable` iterable p0, p1, ... plast, q0, q1, ... ``chain.from_iterable(['ABC', 'DEF']) --> A B C D E F``
52:func:`compress` data, selectors (d[0] if s[0]), (d[1] if s[1]), ... ``compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F``
53:func:`dropwhile` pred, seq seq[n], seq[n+1], starting when pred fails ``dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1``
54:func:`filterfalse` pred, seq elements of seq where pred(elem) is false ``filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8``
55:func:`groupby` iterable[, keyfunc] sub-iterators grouped by value of keyfunc(v)
56:func:`islice` seq, [start,] stop [, step] elements from seq[start:stop:step] ``islice('ABCDEFG', 2, None) --> C D E F G``
57:func:`starmap` func, seq func(\*seq[0]), func(\*seq[1]), ... ``starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000``
58:func:`takewhile` pred, seq seq[0], seq[1], until pred fails ``takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4``
59:func:`tee` it, n it1, it2, ... itn splits one iterator into n
60:func:`zip_longest` p, q, ... (p[0], q[0]), (p[1], q[1]), ... ``zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-``
61============================ ============================ ================================================= =============================================================
Raymond Hettingerf76b9202009-02-17 20:00:59 +000062
63**Combinatoric generators:**
64
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000065============================================== ==================== =============================================================
66Iterator Arguments Results
67============================================== ==================== =============================================================
68:func:`product` p, q, ... [repeat=1] cartesian product, equivalent to a nested for-loop
69:func:`permutations` p[, r] r-length tuples, all possible orderings, no repeated elements
Raymond Hettinger36c3c022009-11-19 01:20:07 +000070:func:`combinations` p, r r-length tuples, in sorted order, no repeated elements
71:func:`combinations_with_replacement` p, r r-length tuples, in sorted order, with repeated elements
Raymond Hettinger7f587cd2009-04-10 19:43:50 +000072``product('ABCD', repeat=2)`` ``AA AB AC AD BA BB BC BD CA CB CC CD DA DB DC DD``
73``permutations('ABCD', 2)`` ``AB AC AD BA BC BD CA CB CD DA DB DC``
74``combinations('ABCD', 2)`` ``AB AC AD BC BD CD``
75``combinations_with_replacement('ABCD', 2)`` ``AA AB AC AD BB BC BD CC CD DD``
76============================================== ==================== =============================================================
Georg Brandl116aa622007-08-15 14:28:22 +000077
78
79.. _itertools-functions:
80
81Itertool functions
82------------------
83
84The following module functions all construct and return iterators. Some provide
85streams of infinite length, so they should only be accessed by functions or
86loops that truncate the stream.
87
Raymond Hettinger5d446132011-03-27 18:52:10 -070088.. function:: accumulate(iterable[, func])
Raymond Hettingeradb81462010-12-01 22:50:36 +000089
Raymond Hettinger2d93e6e2010-12-03 02:33:53 +000090 Make an iterator that returns accumulated sums. Elements may be any addable
Serhiy Storchakabfdcd432013-10-13 23:09:14 +030091 type including :class:`~decimal.Decimal` or :class:`~fractions.Fraction`.
92 If the optional *func* argument is supplied, it should be a function of two
93 arguments and it will be used instead of addition.
Raymond Hettingeradb81462010-12-01 22:50:36 +000094
Raymond Hettinger5d446132011-03-27 18:52:10 -070095 Equivalent to::
96
97 def accumulate(iterable, func=operator.add):
Raymond Hettingeradb81462010-12-01 22:50:36 +000098 'Return running totals'
Raymond Hettingerd8ff4652010-12-03 02:09:34 +000099 # accumulate([1,2,3,4,5]) --> 1 3 6 10 15
Raymond Hettinger5d446132011-03-27 18:52:10 -0700100 # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000101 it = iter(iterable)
102 total = next(it)
103 yield total
104 for element in it:
Raymond Hettinger5d446132011-03-27 18:52:10 -0700105 total = func(total, element)
Raymond Hettingerd8ff4652010-12-03 02:09:34 +0000106 yield total
Raymond Hettingeradb81462010-12-01 22:50:36 +0000107
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700108 There are a number of uses for the *func* argument. It can be set to
109 :func:`min` for a running minimum, :func:`max` for a running maximum, or
110 :func:`operator.mul` for a running product. Amortization tables can be
111 built by accumulating interest and applying payments. First-order
112 `recurrence relations <http://en.wikipedia.org/wiki/Recurrence_relation>`_
113 can be modeled by supplying the initial value in the iterable and using only
114 the accumulated total in *func* argument::
Raymond Hettinger5d446132011-03-27 18:52:10 -0700115
116 >>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8]
117 >>> list(accumulate(data, operator.mul)) # running product
118 [3, 12, 72, 144, 144, 1296, 0, 0, 0, 0]
119 >>> list(accumulate(data, max)) # running maximum
120 [3, 4, 6, 6, 6, 9, 9, 9, 9, 9]
121
122 # Amortize a 5% loan of 1000 with 4 annual payments of 90
123 >>> cashflows = [1000, -90, -90, -90, -90]
124 >>> list(accumulate(cashflows, lambda bal, pmt: bal*1.05 + pmt))
125 [1000, 960.0, 918.0, 873.9000000000001, 827.5950000000001]
126
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700127 # Chaotic recurrence relation http://en.wikipedia.org/wiki/Logistic_map
128 >>> logistic_map = lambda x, _: r * x * (1 - x)
129 >>> r = 3.8
130 >>> x0 = 0.4
131 >>> inputs = repeat(x0, 36) # only the initial value is used
132 >>> [format(x, '.2f') for x in accumulate(inputs, logistic_map)]
133 ['0.40', '0.91', '0.30', '0.81', '0.60', '0.92', '0.29', '0.79', '0.63',
Serhiy Storchakaa4d170d2013-12-23 18:20:51 +0200134 '0.88', '0.39', '0.90', '0.33', '0.84', '0.52', '0.95', '0.18', '0.57',
Raymond Hettinger295c1d42011-04-21 11:09:28 -0700135 '0.93', '0.25', '0.71', '0.79', '0.63', '0.88', '0.39', '0.91', '0.32',
136 '0.83', '0.54', '0.95', '0.20', '0.60', '0.91', '0.30', '0.80', '0.60']
137
Raymond Hettinger64801682013-10-12 16:04:17 -0700138 See :func:`functools.reduce` for a similar function that returns only the
139 final accumulated value.
140
Raymond Hettingeradb81462010-12-01 22:50:36 +0000141 .. versionadded:: 3.2
Georg Brandl116aa622007-08-15 14:28:22 +0000142
Raymond Hettinger5d446132011-03-27 18:52:10 -0700143 .. versionchanged:: 3.3
144 Added the optional *func* parameter.
145
Georg Brandl116aa622007-08-15 14:28:22 +0000146.. function:: chain(*iterables)
147
Raymond Hettinger30c73622010-12-01 23:45:20 +0000148 Make an iterator that returns elements from the first iterable until it is
149 exhausted, then proceeds to the next iterable, until all of the iterables are
150 exhausted. Used for treating consecutive sequences as a single sequence.
151 Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000152
Raymond Hettinger30c73622010-12-01 23:45:20 +0000153 def chain(*iterables):
154 # chain('ABC', 'DEF') --> A B C D E F
155 for it in iterables:
156 for element in it:
157 yield element
Georg Brandl116aa622007-08-15 14:28:22 +0000158
159
Georg Brandl933b9742010-07-29 14:36:11 +0000160.. classmethod:: chain.from_iterable(iterable)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000161
Raymond Hettinger30c73622010-12-01 23:45:20 +0000162 Alternate constructor for :func:`chain`. Gets chained inputs from a
Raymond Hettinger1e21ebc2013-09-09 01:54:27 -0500163 single iterable argument that is evaluated lazily. Roughly equivalent to::
Christian Heimes70e7ea22008-02-28 20:02:27 +0000164
Raymond Hettinger30c73622010-12-01 23:45:20 +0000165 def from_iterable(iterables):
166 # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
167 for it in iterables:
168 for element in it:
169 yield element
Christian Heimes70e7ea22008-02-28 20:02:27 +0000170
Christian Heimes78644762008-03-04 23:39:23 +0000171
Christian Heimes836baa52008-02-26 08:18:30 +0000172.. function:: combinations(iterable, r)
173
Raymond Hettinger30c73622010-12-01 23:45:20 +0000174 Return *r* length subsequences of elements from the input *iterable*.
Christian Heimes836baa52008-02-26 08:18:30 +0000175
Raymond Hettinger30c73622010-12-01 23:45:20 +0000176 Combinations are emitted in lexicographic sort order. So, if the
177 input *iterable* is sorted, the combination tuples will be produced
178 in sorted order.
Christian Heimes836baa52008-02-26 08:18:30 +0000179
Raymond Hettinger30c73622010-12-01 23:45:20 +0000180 Elements are treated as unique based on their position, not on their
181 value. So if the input elements are unique, there will be no repeat
182 values in each combination.
Christian Heimes836baa52008-02-26 08:18:30 +0000183
Raymond Hettinger30c73622010-12-01 23:45:20 +0000184 Equivalent to::
Christian Heimes836baa52008-02-26 08:18:30 +0000185
186 def combinations(iterable, r):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000187 # combinations('ABCD', 2) --> AB AC AD BC BD CD
188 # combinations(range(4), 3) --> 012 013 023 123
Christian Heimes836baa52008-02-26 08:18:30 +0000189 pool = tuple(iterable)
Christian Heimes380f7f22008-02-28 11:19:05 +0000190 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000191 if r > n:
192 return
193 indices = list(range(r))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000194 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000195 while True:
Christian Heimes380f7f22008-02-28 11:19:05 +0000196 for i in reversed(range(r)):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000197 if indices[i] != i + n - r:
Christian Heimes836baa52008-02-26 08:18:30 +0000198 break
Christian Heimes380f7f22008-02-28 11:19:05 +0000199 else:
200 return
Christian Heimesb558a2e2008-03-02 22:46:37 +0000201 indices[i] += 1
Christian Heimes380f7f22008-02-28 11:19:05 +0000202 for j in range(i+1, r):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000203 indices[j] = indices[j-1] + 1
204 yield tuple(pool[i] for i in indices)
Christian Heimes836baa52008-02-26 08:18:30 +0000205
Raymond Hettinger30c73622010-12-01 23:45:20 +0000206 The code for :func:`combinations` can be also expressed as a subsequence
207 of :func:`permutations` after filtering entries where the elements are not
208 in sorted order (according to their position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000209
210 def combinations(iterable, r):
211 pool = tuple(iterable)
212 n = len(pool)
213 for indices in permutations(range(n), r):
214 if sorted(indices) == list(indices):
215 yield tuple(pool[i] for i in indices)
216
Raymond Hettinger30c73622010-12-01 23:45:20 +0000217 The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
218 or zero when ``r > n``.
Christian Heimes836baa52008-02-26 08:18:30 +0000219
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000220.. function:: combinations_with_replacement(iterable, r)
221
Raymond Hettinger30c73622010-12-01 23:45:20 +0000222 Return *r* length subsequences of elements from the input *iterable*
223 allowing individual elements to be repeated more than once.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000224
Raymond Hettinger30c73622010-12-01 23:45:20 +0000225 Combinations are emitted in lexicographic sort order. So, if the
226 input *iterable* is sorted, the combination tuples will be produced
227 in sorted order.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000228
Raymond Hettinger30c73622010-12-01 23:45:20 +0000229 Elements are treated as unique based on their position, not on their
230 value. So if the input elements are unique, the generated combinations
231 will also be unique.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000232
Raymond Hettinger30c73622010-12-01 23:45:20 +0000233 Equivalent to::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000234
235 def combinations_with_replacement(iterable, r):
236 # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
237 pool = tuple(iterable)
238 n = len(pool)
239 if not n and r:
240 return
241 indices = [0] * r
242 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000243 while True:
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000244 for i in reversed(range(r)):
245 if indices[i] != n - 1:
246 break
247 else:
248 return
249 indices[i:] = [indices[i] + 1] * (r - i)
250 yield tuple(pool[i] for i in indices)
251
Raymond Hettinger30c73622010-12-01 23:45:20 +0000252 The code for :func:`combinations_with_replacement` can be also expressed as
253 a subsequence of :func:`product` after filtering entries where the elements
254 are not in sorted order (according to their position in the input pool)::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000255
256 def combinations_with_replacement(iterable, r):
257 pool = tuple(iterable)
258 n = len(pool)
259 for indices in product(range(n), repeat=r):
260 if sorted(indices) == list(indices):
261 yield tuple(pool[i] for i in indices)
262
Raymond Hettinger30c73622010-12-01 23:45:20 +0000263 The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000264
Raymond Hettinger30c73622010-12-01 23:45:20 +0000265 .. versionadded:: 3.1
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000266
Georg Brandl67b21b72010-08-17 15:07:14 +0000267
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000268.. function:: compress(data, selectors)
269
Raymond Hettinger30c73622010-12-01 23:45:20 +0000270 Make an iterator that filters elements from *data* returning only those that
271 have a corresponding element in *selectors* that evaluates to ``True``.
272 Stops when either the *data* or *selectors* iterables has been exhausted.
273 Equivalent to::
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000274
Raymond Hettinger30c73622010-12-01 23:45:20 +0000275 def compress(data, selectors):
276 # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
277 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000278
Raymond Hettinger30c73622010-12-01 23:45:20 +0000279 .. versionadded:: 3.1
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000280
281
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000282.. function:: count(start=0, step=1)
Georg Brandl116aa622007-08-15 14:28:22 +0000283
Andrew Kuchlingedb42602013-06-21 08:00:58 -0400284 Make an iterator that returns evenly spaced values starting with number *start*. Often
Raymond Hettinger30c73622010-12-01 23:45:20 +0000285 used as an argument to :func:`map` to generate consecutive data points.
286 Also, used with :func:`zip` to add sequence numbers. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000287
Raymond Hettinger30c73622010-12-01 23:45:20 +0000288 def count(start=0, step=1):
289 # count(10) --> 10 11 12 13 14 ...
Georg Brandl37a80dc2011-01-13 07:31:18 +0000290 # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
Raymond Hettinger30c73622010-12-01 23:45:20 +0000291 n = start
292 while True:
293 yield n
294 n += step
Georg Brandl116aa622007-08-15 14:28:22 +0000295
Raymond Hettinger30c73622010-12-01 23:45:20 +0000296 When counting with floating point numbers, better accuracy can sometimes be
297 achieved by substituting multiplicative code such as: ``(start + step * i
298 for i in count())``.
Raymond Hettinger5bc472a2009-06-17 01:40:52 +0000299
Raymond Hettinger30c73622010-12-01 23:45:20 +0000300 .. versionchanged:: 3.1
301 Added *step* argument and allowed non-integer arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000302
303.. function:: cycle(iterable)
304
Raymond Hettinger30c73622010-12-01 23:45:20 +0000305 Make an iterator returning elements from the iterable and saving a copy of each.
306 When the iterable is exhausted, return elements from the saved copy. Repeats
307 indefinitely. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000308
Raymond Hettinger30c73622010-12-01 23:45:20 +0000309 def cycle(iterable):
310 # cycle('ABCD') --> A B C D A B C D A B C D ...
311 saved = []
312 for element in iterable:
313 yield element
314 saved.append(element)
315 while saved:
316 for element in saved:
Georg Brandl116aa622007-08-15 14:28:22 +0000317 yield element
318
Raymond Hettinger30c73622010-12-01 23:45:20 +0000319 Note, this member of the toolkit may require significant auxiliary storage
320 (depending on the length of the iterable).
Georg Brandl116aa622007-08-15 14:28:22 +0000321
322
323.. function:: dropwhile(predicate, iterable)
324
Raymond Hettinger30c73622010-12-01 23:45:20 +0000325 Make an iterator that drops elements from the iterable as long as the predicate
326 is true; afterwards, returns every element. Note, the iterator does not produce
327 *any* output until the predicate first becomes false, so it may have a lengthy
328 start-up time. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000329
Raymond Hettinger30c73622010-12-01 23:45:20 +0000330 def dropwhile(predicate, iterable):
331 # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
332 iterable = iter(iterable)
333 for x in iterable:
334 if not predicate(x):
335 yield x
336 break
337 for x in iterable:
338 yield x
Georg Brandl116aa622007-08-15 14:28:22 +0000339
Raymond Hettinger749761e2009-01-27 04:42:48 +0000340.. function:: filterfalse(predicate, iterable)
341
Raymond Hettinger30c73622010-12-01 23:45:20 +0000342 Make an iterator that filters elements from iterable returning only those for
343 which the predicate is ``False``. If *predicate* is ``None``, return the items
344 that are false. Equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000345
Raymond Hettinger30c73622010-12-01 23:45:20 +0000346 def filterfalse(predicate, iterable):
347 # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
348 if predicate is None:
349 predicate = bool
350 for x in iterable:
351 if not predicate(x):
352 yield x
Raymond Hettinger749761e2009-01-27 04:42:48 +0000353
Georg Brandl116aa622007-08-15 14:28:22 +0000354
Georg Brandl3dd33882009-06-01 17:35:27 +0000355.. function:: groupby(iterable, key=None)
Georg Brandl116aa622007-08-15 14:28:22 +0000356
Raymond Hettinger30c73622010-12-01 23:45:20 +0000357 Make an iterator that returns consecutive keys and groups from the *iterable*.
358 The *key* is a function computing a key value for each element. If not
359 specified or is ``None``, *key* defaults to an identity function and returns
360 the element unchanged. Generally, the iterable needs to already be sorted on
361 the same key function.
Georg Brandl116aa622007-08-15 14:28:22 +0000362
Raymond Hettinger30c73622010-12-01 23:45:20 +0000363 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
364 generates a break or new group every time the value of the key function changes
365 (which is why it is usually necessary to have sorted the data using the same key
366 function). That behavior differs from SQL's GROUP BY which aggregates common
367 elements regardless of their input order.
Georg Brandl116aa622007-08-15 14:28:22 +0000368
Raymond Hettinger30c73622010-12-01 23:45:20 +0000369 The returned group is itself an iterator that shares the underlying iterable
370 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
371 object is advanced, the previous group is no longer visible. So, if that data
372 is needed later, it should be stored as a list::
Georg Brandl116aa622007-08-15 14:28:22 +0000373
Raymond Hettinger30c73622010-12-01 23:45:20 +0000374 groups = []
375 uniquekeys = []
376 data = sorted(data, key=keyfunc)
377 for k, g in groupby(data, keyfunc):
378 groups.append(list(g)) # Store group iterator as a list
379 uniquekeys.append(k)
Georg Brandl116aa622007-08-15 14:28:22 +0000380
Raymond Hettinger30c73622010-12-01 23:45:20 +0000381 :func:`groupby` is equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000382
Raymond Hettinger30c73622010-12-01 23:45:20 +0000383 class groupby:
384 # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
385 # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
386 def __init__(self, iterable, key=None):
387 if key is None:
388 key = lambda x: x
389 self.keyfunc = key
390 self.it = iter(iterable)
391 self.tgtkey = self.currkey = self.currvalue = object()
392 def __iter__(self):
393 return self
394 def __next__(self):
395 while self.currkey == self.tgtkey:
396 self.currvalue = next(self.it) # Exit on StopIteration
397 self.currkey = self.keyfunc(self.currvalue)
398 self.tgtkey = self.currkey
399 return (self.currkey, self._grouper(self.tgtkey))
400 def _grouper(self, tgtkey):
401 while self.currkey == tgtkey:
402 yield self.currvalue
403 self.currvalue = next(self.it) # Exit on StopIteration
404 self.currkey = self.keyfunc(self.currvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000405
Georg Brandl116aa622007-08-15 14:28:22 +0000406
Ezio Melottie0add762012-09-14 06:32:35 +0300407.. function:: islice(iterable, stop)
408 islice(iterable, start, stop[, step])
Georg Brandl116aa622007-08-15 14:28:22 +0000409
Raymond Hettinger30c73622010-12-01 23:45:20 +0000410 Make an iterator that returns selected elements from the iterable. If *start* is
411 non-zero, then elements from the iterable are skipped until start is reached.
412 Afterward, elements are returned consecutively unless *step* is set higher than
413 one which results in items being skipped. If *stop* is ``None``, then iteration
414 continues until the iterator is exhausted, if at all; otherwise, it stops at the
415 specified position. Unlike regular slicing, :func:`islice` does not support
416 negative values for *start*, *stop*, or *step*. Can be used to extract related
417 fields from data where the internal structure has been flattened (for example, a
418 multi-line report may list a name field on every third line). Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000419
Raymond Hettinger30c73622010-12-01 23:45:20 +0000420 def islice(iterable, *args):
421 # islice('ABCDEFG', 2) --> A B
422 # islice('ABCDEFG', 2, 4) --> C D
423 # islice('ABCDEFG', 2, None) --> C D E F G
424 # islice('ABCDEFG', 0, None, 2) --> A C E G
425 s = slice(*args)
426 it = iter(range(s.start or 0, s.stop or sys.maxsize, s.step or 1))
427 nexti = next(it)
428 for i, element in enumerate(iterable):
429 if i == nexti:
430 yield element
431 nexti = next(it)
Georg Brandl116aa622007-08-15 14:28:22 +0000432
Raymond Hettinger30c73622010-12-01 23:45:20 +0000433 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
434 then the step defaults to one.
Georg Brandl116aa622007-08-15 14:28:22 +0000435
Georg Brandl116aa622007-08-15 14:28:22 +0000436
Georg Brandl3dd33882009-06-01 17:35:27 +0000437.. function:: permutations(iterable, r=None)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000438
Raymond Hettinger30c73622010-12-01 23:45:20 +0000439 Return successive *r* length permutations of elements in the *iterable*.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000440
Raymond Hettinger30c73622010-12-01 23:45:20 +0000441 If *r* is not specified or is ``None``, then *r* defaults to the length
442 of the *iterable* and all possible full-length permutations
443 are generated.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000444
Raymond Hettinger30c73622010-12-01 23:45:20 +0000445 Permutations are emitted in lexicographic sort order. So, if the
446 input *iterable* is sorted, the permutation tuples will be produced
447 in sorted order.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000448
Raymond Hettinger30c73622010-12-01 23:45:20 +0000449 Elements are treated as unique based on their position, not on their
450 value. So if the input elements are unique, there will be no repeat
451 values in each permutation.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000452
Raymond Hettinger30c73622010-12-01 23:45:20 +0000453 Equivalent to::
Christian Heimesb558a2e2008-03-02 22:46:37 +0000454
455 def permutations(iterable, r=None):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000456 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
457 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesb558a2e2008-03-02 22:46:37 +0000458 pool = tuple(iterable)
459 n = len(pool)
460 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000461 if r > n:
462 return
463 indices = list(range(n))
Sandro Tosi73866622011-12-25 17:25:45 +0100464 cycles = list(range(n, n-r, -1))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000465 yield tuple(pool[i] for i in indices[:r])
466 while n:
467 for i in reversed(range(r)):
468 cycles[i] -= 1
469 if cycles[i] == 0:
470 indices[i:] = indices[i+1:] + indices[i:i+1]
471 cycles[i] = n - i
472 else:
473 j = cycles[i]
474 indices[i], indices[-j] = indices[-j], indices[i]
475 yield tuple(pool[i] for i in indices[:r])
476 break
477 else:
478 return
Christian Heimes70e7ea22008-02-28 20:02:27 +0000479
Raymond Hettinger30c73622010-12-01 23:45:20 +0000480 The code for :func:`permutations` can be also expressed as a subsequence of
481 :func:`product`, filtered to exclude entries with repeated elements (those
482 from the same position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000483
484 def permutations(iterable, r=None):
485 pool = tuple(iterable)
486 n = len(pool)
487 r = n if r is None else r
488 for indices in product(range(n), repeat=r):
489 if len(set(indices)) == r:
490 yield tuple(pool[i] for i in indices)
491
Raymond Hettinger30c73622010-12-01 23:45:20 +0000492 The number of items returned is ``n! / (n-r)!`` when ``0 <= r <= n``
493 or zero when ``r > n``.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000494
Georg Brandl3dd33882009-06-01 17:35:27 +0000495.. function:: product(*iterables, repeat=1)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000496
Raymond Hettinger30c73622010-12-01 23:45:20 +0000497 Cartesian product of input iterables.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000498
Raymond Hettinger30c73622010-12-01 23:45:20 +0000499 Equivalent to nested for-loops in a generator expression. For example,
500 ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000501
Raymond Hettinger30c73622010-12-01 23:45:20 +0000502 The nested loops cycle like an odometer with the rightmost element advancing
503 on every iteration. This pattern creates a lexicographic ordering so that if
504 the input's iterables are sorted, the product tuples are emitted in sorted
505 order.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000506
Raymond Hettinger30c73622010-12-01 23:45:20 +0000507 To compute the product of an iterable with itself, specify the number of
508 repetitions with the optional *repeat* keyword argument. For example,
509 ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``.
Christian Heimes9e7f1d22008-02-28 12:27:11 +0000510
Raymond Hettinger30c73622010-12-01 23:45:20 +0000511 This function is equivalent to the following code, except that the
512 actual implementation does not build up intermediate results in memory::
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000513
Raymond Hettinger30c73622010-12-01 23:45:20 +0000514 def product(*args, repeat=1):
515 # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
516 # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
517 pools = [tuple(pool) for pool in args] * repeat
518 result = [[]]
519 for pool in pools:
520 result = [x+[y] for x in result for y in pool]
521 for prod in result:
522 yield tuple(prod)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000523
524
Raymond Hettingerd75ad442009-06-01 19:16:52 +0000525.. function:: repeat(object[, times])
Georg Brandl116aa622007-08-15 14:28:22 +0000526
Raymond Hettinger30c73622010-12-01 23:45:20 +0000527 Make an iterator that returns *object* over and over again. Runs indefinitely
528 unless the *times* argument is specified. Used as argument to :func:`map` for
529 invariant parameters to the called function. Also used with :func:`zip` to
530 create an invariant part of a tuple record. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000531
Raymond Hettinger30c73622010-12-01 23:45:20 +0000532 def repeat(object, times=None):
533 # repeat(10, 3) --> 10 10 10
534 if times is None:
535 while True:
536 yield object
537 else:
538 for i in range(times):
539 yield object
Georg Brandl116aa622007-08-15 14:28:22 +0000540
Raymond Hettingerfc3ba6b2012-02-01 09:07:40 -0800541 A common use for *repeat* is to supply a stream of constant values to *map*
542 or *zip*::
543
544 >>> list(map(pow, range(10), repeat(2)))
545 [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
Georg Brandl116aa622007-08-15 14:28:22 +0000546
547.. function:: starmap(function, iterable)
548
Raymond Hettinger30c73622010-12-01 23:45:20 +0000549 Make an iterator that computes the function using arguments obtained from
550 the iterable. Used instead of :func:`map` when argument parameters are already
551 grouped in tuples from a single iterable (the data has been "pre-zipped"). The
552 difference between :func:`map` and :func:`starmap` parallels the distinction
553 between ``function(a,b)`` and ``function(*c)``. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000554
Raymond Hettinger30c73622010-12-01 23:45:20 +0000555 def starmap(function, iterable):
556 # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
557 for args in iterable:
558 yield function(*args)
Christian Heimes679db4a2008-01-18 09:56:22 +0000559
Georg Brandl116aa622007-08-15 14:28:22 +0000560
561.. function:: takewhile(predicate, iterable)
562
Raymond Hettinger30c73622010-12-01 23:45:20 +0000563 Make an iterator that returns elements from the iterable as long as the
564 predicate is true. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000565
Raymond Hettinger30c73622010-12-01 23:45:20 +0000566 def takewhile(predicate, iterable):
567 # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
568 for x in iterable:
569 if predicate(x):
570 yield x
571 else:
572 break
Georg Brandl116aa622007-08-15 14:28:22 +0000573
574
Georg Brandl3dd33882009-06-01 17:35:27 +0000575.. function:: tee(iterable, n=2)
Georg Brandl116aa622007-08-15 14:28:22 +0000576
Raymond Hettinger30c73622010-12-01 23:45:20 +0000577 Return *n* independent iterators from a single iterable. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000578
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000579 def tee(iterable, n=2):
580 it = iter(iterable)
581 deques = [collections.deque() for i in range(n)]
582 def gen(mydeque):
583 while True:
584 if not mydeque: # when the local deque is empty
585 newval = next(it) # fetch a new value and
586 for d in deques: # load it to all the deques
587 d.append(newval)
588 yield mydeque.popleft()
589 return tuple(gen(d) for d in deques)
Georg Brandl116aa622007-08-15 14:28:22 +0000590
Raymond Hettinger30c73622010-12-01 23:45:20 +0000591 Once :func:`tee` has made a split, the original *iterable* should not be
592 used anywhere else; otherwise, the *iterable* could get advanced without
593 the tee objects being informed.
Georg Brandl116aa622007-08-15 14:28:22 +0000594
Raymond Hettinger30c73622010-12-01 23:45:20 +0000595 This itertool may require significant auxiliary storage (depending on how
596 much temporary data needs to be stored). In general, if one iterator uses
597 most or all of the data before another iterator starts, it is faster to use
598 :func:`list` instead of :func:`tee`.
Georg Brandl116aa622007-08-15 14:28:22 +0000599
Georg Brandl116aa622007-08-15 14:28:22 +0000600
Georg Brandl3dd33882009-06-01 17:35:27 +0000601.. function:: zip_longest(*iterables, fillvalue=None)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000602
Raymond Hettinger30c73622010-12-01 23:45:20 +0000603 Make an iterator that aggregates elements from each of the iterables. If the
604 iterables are of uneven length, missing values are filled-in with *fillvalue*.
605 Iteration continues until the longest iterable is exhausted. Equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000606
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700607 class ZipExhausted(Exception):
608 pass
609
610 def zip_longest(*args, **kwds):
Raymond Hettinger30c73622010-12-01 23:45:20 +0000611 # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700612 fillvalue = kwds.get('fillvalue')
613 counter = len(args) - 1
614 def sentinel():
615 nonlocal counter
616 if not counter:
617 raise ZipExhausted
618 counter -= 1
619 yield fillvalue
Raymond Hettinger30c73622010-12-01 23:45:20 +0000620 fillers = repeat(fillvalue)
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700621 iterators = [chain(it, sentinel(), fillers) for it in args]
Raymond Hettinger30c73622010-12-01 23:45:20 +0000622 try:
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700623 while iterators:
624 yield tuple(map(next, iterators))
625 except ZipExhausted:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000626 pass
Raymond Hettinger749761e2009-01-27 04:42:48 +0000627
Raymond Hettinger30c73622010-12-01 23:45:20 +0000628 If one of the iterables is potentially infinite, then the :func:`zip_longest`
629 function should be wrapped with something that limits the number of calls
630 (for example :func:`islice` or :func:`takewhile`). If not specified,
631 *fillvalue* defaults to ``None``.
Raymond Hettinger749761e2009-01-27 04:42:48 +0000632
633
Georg Brandl116aa622007-08-15 14:28:22 +0000634.. _itertools-recipes:
635
Raymond Hettinger1fa76822010-12-06 23:31:36 +0000636Itertools Recipes
637-----------------
Georg Brandl116aa622007-08-15 14:28:22 +0000638
639This section shows recipes for creating an extended toolset using the existing
640itertools as building blocks.
641
642The extended tools offer the same high performance as the underlying toolset.
643The superior memory performance is kept by processing elements one at a time
644rather than bringing the whole iterable into memory all at once. Code volume is
645kept small by linking the tools together in a functional style which helps
646eliminate temporary variables. High speed is retained by preferring
Georg Brandl9afde1c2007-11-01 20:32:30 +0000647"vectorized" building blocks over the use of for-loops and :term:`generator`\s
Christian Heimesfe337bf2008-03-23 21:54:12 +0000648which incur interpreter overhead.
649
650.. testcode::
Georg Brandl116aa622007-08-15 14:28:22 +0000651
Raymond Hettinger30c73622010-12-01 23:45:20 +0000652 def take(n, iterable):
653 "Return first n items of the iterable as a list"
654 return list(islice(iterable, n))
Georg Brandl116aa622007-08-15 14:28:22 +0000655
Raymond Hettinger30c73622010-12-01 23:45:20 +0000656 def tabulate(function, start=0):
657 "Return function(0), function(1), ..."
658 return map(function, count(start))
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000659
Raymond Hettinger30c73622010-12-01 23:45:20 +0000660 def consume(iterator, n):
661 "Advance the iterator n-steps ahead. If n is none, consume entirely."
662 # Use functions that consume iterators at C speed.
663 if n is None:
664 # feed the entire iterator into a zero-length deque
665 collections.deque(iterator, maxlen=0)
666 else:
667 # advance to the empty slice starting at position n
668 next(islice(iterator, n, n), None)
Raymond Hettingerfa007962009-03-09 11:55:25 +0000669
Raymond Hettinger30c73622010-12-01 23:45:20 +0000670 def nth(iterable, n, default=None):
671 "Returns the nth item or a default value"
672 return next(islice(iterable, n, None), default)
Georg Brandl116aa622007-08-15 14:28:22 +0000673
Raymond Hettinger30c73622010-12-01 23:45:20 +0000674 def quantify(iterable, pred=bool):
675 "Count how many times the predicate is true"
676 return sum(map(pred, iterable))
Georg Brandl116aa622007-08-15 14:28:22 +0000677
Raymond Hettinger30c73622010-12-01 23:45:20 +0000678 def padnone(iterable):
679 """Returns the sequence elements and then returns None indefinitely.
Georg Brandl116aa622007-08-15 14:28:22 +0000680
Raymond Hettinger30c73622010-12-01 23:45:20 +0000681 Useful for emulating the behavior of the built-in map() function.
682 """
683 return chain(iterable, repeat(None))
Georg Brandl116aa622007-08-15 14:28:22 +0000684
Raymond Hettinger30c73622010-12-01 23:45:20 +0000685 def ncycles(iterable, n):
686 "Returns the sequence elements n times"
687 return chain.from_iterable(repeat(tuple(iterable), n))
Georg Brandl116aa622007-08-15 14:28:22 +0000688
Raymond Hettinger30c73622010-12-01 23:45:20 +0000689 def dotproduct(vec1, vec2):
690 return sum(map(operator.mul, vec1, vec2))
Georg Brandl116aa622007-08-15 14:28:22 +0000691
Raymond Hettinger30c73622010-12-01 23:45:20 +0000692 def flatten(listOfLists):
693 "Flatten one level of nesting"
694 return chain.from_iterable(listOfLists)
Georg Brandl116aa622007-08-15 14:28:22 +0000695
Raymond Hettinger30c73622010-12-01 23:45:20 +0000696 def repeatfunc(func, times=None, *args):
697 """Repeat calls to func with specified arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000698
Raymond Hettinger30c73622010-12-01 23:45:20 +0000699 Example: repeatfunc(random.random)
700 """
701 if times is None:
702 return starmap(func, repeat(args))
703 return starmap(func, repeat(args, times))
Georg Brandl116aa622007-08-15 14:28:22 +0000704
Raymond Hettinger30c73622010-12-01 23:45:20 +0000705 def pairwise(iterable):
706 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
707 a, b = tee(iterable)
708 next(b, None)
709 return zip(a, b)
Georg Brandl116aa622007-08-15 14:28:22 +0000710
Raymond Hettinger44571da2013-05-05 19:53:41 -0700711 def grouper(iterable, n, fillvalue=None):
Raymond Hettinger9ae94732012-07-08 16:04:03 -0700712 "Collect data into fixed-length chunks or blocks"
Raymond Hettinger44571da2013-05-05 19:53:41 -0700713 # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
Raymond Hettinger30c73622010-12-01 23:45:20 +0000714 args = [iter(iterable)] * n
715 return zip_longest(*args, fillvalue=fillvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000716
Raymond Hettinger30c73622010-12-01 23:45:20 +0000717 def roundrobin(*iterables):
718 "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
719 # Recipe credited to George Sakkis
720 pending = len(iterables)
721 nexts = cycle(iter(it).__next__ for it in iterables)
722 while pending:
723 try:
724 for next in nexts:
725 yield next()
726 except StopIteration:
727 pending -= 1
728 nexts = cycle(islice(nexts, pending))
Georg Brandl116aa622007-08-15 14:28:22 +0000729
Raymond Hettinger30c73622010-12-01 23:45:20 +0000730 def partition(pred, iterable):
731 'Use a predicate to partition entries into false entries and true entries'
732 # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
733 t1, t2 = tee(iterable)
734 return filterfalse(pred, t1), filter(pred, t2)
Raymond Hettinger5ce0aa22010-12-01 10:49:19 +0000735
Raymond Hettinger30c73622010-12-01 23:45:20 +0000736 def powerset(iterable):
737 "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
738 s = list(iterable)
739 return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettinger08d01ee2010-08-07 05:36:53 +0000740
Raymond Hettinger30c73622010-12-01 23:45:20 +0000741 def unique_everseen(iterable, key=None):
742 "List unique elements, preserving order. Remember all elements ever seen."
743 # unique_everseen('AAAABBBCCDAABBB') --> A B C D
744 # unique_everseen('ABBCcAD', str.lower) --> A B C D
745 seen = set()
746 seen_add = seen.add
747 if key is None:
748 for element in filterfalse(seen.__contains__, iterable):
749 seen_add(element)
750 yield element
751 else:
752 for element in iterable:
753 k = key(element)
754 if k not in seen:
755 seen_add(k)
756 yield element
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000757
Raymond Hettinger30c73622010-12-01 23:45:20 +0000758 def unique_justseen(iterable, key=None):
759 "List unique elements, preserving order. Remember only the element just seen."
760 # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
761 # unique_justseen('ABBCcAD', str.lower) --> A B C A D
762 return map(next, map(itemgetter(1), groupby(iterable, key)))
Raymond Hettingerad9d96b2009-01-02 21:39:07 +0000763
Raymond Hettinger30c73622010-12-01 23:45:20 +0000764 def iter_except(func, exception, first=None):
765 """ Call a function repeatedly until an exception is raised.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000766
Raymond Hettinger30c73622010-12-01 23:45:20 +0000767 Converts a call-until-exception interface to an iterator interface.
Andrew Kuchling1d7d5802013-06-20 21:17:41 -0400768 Like builtins.iter(func, sentinel) but uses an exception instead
Raymond Hettinger30c73622010-12-01 23:45:20 +0000769 of a sentinel to end the loop.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000770
Raymond Hettinger30c73622010-12-01 23:45:20 +0000771 Examples:
772 iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
773 iter_except(d.popitem, KeyError) # non-blocking dict iterator
774 iter_except(d.popleft, IndexError) # non-blocking deque iterator
775 iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
776 iter_except(s.pop, KeyError) # non-blocking set iterator
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000777
Raymond Hettinger30c73622010-12-01 23:45:20 +0000778 """
779 try:
780 if first is not None:
781 yield first() # For database APIs needing an initial cast to db.first()
782 while 1:
783 yield func()
784 except exception:
785 pass
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000786
Raymond Hettinger31b26f62014-04-02 03:16:42 -0700787 def first_true(iterable, default=False, pred=None):
788 """Returns the first true value in the iterable.
789
790 If no true value is found, returns *default*
791
792 If *pred* is not None, returns the first item
793 for which pred(item) is true.
794
795 """
796 # first_true([a,b,c], x) --> a or b or c or x
797 # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
798 return next(filter(pred, iterable), default)
799
Raymond Hettinger30c73622010-12-01 23:45:20 +0000800 def random_product(*args, repeat=1):
801 "Random selection from itertools.product(*args, **kwds)"
802 pools = [tuple(pool) for pool in args] * repeat
803 return tuple(random.choice(pool) for pool in pools)
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000804
Raymond Hettinger30c73622010-12-01 23:45:20 +0000805 def random_permutation(iterable, r=None):
806 "Random selection from itertools.permutations(iterable, r)"
807 pool = tuple(iterable)
808 r = len(pool) if r is None else r
809 return tuple(random.sample(pool, r))
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000810
Raymond Hettinger30c73622010-12-01 23:45:20 +0000811 def random_combination(iterable, r):
812 "Random selection from itertools.combinations(iterable, r)"
813 pool = tuple(iterable)
814 n = len(pool)
815 indices = sorted(random.sample(range(n), r))
816 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000817
Raymond Hettinger30c73622010-12-01 23:45:20 +0000818 def random_combination_with_replacement(iterable, r):
819 "Random selection from itertools.combinations_with_replacement(iterable, r)"
820 pool = tuple(iterable)
821 n = len(pool)
822 indices = sorted(random.randrange(n) for i in range(r))
823 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000824
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000825Note, many of the above recipes can be optimized by replacing global lookups
826with local variables defined as default values. For example, the
827*dotproduct* recipe can be written as::
828
Raymond Hettinger30c73622010-12-01 23:45:20 +0000829 def dotproduct(vec1, vec2, sum=sum, map=map, mul=operator.mul):
830 return sum(map(mul, vec1, vec2))