blob: aa793e26c5040a8372791048e82f38eeb0affe59 [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
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000046==================== ============================ ================================================= =============================================================
47Iterator Arguments Results Example
48==================== ============================ ================================================= =============================================================
Raymond Hettinger5d446132011-03-27 18:52:10 -070049:func:`accumulate` p [,func] p0, p0+p1, p0+p1+p2, ... ``accumulate([1,2,3,4,5]) --> 1 3 6 10 15``
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000050:func:`chain` p, q, ... p0, p1, ... plast, q0, q1, ... ``chain('ABC', 'DEF') --> A B C D E F``
Raymond Hettinger8df58f72013-09-09 01:29:40 -050051chain.from_iterable iterable p0, p1, ... plast, q0, q1, ... ``chain.from_iterable(['ABC', 'DEF']) --> A B C D E F``
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000052: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``
Georg Brandlc2da5ce2009-08-13 09:16:39 +000054:func:`filterfalse` pred, seq elements of seq where pred(elem) is False ``filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8``
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000055: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
Georg Brandlc2da5ce2009-08-13 09:16:39 +000060:func:`zip_longest` p, q, ... (p[0], q[0]), (p[1], q[1]), ... ``zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-``
Raymond Hettinger5bfd8ce2009-04-10 19:02:36 +000061==================== ============================ ================================================= =============================================================
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
Raymond Hettinger5d446132011-03-27 18:52:10 -070091 type including :class:`Decimal` or :class:`Fraction`. If the optional
92 *func* argument is supplied, it should be a function of two arguments
93 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',
134 '0.88' ,'0.39', '0.90', '0.33', '0.84', '0.52', '0.95', '0.18', '0.57',
135 '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 Hettingeradb81462010-12-01 22:50:36 +0000138 .. versionadded:: 3.2
Georg Brandl116aa622007-08-15 14:28:22 +0000139
Raymond Hettinger5d446132011-03-27 18:52:10 -0700140 .. versionchanged:: 3.3
141 Added the optional *func* parameter.
142
Georg Brandl116aa622007-08-15 14:28:22 +0000143.. function:: chain(*iterables)
144
Raymond Hettinger30c73622010-12-01 23:45:20 +0000145 Make an iterator that returns elements from the first iterable until it is
146 exhausted, then proceeds to the next iterable, until all of the iterables are
147 exhausted. Used for treating consecutive sequences as a single sequence.
148 Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000149
Raymond Hettinger30c73622010-12-01 23:45:20 +0000150 def chain(*iterables):
151 # chain('ABC', 'DEF') --> A B C D E F
152 for it in iterables:
153 for element in it:
154 yield element
Georg Brandl116aa622007-08-15 14:28:22 +0000155
156
Georg Brandl933b9742010-07-29 14:36:11 +0000157.. classmethod:: chain.from_iterable(iterable)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000158
Raymond Hettinger30c73622010-12-01 23:45:20 +0000159 Alternate constructor for :func:`chain`. Gets chained inputs from a
160 single iterable argument that is evaluated lazily. Equivalent to::
Christian Heimes70e7ea22008-02-28 20:02:27 +0000161
Raymond Hettinger30c73622010-12-01 23:45:20 +0000162 @classmethod
163 def from_iterable(iterables):
164 # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
165 for it in iterables:
166 for element in it:
167 yield element
Christian Heimes70e7ea22008-02-28 20:02:27 +0000168
Christian Heimes78644762008-03-04 23:39:23 +0000169
Christian Heimes836baa52008-02-26 08:18:30 +0000170.. function:: combinations(iterable, r)
171
Raymond Hettinger30c73622010-12-01 23:45:20 +0000172 Return *r* length subsequences of elements from the input *iterable*.
Christian Heimes836baa52008-02-26 08:18:30 +0000173
Raymond Hettinger30c73622010-12-01 23:45:20 +0000174 Combinations are emitted in lexicographic sort order. So, if the
175 input *iterable* is sorted, the combination tuples will be produced
176 in sorted order.
Christian Heimes836baa52008-02-26 08:18:30 +0000177
Raymond Hettinger30c73622010-12-01 23:45:20 +0000178 Elements are treated as unique based on their position, not on their
179 value. So if the input elements are unique, there will be no repeat
180 values in each combination.
Christian Heimes836baa52008-02-26 08:18:30 +0000181
Raymond Hettinger30c73622010-12-01 23:45:20 +0000182 Equivalent to::
Christian Heimes836baa52008-02-26 08:18:30 +0000183
184 def combinations(iterable, r):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000185 # combinations('ABCD', 2) --> AB AC AD BC BD CD
186 # combinations(range(4), 3) --> 012 013 023 123
Christian Heimes836baa52008-02-26 08:18:30 +0000187 pool = tuple(iterable)
Christian Heimes380f7f22008-02-28 11:19:05 +0000188 n = len(pool)
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000189 if r > n:
190 return
191 indices = list(range(r))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000192 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000193 while True:
Christian Heimes380f7f22008-02-28 11:19:05 +0000194 for i in reversed(range(r)):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000195 if indices[i] != i + n - r:
Christian Heimes836baa52008-02-26 08:18:30 +0000196 break
Christian Heimes380f7f22008-02-28 11:19:05 +0000197 else:
198 return
Christian Heimesb558a2e2008-03-02 22:46:37 +0000199 indices[i] += 1
Christian Heimes380f7f22008-02-28 11:19:05 +0000200 for j in range(i+1, r):
Christian Heimesb558a2e2008-03-02 22:46:37 +0000201 indices[j] = indices[j-1] + 1
202 yield tuple(pool[i] for i in indices)
Christian Heimes836baa52008-02-26 08:18:30 +0000203
Raymond Hettinger30c73622010-12-01 23:45:20 +0000204 The code for :func:`combinations` can be also expressed as a subsequence
205 of :func:`permutations` after filtering entries where the elements are not
206 in sorted order (according to their position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000207
208 def combinations(iterable, r):
209 pool = tuple(iterable)
210 n = len(pool)
211 for indices in permutations(range(n), r):
212 if sorted(indices) == list(indices):
213 yield tuple(pool[i] for i in indices)
214
Raymond Hettinger30c73622010-12-01 23:45:20 +0000215 The number of items returned is ``n! / r! / (n-r)!`` when ``0 <= r <= n``
216 or zero when ``r > n``.
Christian Heimes836baa52008-02-26 08:18:30 +0000217
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000218.. function:: combinations_with_replacement(iterable, r)
219
Raymond Hettinger30c73622010-12-01 23:45:20 +0000220 Return *r* length subsequences of elements from the input *iterable*
221 allowing individual elements to be repeated more than once.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000222
Raymond Hettinger30c73622010-12-01 23:45:20 +0000223 Combinations are emitted in lexicographic sort order. So, if the
224 input *iterable* is sorted, the combination tuples will be produced
225 in sorted order.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000226
Raymond Hettinger30c73622010-12-01 23:45:20 +0000227 Elements are treated as unique based on their position, not on their
228 value. So if the input elements are unique, the generated combinations
229 will also be unique.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000230
Raymond Hettinger30c73622010-12-01 23:45:20 +0000231 Equivalent to::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000232
233 def combinations_with_replacement(iterable, r):
234 # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
235 pool = tuple(iterable)
236 n = len(pool)
237 if not n and r:
238 return
239 indices = [0] * r
240 yield tuple(pool[i] for i in indices)
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000241 while True:
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000242 for i in reversed(range(r)):
243 if indices[i] != n - 1:
244 break
245 else:
246 return
247 indices[i:] = [indices[i] + 1] * (r - i)
248 yield tuple(pool[i] for i in indices)
249
Raymond Hettinger30c73622010-12-01 23:45:20 +0000250 The code for :func:`combinations_with_replacement` can be also expressed as
251 a subsequence of :func:`product` after filtering entries where the elements
252 are not in sorted order (according to their position in the input pool)::
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000253
254 def combinations_with_replacement(iterable, r):
255 pool = tuple(iterable)
256 n = len(pool)
257 for indices in product(range(n), repeat=r):
258 if sorted(indices) == list(indices):
259 yield tuple(pool[i] for i in indices)
260
Raymond Hettinger30c73622010-12-01 23:45:20 +0000261 The number of items returned is ``(n+r-1)! / r! / (n-1)!`` when ``n > 0``.
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000262
Raymond Hettinger30c73622010-12-01 23:45:20 +0000263 .. versionadded:: 3.1
Raymond Hettingerd07d9392009-01-27 04:20:44 +0000264
Georg Brandl67b21b72010-08-17 15:07:14 +0000265
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000266.. function:: compress(data, selectors)
267
Raymond Hettinger30c73622010-12-01 23:45:20 +0000268 Make an iterator that filters elements from *data* returning only those that
269 have a corresponding element in *selectors* that evaluates to ``True``.
270 Stops when either the *data* or *selectors* iterables has been exhausted.
271 Equivalent to::
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000272
Raymond Hettinger30c73622010-12-01 23:45:20 +0000273 def compress(data, selectors):
274 # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
275 return (d for d, s in zip(data, selectors) if s)
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000276
Raymond Hettinger30c73622010-12-01 23:45:20 +0000277 .. versionadded:: 3.1
Raymond Hettinger6b3b0fc2009-01-26 02:56:58 +0000278
279
Raymond Hettinger9e8dbbc2009-02-14 04:21:49 +0000280.. function:: count(start=0, step=1)
Georg Brandl116aa622007-08-15 14:28:22 +0000281
Andrew Kuchlingedb42602013-06-21 08:00:58 -0400282 Make an iterator that returns evenly spaced values starting with number *start*. Often
Raymond Hettinger30c73622010-12-01 23:45:20 +0000283 used as an argument to :func:`map` to generate consecutive data points.
284 Also, used with :func:`zip` to add sequence numbers. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000285
Raymond Hettinger30c73622010-12-01 23:45:20 +0000286 def count(start=0, step=1):
287 # count(10) --> 10 11 12 13 14 ...
Georg Brandl37a80dc2011-01-13 07:31:18 +0000288 # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
Raymond Hettinger30c73622010-12-01 23:45:20 +0000289 n = start
290 while True:
291 yield n
292 n += step
Georg Brandl116aa622007-08-15 14:28:22 +0000293
Raymond Hettinger30c73622010-12-01 23:45:20 +0000294 When counting with floating point numbers, better accuracy can sometimes be
295 achieved by substituting multiplicative code such as: ``(start + step * i
296 for i in count())``.
Raymond Hettinger5bc472a2009-06-17 01:40:52 +0000297
Raymond Hettinger30c73622010-12-01 23:45:20 +0000298 .. versionchanged:: 3.1
299 Added *step* argument and allowed non-integer arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000300
301.. function:: cycle(iterable)
302
Raymond Hettinger30c73622010-12-01 23:45:20 +0000303 Make an iterator returning elements from the iterable and saving a copy of each.
304 When the iterable is exhausted, return elements from the saved copy. Repeats
305 indefinitely. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000306
Raymond Hettinger30c73622010-12-01 23:45:20 +0000307 def cycle(iterable):
308 # cycle('ABCD') --> A B C D A B C D A B C D ...
309 saved = []
310 for element in iterable:
311 yield element
312 saved.append(element)
313 while saved:
314 for element in saved:
Georg Brandl116aa622007-08-15 14:28:22 +0000315 yield element
316
Raymond Hettinger30c73622010-12-01 23:45:20 +0000317 Note, this member of the toolkit may require significant auxiliary storage
318 (depending on the length of the iterable).
Georg Brandl116aa622007-08-15 14:28:22 +0000319
320
321.. function:: dropwhile(predicate, iterable)
322
Raymond Hettinger30c73622010-12-01 23:45:20 +0000323 Make an iterator that drops elements from the iterable as long as the predicate
324 is true; afterwards, returns every element. Note, the iterator does not produce
325 *any* output until the predicate first becomes false, so it may have a lengthy
326 start-up time. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000327
Raymond Hettinger30c73622010-12-01 23:45:20 +0000328 def dropwhile(predicate, iterable):
329 # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
330 iterable = iter(iterable)
331 for x in iterable:
332 if not predicate(x):
333 yield x
334 break
335 for x in iterable:
336 yield x
Georg Brandl116aa622007-08-15 14:28:22 +0000337
Raymond Hettinger749761e2009-01-27 04:42:48 +0000338.. function:: filterfalse(predicate, iterable)
339
Raymond Hettinger30c73622010-12-01 23:45:20 +0000340 Make an iterator that filters elements from iterable returning only those for
341 which the predicate is ``False``. If *predicate* is ``None``, return the items
342 that are false. Equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000343
Raymond Hettinger30c73622010-12-01 23:45:20 +0000344 def filterfalse(predicate, iterable):
345 # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
346 if predicate is None:
347 predicate = bool
348 for x in iterable:
349 if not predicate(x):
350 yield x
Raymond Hettinger749761e2009-01-27 04:42:48 +0000351
Georg Brandl116aa622007-08-15 14:28:22 +0000352
Georg Brandl3dd33882009-06-01 17:35:27 +0000353.. function:: groupby(iterable, key=None)
Georg Brandl116aa622007-08-15 14:28:22 +0000354
Raymond Hettinger30c73622010-12-01 23:45:20 +0000355 Make an iterator that returns consecutive keys and groups from the *iterable*.
356 The *key* is a function computing a key value for each element. If not
357 specified or is ``None``, *key* defaults to an identity function and returns
358 the element unchanged. Generally, the iterable needs to already be sorted on
359 the same key function.
Georg Brandl116aa622007-08-15 14:28:22 +0000360
Raymond Hettinger30c73622010-12-01 23:45:20 +0000361 The operation of :func:`groupby` is similar to the ``uniq`` filter in Unix. It
362 generates a break or new group every time the value of the key function changes
363 (which is why it is usually necessary to have sorted the data using the same key
364 function). That behavior differs from SQL's GROUP BY which aggregates common
365 elements regardless of their input order.
Georg Brandl116aa622007-08-15 14:28:22 +0000366
Raymond Hettinger30c73622010-12-01 23:45:20 +0000367 The returned group is itself an iterator that shares the underlying iterable
368 with :func:`groupby`. Because the source is shared, when the :func:`groupby`
369 object is advanced, the previous group is no longer visible. So, if that data
370 is needed later, it should be stored as a list::
Georg Brandl116aa622007-08-15 14:28:22 +0000371
Raymond Hettinger30c73622010-12-01 23:45:20 +0000372 groups = []
373 uniquekeys = []
374 data = sorted(data, key=keyfunc)
375 for k, g in groupby(data, keyfunc):
376 groups.append(list(g)) # Store group iterator as a list
377 uniquekeys.append(k)
Georg Brandl116aa622007-08-15 14:28:22 +0000378
Raymond Hettinger30c73622010-12-01 23:45:20 +0000379 :func:`groupby` is equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000380
Raymond Hettinger30c73622010-12-01 23:45:20 +0000381 class groupby:
382 # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
383 # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
384 def __init__(self, iterable, key=None):
385 if key is None:
386 key = lambda x: x
387 self.keyfunc = key
388 self.it = iter(iterable)
389 self.tgtkey = self.currkey = self.currvalue = object()
390 def __iter__(self):
391 return self
392 def __next__(self):
393 while self.currkey == self.tgtkey:
394 self.currvalue = next(self.it) # Exit on StopIteration
395 self.currkey = self.keyfunc(self.currvalue)
396 self.tgtkey = self.currkey
397 return (self.currkey, self._grouper(self.tgtkey))
398 def _grouper(self, tgtkey):
399 while self.currkey == tgtkey:
400 yield self.currvalue
401 self.currvalue = next(self.it) # Exit on StopIteration
402 self.currkey = self.keyfunc(self.currvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000403
Georg Brandl116aa622007-08-15 14:28:22 +0000404
Ezio Melottie0add762012-09-14 06:32:35 +0300405.. function:: islice(iterable, stop)
406 islice(iterable, start, stop[, step])
Georg Brandl116aa622007-08-15 14:28:22 +0000407
Raymond Hettinger30c73622010-12-01 23:45:20 +0000408 Make an iterator that returns selected elements from the iterable. If *start* is
409 non-zero, then elements from the iterable are skipped until start is reached.
410 Afterward, elements are returned consecutively unless *step* is set higher than
411 one which results in items being skipped. If *stop* is ``None``, then iteration
412 continues until the iterator is exhausted, if at all; otherwise, it stops at the
413 specified position. Unlike regular slicing, :func:`islice` does not support
414 negative values for *start*, *stop*, or *step*. Can be used to extract related
415 fields from data where the internal structure has been flattened (for example, a
416 multi-line report may list a name field on every third line). Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000417
Raymond Hettinger30c73622010-12-01 23:45:20 +0000418 def islice(iterable, *args):
419 # islice('ABCDEFG', 2) --> A B
420 # islice('ABCDEFG', 2, 4) --> C D
421 # islice('ABCDEFG', 2, None) --> C D E F G
422 # islice('ABCDEFG', 0, None, 2) --> A C E G
423 s = slice(*args)
424 it = iter(range(s.start or 0, s.stop or sys.maxsize, s.step or 1))
425 nexti = next(it)
426 for i, element in enumerate(iterable):
427 if i == nexti:
428 yield element
429 nexti = next(it)
Georg Brandl116aa622007-08-15 14:28:22 +0000430
Raymond Hettinger30c73622010-12-01 23:45:20 +0000431 If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
432 then the step defaults to one.
Georg Brandl116aa622007-08-15 14:28:22 +0000433
Georg Brandl116aa622007-08-15 14:28:22 +0000434
Georg Brandl3dd33882009-06-01 17:35:27 +0000435.. function:: permutations(iterable, r=None)
Christian Heimes70e7ea22008-02-28 20:02:27 +0000436
Raymond Hettinger30c73622010-12-01 23:45:20 +0000437 Return successive *r* length permutations of elements in the *iterable*.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000438
Raymond Hettinger30c73622010-12-01 23:45:20 +0000439 If *r* is not specified or is ``None``, then *r* defaults to the length
440 of the *iterable* and all possible full-length permutations
441 are generated.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000442
Raymond Hettinger30c73622010-12-01 23:45:20 +0000443 Permutations are emitted in lexicographic sort order. So, if the
444 input *iterable* is sorted, the permutation tuples will be produced
445 in sorted order.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000446
Raymond Hettinger30c73622010-12-01 23:45:20 +0000447 Elements are treated as unique based on their position, not on their
448 value. So if the input elements are unique, there will be no repeat
449 values in each permutation.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000450
Raymond Hettinger30c73622010-12-01 23:45:20 +0000451 Equivalent to::
Christian Heimesb558a2e2008-03-02 22:46:37 +0000452
453 def permutations(iterable, r=None):
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000454 # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
455 # permutations(range(3)) --> 012 021 102 120 201 210
Christian Heimesb558a2e2008-03-02 22:46:37 +0000456 pool = tuple(iterable)
457 n = len(pool)
458 r = n if r is None else r
Raymond Hettinger5bad41e2009-01-08 21:01:54 +0000459 if r > n:
460 return
461 indices = list(range(n))
Sandro Tosi73866622011-12-25 17:25:45 +0100462 cycles = list(range(n, n-r, -1))
Christian Heimesb558a2e2008-03-02 22:46:37 +0000463 yield tuple(pool[i] for i in indices[:r])
464 while n:
465 for i in reversed(range(r)):
466 cycles[i] -= 1
467 if cycles[i] == 0:
468 indices[i:] = indices[i+1:] + indices[i:i+1]
469 cycles[i] = n - i
470 else:
471 j = cycles[i]
472 indices[i], indices[-j] = indices[-j], indices[i]
473 yield tuple(pool[i] for i in indices[:r])
474 break
475 else:
476 return
Christian Heimes70e7ea22008-02-28 20:02:27 +0000477
Raymond Hettinger30c73622010-12-01 23:45:20 +0000478 The code for :func:`permutations` can be also expressed as a subsequence of
479 :func:`product`, filtered to exclude entries with repeated elements (those
480 from the same position in the input pool)::
Christian Heimes78644762008-03-04 23:39:23 +0000481
482 def permutations(iterable, r=None):
483 pool = tuple(iterable)
484 n = len(pool)
485 r = n if r is None else r
486 for indices in product(range(n), repeat=r):
487 if len(set(indices)) == r:
488 yield tuple(pool[i] for i in indices)
489
Raymond Hettinger30c73622010-12-01 23:45:20 +0000490 The number of items returned is ``n! / (n-r)!`` when ``0 <= r <= n``
491 or zero when ``r > n``.
Christian Heimes70e7ea22008-02-28 20:02:27 +0000492
Georg Brandl3dd33882009-06-01 17:35:27 +0000493.. function:: product(*iterables, repeat=1)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000494
Raymond Hettinger30c73622010-12-01 23:45:20 +0000495 Cartesian product of input iterables.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000496
Raymond Hettinger30c73622010-12-01 23:45:20 +0000497 Equivalent to nested for-loops in a generator expression. For example,
498 ``product(A, B)`` returns the same as ``((x,y) for x in A for y in B)``.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000499
Raymond Hettinger30c73622010-12-01 23:45:20 +0000500 The nested loops cycle like an odometer with the rightmost element advancing
501 on every iteration. This pattern creates a lexicographic ordering so that if
502 the input's iterables are sorted, the product tuples are emitted in sorted
503 order.
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000504
Raymond Hettinger30c73622010-12-01 23:45:20 +0000505 To compute the product of an iterable with itself, specify the number of
506 repetitions with the optional *repeat* keyword argument. For example,
507 ``product(A, repeat=4)`` means the same as ``product(A, A, A, A)``.
Christian Heimes9e7f1d22008-02-28 12:27:11 +0000508
Raymond Hettinger30c73622010-12-01 23:45:20 +0000509 This function is equivalent to the following code, except that the
510 actual implementation does not build up intermediate results in memory::
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000511
Raymond Hettinger30c73622010-12-01 23:45:20 +0000512 def product(*args, repeat=1):
513 # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
514 # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
515 pools = [tuple(pool) for pool in args] * repeat
516 result = [[]]
517 for pool in pools:
518 result = [x+[y] for x in result for y in pool]
519 for prod in result:
520 yield tuple(prod)
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000521
522
Raymond Hettingerd75ad442009-06-01 19:16:52 +0000523.. function:: repeat(object[, times])
Georg Brandl116aa622007-08-15 14:28:22 +0000524
Raymond Hettinger30c73622010-12-01 23:45:20 +0000525 Make an iterator that returns *object* over and over again. Runs indefinitely
526 unless the *times* argument is specified. Used as argument to :func:`map` for
527 invariant parameters to the called function. Also used with :func:`zip` to
528 create an invariant part of a tuple record. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000529
Raymond Hettinger30c73622010-12-01 23:45:20 +0000530 def repeat(object, times=None):
531 # repeat(10, 3) --> 10 10 10
532 if times is None:
533 while True:
534 yield object
535 else:
536 for i in range(times):
537 yield object
Georg Brandl116aa622007-08-15 14:28:22 +0000538
Raymond Hettingerfc3ba6b2012-02-01 09:07:40 -0800539 A common use for *repeat* is to supply a stream of constant values to *map*
540 or *zip*::
541
542 >>> list(map(pow, range(10), repeat(2)))
543 [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
Georg Brandl116aa622007-08-15 14:28:22 +0000544
545.. function:: starmap(function, iterable)
546
Raymond Hettinger30c73622010-12-01 23:45:20 +0000547 Make an iterator that computes the function using arguments obtained from
548 the iterable. Used instead of :func:`map` when argument parameters are already
549 grouped in tuples from a single iterable (the data has been "pre-zipped"). The
550 difference between :func:`map` and :func:`starmap` parallels the distinction
551 between ``function(a,b)`` and ``function(*c)``. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000552
Raymond Hettinger30c73622010-12-01 23:45:20 +0000553 def starmap(function, iterable):
554 # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
555 for args in iterable:
556 yield function(*args)
Christian Heimes679db4a2008-01-18 09:56:22 +0000557
Georg Brandl116aa622007-08-15 14:28:22 +0000558
559.. function:: takewhile(predicate, iterable)
560
Raymond Hettinger30c73622010-12-01 23:45:20 +0000561 Make an iterator that returns elements from the iterable as long as the
562 predicate is true. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000563
Raymond Hettinger30c73622010-12-01 23:45:20 +0000564 def takewhile(predicate, iterable):
565 # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
566 for x in iterable:
567 if predicate(x):
568 yield x
569 else:
570 break
Georg Brandl116aa622007-08-15 14:28:22 +0000571
572
Georg Brandl3dd33882009-06-01 17:35:27 +0000573.. function:: tee(iterable, n=2)
Georg Brandl116aa622007-08-15 14:28:22 +0000574
Raymond Hettinger30c73622010-12-01 23:45:20 +0000575 Return *n* independent iterators from a single iterable. Equivalent to::
Georg Brandl116aa622007-08-15 14:28:22 +0000576
Raymond Hettingercf984ce2009-02-18 20:56:51 +0000577 def tee(iterable, n=2):
578 it = iter(iterable)
579 deques = [collections.deque() for i in range(n)]
580 def gen(mydeque):
581 while True:
582 if not mydeque: # when the local deque is empty
583 newval = next(it) # fetch a new value and
584 for d in deques: # load it to all the deques
585 d.append(newval)
586 yield mydeque.popleft()
587 return tuple(gen(d) for d in deques)
Georg Brandl116aa622007-08-15 14:28:22 +0000588
Raymond Hettinger30c73622010-12-01 23:45:20 +0000589 Once :func:`tee` has made a split, the original *iterable* should not be
590 used anywhere else; otherwise, the *iterable* could get advanced without
591 the tee objects being informed.
Georg Brandl116aa622007-08-15 14:28:22 +0000592
Raymond Hettinger30c73622010-12-01 23:45:20 +0000593 This itertool may require significant auxiliary storage (depending on how
594 much temporary data needs to be stored). In general, if one iterator uses
595 most or all of the data before another iterator starts, it is faster to use
596 :func:`list` instead of :func:`tee`.
Georg Brandl116aa622007-08-15 14:28:22 +0000597
Georg Brandl116aa622007-08-15 14:28:22 +0000598
Georg Brandl3dd33882009-06-01 17:35:27 +0000599.. function:: zip_longest(*iterables, fillvalue=None)
Raymond Hettinger749761e2009-01-27 04:42:48 +0000600
Raymond Hettinger30c73622010-12-01 23:45:20 +0000601 Make an iterator that aggregates elements from each of the iterables. If the
602 iterables are of uneven length, missing values are filled-in with *fillvalue*.
603 Iteration continues until the longest iterable is exhausted. Equivalent to::
Raymond Hettinger749761e2009-01-27 04:42:48 +0000604
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700605 class ZipExhausted(Exception):
606 pass
607
608 def zip_longest(*args, **kwds):
Raymond Hettinger30c73622010-12-01 23:45:20 +0000609 # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700610 fillvalue = kwds.get('fillvalue')
611 counter = len(args) - 1
612 def sentinel():
613 nonlocal counter
614 if not counter:
615 raise ZipExhausted
616 counter -= 1
617 yield fillvalue
Raymond Hettinger30c73622010-12-01 23:45:20 +0000618 fillers = repeat(fillvalue)
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700619 iterators = [chain(it, sentinel(), fillers) for it in args]
Raymond Hettinger30c73622010-12-01 23:45:20 +0000620 try:
Raymond Hettinger6f45d182011-10-30 15:06:14 -0700621 while iterators:
622 yield tuple(map(next, iterators))
623 except ZipExhausted:
Raymond Hettinger30c73622010-12-01 23:45:20 +0000624 pass
Raymond Hettinger749761e2009-01-27 04:42:48 +0000625
Raymond Hettinger30c73622010-12-01 23:45:20 +0000626 If one of the iterables is potentially infinite, then the :func:`zip_longest`
627 function should be wrapped with something that limits the number of calls
628 (for example :func:`islice` or :func:`takewhile`). If not specified,
629 *fillvalue* defaults to ``None``.
Raymond Hettinger749761e2009-01-27 04:42:48 +0000630
631
Georg Brandl116aa622007-08-15 14:28:22 +0000632.. _itertools-recipes:
633
Raymond Hettinger1fa76822010-12-06 23:31:36 +0000634Itertools Recipes
635-----------------
Georg Brandl116aa622007-08-15 14:28:22 +0000636
637This section shows recipes for creating an extended toolset using the existing
638itertools as building blocks.
639
640The extended tools offer the same high performance as the underlying toolset.
641The superior memory performance is kept by processing elements one at a time
642rather than bringing the whole iterable into memory all at once. Code volume is
643kept small by linking the tools together in a functional style which helps
644eliminate temporary variables. High speed is retained by preferring
Georg Brandl9afde1c2007-11-01 20:32:30 +0000645"vectorized" building blocks over the use of for-loops and :term:`generator`\s
Christian Heimesfe337bf2008-03-23 21:54:12 +0000646which incur interpreter overhead.
647
648.. testcode::
Georg Brandl116aa622007-08-15 14:28:22 +0000649
Raymond Hettinger30c73622010-12-01 23:45:20 +0000650 def take(n, iterable):
651 "Return first n items of the iterable as a list"
652 return list(islice(iterable, n))
Georg Brandl116aa622007-08-15 14:28:22 +0000653
Raymond Hettinger30c73622010-12-01 23:45:20 +0000654 def tabulate(function, start=0):
655 "Return function(0), function(1), ..."
656 return map(function, count(start))
Raymond Hettingerdd1150e2008-03-13 02:39:40 +0000657
Raymond Hettinger30c73622010-12-01 23:45:20 +0000658 def consume(iterator, n):
659 "Advance the iterator n-steps ahead. If n is none, consume entirely."
660 # Use functions that consume iterators at C speed.
661 if n is None:
662 # feed the entire iterator into a zero-length deque
663 collections.deque(iterator, maxlen=0)
664 else:
665 # advance to the empty slice starting at position n
666 next(islice(iterator, n, n), None)
Raymond Hettingerfa007962009-03-09 11:55:25 +0000667
Raymond Hettinger30c73622010-12-01 23:45:20 +0000668 def nth(iterable, n, default=None):
669 "Returns the nth item or a default value"
670 return next(islice(iterable, n, None), default)
Georg Brandl116aa622007-08-15 14:28:22 +0000671
Raymond Hettinger30c73622010-12-01 23:45:20 +0000672 def quantify(iterable, pred=bool):
673 "Count how many times the predicate is true"
674 return sum(map(pred, iterable))
Georg Brandl116aa622007-08-15 14:28:22 +0000675
Raymond Hettinger30c73622010-12-01 23:45:20 +0000676 def padnone(iterable):
677 """Returns the sequence elements and then returns None indefinitely.
Georg Brandl116aa622007-08-15 14:28:22 +0000678
Raymond Hettinger30c73622010-12-01 23:45:20 +0000679 Useful for emulating the behavior of the built-in map() function.
680 """
681 return chain(iterable, repeat(None))
Georg Brandl116aa622007-08-15 14:28:22 +0000682
Raymond Hettinger30c73622010-12-01 23:45:20 +0000683 def ncycles(iterable, n):
684 "Returns the sequence elements n times"
685 return chain.from_iterable(repeat(tuple(iterable), n))
Georg Brandl116aa622007-08-15 14:28:22 +0000686
Raymond Hettinger30c73622010-12-01 23:45:20 +0000687 def dotproduct(vec1, vec2):
688 return sum(map(operator.mul, vec1, vec2))
Georg Brandl116aa622007-08-15 14:28:22 +0000689
Raymond Hettinger30c73622010-12-01 23:45:20 +0000690 def flatten(listOfLists):
691 "Flatten one level of nesting"
692 return chain.from_iterable(listOfLists)
Georg Brandl116aa622007-08-15 14:28:22 +0000693
Raymond Hettinger30c73622010-12-01 23:45:20 +0000694 def repeatfunc(func, times=None, *args):
695 """Repeat calls to func with specified arguments.
Georg Brandl116aa622007-08-15 14:28:22 +0000696
Raymond Hettinger30c73622010-12-01 23:45:20 +0000697 Example: repeatfunc(random.random)
698 """
699 if times is None:
700 return starmap(func, repeat(args))
701 return starmap(func, repeat(args, times))
Georg Brandl116aa622007-08-15 14:28:22 +0000702
Raymond Hettinger30c73622010-12-01 23:45:20 +0000703 def pairwise(iterable):
704 "s -> (s0,s1), (s1,s2), (s2, s3), ..."
705 a, b = tee(iterable)
706 next(b, None)
707 return zip(a, b)
Georg Brandl116aa622007-08-15 14:28:22 +0000708
Raymond Hettinger44571da2013-05-05 19:53:41 -0700709 def grouper(iterable, n, fillvalue=None):
Raymond Hettinger9ae94732012-07-08 16:04:03 -0700710 "Collect data into fixed-length chunks or blocks"
Raymond Hettinger44571da2013-05-05 19:53:41 -0700711 # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
Raymond Hettinger30c73622010-12-01 23:45:20 +0000712 args = [iter(iterable)] * n
713 return zip_longest(*args, fillvalue=fillvalue)
Georg Brandl116aa622007-08-15 14:28:22 +0000714
Raymond Hettinger30c73622010-12-01 23:45:20 +0000715 def roundrobin(*iterables):
716 "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
717 # Recipe credited to George Sakkis
718 pending = len(iterables)
719 nexts = cycle(iter(it).__next__ for it in iterables)
720 while pending:
721 try:
722 for next in nexts:
723 yield next()
724 except StopIteration:
725 pending -= 1
726 nexts = cycle(islice(nexts, pending))
Georg Brandl116aa622007-08-15 14:28:22 +0000727
Raymond Hettinger30c73622010-12-01 23:45:20 +0000728 def partition(pred, iterable):
729 'Use a predicate to partition entries into false entries and true entries'
730 # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9
731 t1, t2 = tee(iterable)
732 return filterfalse(pred, t1), filter(pred, t2)
Raymond Hettinger5ce0aa22010-12-01 10:49:19 +0000733
Raymond Hettinger30c73622010-12-01 23:45:20 +0000734 def powerset(iterable):
735 "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
736 s = list(iterable)
737 return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
Raymond Hettinger08d01ee2010-08-07 05:36:53 +0000738
Raymond Hettinger30c73622010-12-01 23:45:20 +0000739 def unique_everseen(iterable, key=None):
740 "List unique elements, preserving order. Remember all elements ever seen."
741 # unique_everseen('AAAABBBCCDAABBB') --> A B C D
742 # unique_everseen('ABBCcAD', str.lower) --> A B C D
743 seen = set()
744 seen_add = seen.add
745 if key is None:
746 for element in filterfalse(seen.__contains__, iterable):
747 seen_add(element)
748 yield element
749 else:
750 for element in iterable:
751 k = key(element)
752 if k not in seen:
753 seen_add(k)
754 yield element
Christian Heimes90c3d9b2008-02-23 13:18:03 +0000755
Raymond Hettinger30c73622010-12-01 23:45:20 +0000756 def unique_justseen(iterable, key=None):
757 "List unique elements, preserving order. Remember only the element just seen."
758 # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
759 # unique_justseen('ABBCcAD', str.lower) --> A B C A D
760 return map(next, map(itemgetter(1), groupby(iterable, key)))
Raymond Hettingerad9d96b2009-01-02 21:39:07 +0000761
Raymond Hettinger30c73622010-12-01 23:45:20 +0000762 def iter_except(func, exception, first=None):
763 """ Call a function repeatedly until an exception is raised.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000764
Raymond Hettinger30c73622010-12-01 23:45:20 +0000765 Converts a call-until-exception interface to an iterator interface.
Andrew Kuchling1d7d5802013-06-20 21:17:41 -0400766 Like builtins.iter(func, sentinel) but uses an exception instead
Raymond Hettinger30c73622010-12-01 23:45:20 +0000767 of a sentinel to end the loop.
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000768
Raymond Hettinger30c73622010-12-01 23:45:20 +0000769 Examples:
770 iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
771 iter_except(d.popitem, KeyError) # non-blocking dict iterator
772 iter_except(d.popleft, IndexError) # non-blocking deque iterator
773 iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
774 iter_except(s.pop, KeyError) # non-blocking set iterator
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000775
Raymond Hettinger30c73622010-12-01 23:45:20 +0000776 """
777 try:
778 if first is not None:
779 yield first() # For database APIs needing an initial cast to db.first()
780 while 1:
781 yield func()
782 except exception:
783 pass
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000784
Raymond Hettinger30c73622010-12-01 23:45:20 +0000785 def random_product(*args, repeat=1):
786 "Random selection from itertools.product(*args, **kwds)"
787 pools = [tuple(pool) for pool in args] * repeat
788 return tuple(random.choice(pool) for pool in pools)
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000789
Raymond Hettinger30c73622010-12-01 23:45:20 +0000790 def random_permutation(iterable, r=None):
791 "Random selection from itertools.permutations(iterable, r)"
792 pool = tuple(iterable)
793 r = len(pool) if r is None else r
794 return tuple(random.sample(pool, r))
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000795
Raymond Hettinger30c73622010-12-01 23:45:20 +0000796 def random_combination(iterable, r):
797 "Random selection from itertools.combinations(iterable, r)"
798 pool = tuple(iterable)
799 n = len(pool)
800 indices = sorted(random.sample(range(n), r))
801 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000802
Raymond Hettinger30c73622010-12-01 23:45:20 +0000803 def random_combination_with_replacement(iterable, r):
804 "Random selection from itertools.combinations_with_replacement(iterable, r)"
805 pool = tuple(iterable)
806 n = len(pool)
807 indices = sorted(random.randrange(n) for i in range(r))
808 return tuple(pool[i] for i in indices)
Raymond Hettinger063a4b62010-04-02 04:18:18 +0000809
Raymond Hettingerfc91aa22010-03-28 18:27:13 +0000810Note, many of the above recipes can be optimized by replacing global lookups
811with local variables defined as default values. For example, the
812*dotproduct* recipe can be written as::
813
Raymond Hettinger30c73622010-12-01 23:45:20 +0000814 def dotproduct(vec1, vec2, sum=sum, map=map, mul=operator.mul):
815 return sum(map(mul, vec1, vec2))