Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1from collections import Counter, defaultdict, deque 

2from functools import partial, wraps 

3from heapq import merge 

4from itertools import ( 

5 chain, 

6 compress, 

7 count, 

8 cycle, 

9 dropwhile, 

10 groupby, 

11 islice, 

12 repeat, 

13 starmap, 

14 takewhile, 

15 tee, 

16 zip_longest, 

17) 

18from operator import itemgetter, lt, gt, sub 

19from sys import maxsize, version_info 

20from collections.abc import Sequence 

21 

22from .recipes import consume, flatten, take 

23 

24__all__ = [ 

25 'adjacent', 

26 'always_iterable', 

27 'always_reversible', 

28 'bucket', 

29 'chunked', 

30 'circular_shifts', 

31 'collapse', 

32 'collate', 

33 'consecutive_groups', 

34 'consumer', 

35 'count_cycle', 

36 'difference', 

37 'distinct_permutations', 

38 'distribute', 

39 'divide', 

40 'exactly_n', 

41 'first', 

42 'groupby_transform', 

43 'ilen', 

44 'interleave_longest', 

45 'interleave', 

46 'intersperse', 

47 'islice_extended', 

48 'iterate', 

49 'last', 

50 'locate', 

51 'lstrip', 

52 'make_decorator', 

53 'map_reduce', 

54 'numeric_range', 

55 'one', 

56 'padded', 

57 'peekable', 

58 'replace', 

59 'rlocate', 

60 'rstrip', 

61 'run_length', 

62 'seekable', 

63 'SequenceView', 

64 'side_effect', 

65 'sliced', 

66 'sort_together', 

67 'split_at', 

68 'split_after', 

69 'split_before', 

70 'split_into', 

71 'spy', 

72 'stagger', 

73 'strip', 

74 'substrings', 

75 'unique_to_each', 

76 'unzip', 

77 'windowed', 

78 'with_iter', 

79 'zip_offset', 

80] 

81 

82_marker = object() 

83 

84 

85def chunked(iterable, n): 

86 """Break *iterable* into lists of length *n*: 

87 

88 >>> list(chunked([1, 2, 3, 4, 5, 6], 3)) 

89 [[1, 2, 3], [4, 5, 6]] 

90 

91 If the length of *iterable* is not evenly divisible by *n*, the last 

92 returned list will be shorter: 

93 

94 >>> list(chunked([1, 2, 3, 4, 5, 6, 7, 8], 3)) 

95 [[1, 2, 3], [4, 5, 6], [7, 8]] 

96 

97 To use a fill-in value instead, see the :func:`grouper` recipe. 

98 

99 :func:`chunked` is useful for splitting up a computation on a large number 

100 of keys into batches, to be pickled and sent off to worker processes. One 

101 example is operations on rows in MySQL, which does not implement 

102 server-side cursors properly and would otherwise load the entire dataset 

103 into RAM on the client. 

104 

105 """ 

106 return iter(partial(take, n, iter(iterable)), []) 

107 

108 

109def first(iterable, default=_marker): 

110 """Return the first item of *iterable*, or *default* if *iterable* is 

111 empty. 

112 

113 >>> first([0, 1, 2, 3]) 

114 0 

115 >>> first([], 'some default') 

116 'some default' 

117 

118 If *default* is not provided and there are no items in the iterable, 

119 raise ``ValueError``. 

120 

121 :func:`first` is useful when you have a generator of expensive-to-retrieve 

122 values and want any arbitrary one. It is marginally shorter than 

123 ``next(iter(iterable), default)``. 

124 

125 """ 

126 try: 

127 return next(iter(iterable)) 

128 except StopIteration: 

129 # I'm on the edge about raising ValueError instead of StopIteration. At 

130 # the moment, ValueError wins, because the caller could conceivably 

131 # want to do something different with flow control when I raise the 

132 # exception, and it's weird to explicitly catch StopIteration. 

133 if default is _marker: 

134 raise ValueError('first() was called on an empty iterable, and no ' 

135 'default value was provided.') 

136 return default 

137 

138 

139def last(iterable, default=_marker): 

140 """Return the last item of *iterable*, or *default* if *iterable* is 

141 empty. 

142 

143 >>> last([0, 1, 2, 3]) 

144 3 

145 >>> last([], 'some default') 

146 'some default' 

147 

148 If *default* is not provided and there are no items in the iterable, 

149 raise ``ValueError``. 

150 """ 

151 try: 

152 try: 

153 # Try to access the last item directly 

154 return iterable[-1] 

155 except (TypeError, AttributeError, KeyError): 

156 # If not slice-able, iterate entirely using length-1 deque 

157 return deque(iterable, maxlen=1)[0] 

158 except IndexError: # If the iterable was empty 

159 if default is _marker: 

160 raise ValueError('last() was called on an empty iterable, and no ' 

161 'default value was provided.') 

162 return default 

163 

164 

165class peekable: 

166 """Wrap an iterator to allow lookahead and prepending elements. 

167 

168 Call :meth:`peek` on the result to get the value that will be returned 

169 by :func:`next`. This won't advance the iterator: 

170 

171 >>> p = peekable(['a', 'b']) 

172 >>> p.peek() 

173 'a' 

174 >>> next(p) 

175 'a' 

176 

177 Pass :meth:`peek` a default value to return that instead of raising 

178 ``StopIteration`` when the iterator is exhausted. 

179 

180 >>> p = peekable([]) 

181 >>> p.peek('hi') 

182 'hi' 

183 

184 peekables also offer a :meth:`prepend` method, which "inserts" items 

185 at the head of the iterable: 

186 

187 >>> p = peekable([1, 2, 3]) 

188 >>> p.prepend(10, 11, 12) 

189 >>> next(p) 

190 10 

191 >>> p.peek() 

192 11 

193 >>> list(p) 

194 [11, 12, 1, 2, 3] 

195 

196 peekables can be indexed. Index 0 is the item that will be returned by 

197 :func:`next`, index 1 is the item after that, and so on: 

198 The values up to the given index will be cached. 

199 

200 >>> p = peekable(['a', 'b', 'c', 'd']) 

201 >>> p[0] 

202 'a' 

203 >>> p[1] 

204 'b' 

205 >>> next(p) 

206 'a' 

207 

208 Negative indexes are supported, but be aware that they will cache the 

209 remaining items in the source iterator, which may require significant 

210 storage. 

211 

212 To check whether a peekable is exhausted, check its truth value: 

213 

214 >>> p = peekable(['a', 'b']) 

215 >>> if p: # peekable has items 

216 ... list(p) 

217 ['a', 'b'] 

218 >>> if not p: # peekable is exhaused 

219 ... list(p) 

220 [] 

221 

222 """ 

223 def __init__(self, iterable): 

224 self._it = iter(iterable) 

225 self._cache = deque() 

226 

227 def __iter__(self): 

228 return self 

229 

230 def __bool__(self): 

231 try: 

232 self.peek() 

233 except StopIteration: 

234 return False 

235 return True 

236 

237 def peek(self, default=_marker): 

238 """Return the item that will be next returned from ``next()``. 

239 

240 Return ``default`` if there are no items left. If ``default`` is not 

241 provided, raise ``StopIteration``. 

242 

243 """ 

244 if not self._cache: 

245 try: 

246 self._cache.append(next(self._it)) 

247 except StopIteration: 

248 if default is _marker: 

249 raise 

250 return default 

251 return self._cache[0] 

252 

253 def prepend(self, *items): 

254 """Stack up items to be the next ones returned from ``next()`` or 

255 ``self.peek()``. The items will be returned in 

256 first in, first out order:: 

257 

258 >>> p = peekable([1, 2, 3]) 

259 >>> p.prepend(10, 11, 12) 

260 >>> next(p) 

261 10 

262 >>> list(p) 

263 [11, 12, 1, 2, 3] 

264 

265 It is possible, by prepending items, to "resurrect" a peekable that 

266 previously raised ``StopIteration``. 

267 

268 >>> p = peekable([]) 

269 >>> next(p) 

270 Traceback (most recent call last): 

271 ... 

272 StopIteration 

273 >>> p.prepend(1) 

274 >>> next(p) 

275 1 

276 >>> next(p) 

277 Traceback (most recent call last): 

278 ... 

279 StopIteration 

280 

281 """ 

282 self._cache.extendleft(reversed(items)) 

283 

284 def __next__(self): 

285 if self._cache: 

286 return self._cache.popleft() 

287 

288 return next(self._it) 

289 

290 def _get_slice(self, index): 

291 # Normalize the slice's arguments 

292 step = 1 if (index.step is None) else index.step 

293 if step > 0: 

294 start = 0 if (index.start is None) else index.start 

295 stop = maxsize if (index.stop is None) else index.stop 

296 elif step < 0: 

297 start = -1 if (index.start is None) else index.start 

298 stop = (-maxsize - 1) if (index.stop is None) else index.stop 

299 else: 

300 raise ValueError('slice step cannot be zero') 

301 

302 # If either the start or stop index is negative, we'll need to cache 

303 # the rest of the iterable in order to slice from the right side. 

304 if (start < 0) or (stop < 0): 

305 self._cache.extend(self._it) 

306 # Otherwise we'll need to find the rightmost index and cache to that 

307 # point. 

308 else: 

309 n = min(max(start, stop) + 1, maxsize) 

310 cache_len = len(self._cache) 

311 if n >= cache_len: 

312 self._cache.extend(islice(self._it, n - cache_len)) 

313 

314 return list(self._cache)[index] 

315 

316 def __getitem__(self, index): 

317 if isinstance(index, slice): 

318 return self._get_slice(index) 

319 

320 cache_len = len(self._cache) 

321 if index < 0: 

322 self._cache.extend(self._it) 

323 elif index >= cache_len: 

324 self._cache.extend(islice(self._it, index + 1 - cache_len)) 

325 

326 return self._cache[index] 

327 

328 

329def _collate(*iterables, key=lambda a: a, reverse=False): 

330 """Helper for ``collate()``, called when the user is using the ``reverse`` 

331 or ``key`` keyword arguments on Python versions below 3.5. 

332 

333 """ 

334 min_or_max = partial(max if reverse else min, key=itemgetter(0)) 

335 peekables = [peekable(it) for it in iterables] 

336 peekables = [p for p in peekables if p] # Kill empties. 

337 while peekables: 

338 _, p = min_or_max((key(p.peek()), p) for p in peekables) 

339 yield next(p) 

340 peekables = [x for x in peekables if x] 

341 

342 

343def collate(*iterables, **kwargs): 

344 """Return a sorted merge of the items from each of several already-sorted 

345 *iterables*. 

346 

347 >>> list(collate('ACDZ', 'AZ', 'JKL')) 

348 ['A', 'A', 'C', 'D', 'J', 'K', 'L', 'Z', 'Z'] 

349 

350 Works lazily, keeping only the next value from each iterable in memory. Use 

351 :func:`collate` to, for example, perform a n-way mergesort of items that 

352 don't fit in memory. 

353 

354 If a *key* function is specified, the iterables will be sorted according 

355 to its result: 

356 

357 >>> key = lambda s: int(s) # Sort by numeric value, not by string 

358 >>> list(collate(['1', '10'], ['2', '11'], key=key)) 

359 ['1', '2', '10', '11'] 

360 

361 

362 If the *iterables* are sorted in descending order, set *reverse* to 

363 ``True``: 

364 

365 >>> list(collate([5, 3, 1], [4, 2, 0], reverse=True)) 

366 [5, 4, 3, 2, 1, 0] 

367 

368 If the elements of the passed-in iterables are out of order, you might get 

369 unexpected results. 

370 

371 On Python 3.5+, this function is an alias for :func:`heapq.merge`. 

372 

373 """ 

374 if not kwargs: 

375 return merge(*iterables) 

376 

377 return _collate(*iterables, **kwargs) 

378 

379 

380# If using Python version 3.5 or greater, heapq.merge() will be faster than 

381# collate - use that instead. 

382if version_info >= (3, 5, 0): 

383 _collate_docstring = collate.__doc__ 

384 collate = partial(merge) 

385 collate.__doc__ = _collate_docstring 

386 

387 

388def consumer(func): 

389 """Decorator that automatically advances a PEP-342-style "reverse iterator" 

390 to its first yield point so you don't have to call ``next()`` on it 

391 manually. 

392 

393 >>> @consumer 

394 ... def tally(): 

395 ... i = 0 

396 ... while True: 

397 ... print('Thing number %s is %s.' % (i, (yield))) 

398 ... i += 1 

399 ... 

400 >>> t = tally() 

401 >>> t.send('red') 

402 Thing number 0 is red. 

403 >>> t.send('fish') 

404 Thing number 1 is fish. 

405 

406 Without the decorator, you would have to call ``next(t)`` before 

407 ``t.send()`` could be used. 

408 

409 """ 

410 @wraps(func) 

411 def wrapper(*args, **kwargs): 

412 gen = func(*args, **kwargs) 

413 next(gen) 

414 return gen 

415 return wrapper 

416 

417 

418def ilen(iterable): 

419 """Return the number of items in *iterable*. 

420 

421 >>> ilen(x for x in range(1000000) if x % 3 == 0) 

422 333334 

423 

424 This consumes the iterable, so handle with care. 

425 

426 """ 

427 # This approach was selected because benchmarks showed it's likely the 

428 # fastest of the known implementations at the time of writing. 

429 # See GitHub tracker: #236, #230. 

430 counter = count() 

431 deque(zip(iterable, counter), maxlen=0) 

432 return next(counter) 

433 

434 

435def iterate(func, start): 

436 """Return ``start``, ``func(start)``, ``func(func(start))``, ... 

437 

438 >>> from itertools import islice 

439 >>> list(islice(iterate(lambda x: 2*x, 1), 10)) 

440 [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] 

441 

442 """ 

443 while True: 

444 yield start 

445 start = func(start) 

446 

447 

448def with_iter(context_manager): 

449 """Wrap an iterable in a ``with`` statement, so it closes once exhausted. 

450 

451 For example, this will close the file when the iterator is exhausted:: 

452 

453 upper_lines = (line.upper() for line in with_iter(open('foo'))) 

454 

455 Any context manager which returns an iterable is a candidate for 

456 ``with_iter``. 

457 

458 """ 

459 with context_manager as iterable: 

460 yield from iterable 

461 

462 

463def one(iterable, too_short=None, too_long=None): 

464 """Return the first item from *iterable*, which is expected to contain only 

465 that item. Raise an exception if *iterable* is empty or has more than one 

466 item. 

467 

468 :func:`one` is useful for ensuring that an iterable contains only one item. 

469 For example, it can be used to retrieve the result of a database query 

470 that is expected to return a single row. 

471 

472 If *iterable* is empty, ``ValueError`` will be raised. You may specify a 

473 different exception with the *too_short* keyword: 

474 

475 >>> it = [] 

476 >>> one(it) # doctest: +IGNORE_EXCEPTION_DETAIL 

477 Traceback (most recent call last): 

478 ... 

479 ValueError: too many items in iterable (expected 1)' 

480 >>> too_short = IndexError('too few items') 

481 >>> one(it, too_short=too_short) # doctest: +IGNORE_EXCEPTION_DETAIL 

482 Traceback (most recent call last): 

483 ... 

484 IndexError: too few items 

485 

486 Similarly, if *iterable* contains more than one item, ``ValueError`` will 

487 be raised. You may specify a different exception with the *too_long* 

488 keyword: 

489 

490 >>> it = ['too', 'many'] 

491 >>> one(it) # doctest: +IGNORE_EXCEPTION_DETAIL 

492 Traceback (most recent call last): 

493 ... 

494 ValueError: too many items in iterable (expected 1)' 

495 >>> too_long = RuntimeError 

496 >>> one(it, too_long=too_long) # doctest: +IGNORE_EXCEPTION_DETAIL 

497 Traceback (most recent call last): 

498 ... 

499 RuntimeError 

500 

501 Note that :func:`one` attempts to advance *iterable* twice to ensure there 

502 is only one item. If there is more than one, both items will be discarded. 

503 See :func:`spy` or :func:`peekable` to check iterable contents less 

504 destructively. 

505 

506 """ 

507 it = iter(iterable) 

508 

509 try: 

510 value = next(it) 

511 except StopIteration: 

512 raise too_short or ValueError('too few items in iterable (expected 1)') 

513 

514 try: 

515 next(it) 

516 except StopIteration: 

517 pass 

518 else: 

519 raise too_long or ValueError('too many items in iterable (expected 1)') 

520 

521 return value 

522 

523 

524def distinct_permutations(iterable): 

525 """Yield successive distinct permutations of the elements in *iterable*. 

526 

527 >>> sorted(distinct_permutations([1, 0, 1])) 

528 [(0, 1, 1), (1, 0, 1), (1, 1, 0)] 

529 

530 Equivalent to ``set(permutations(iterable))``, except duplicates are not 

531 generated and thrown away. For larger input sequences this is much more 

532 efficient. 

533 

534 Duplicate permutations arise when there are duplicated elements in the 

535 input iterable. The number of items returned is 

536 `n! / (x_1! * x_2! * ... * x_n!)`, where `n` is the total number of 

537 items input, and each `x_i` is the count of a distinct item in the input 

538 sequence. 

539 

540 """ 

541 def make_new_permutations(permutations, e): 

542 """Internal helper function. 

543 The output permutations are built up by adding element *e* to the 

544 current *permutations* at every possible position. 

545 The key idea is to keep repeated elements (reverse) ordered: 

546 if e1 == e2 and e1 is before e2 in the iterable, then all permutations 

547 with e1 before e2 are ignored. 

548 

549 """ 

550 for permutation in permutations: 

551 for j in range(len(permutation)): 

552 yield permutation[:j] + [e] + permutation[j:] 

553 if permutation[j] == e: 

554 break 

555 else: 

556 yield permutation + [e] 

557 

558 permutations = [[]] 

559 for e in iterable: 

560 permutations = make_new_permutations(permutations, e) 

561 

562 return (tuple(t) for t in permutations) 

563 

564 

565def intersperse(e, iterable, n=1): 

566 """Intersperse filler element *e* among the items in *iterable*, leaving 

567 *n* items between each filler element. 

568 

569 >>> list(intersperse('!', [1, 2, 3, 4, 5])) 

570 [1, '!', 2, '!', 3, '!', 4, '!', 5] 

571 

572 >>> list(intersperse(None, [1, 2, 3, 4, 5], n=2)) 

573 [1, 2, None, 3, 4, None, 5] 

574 

575 """ 

576 if n == 0: 

577 raise ValueError('n must be > 0') 

578 elif n == 1: 

579 # interleave(repeat(e), iterable) -> e, x_0, e, e, x_1, e, x_2... 

580 # islice(..., 1, None) -> x_0, e, e, x_1, e, x_2... 

581 return islice(interleave(repeat(e), iterable), 1, None) 

582 else: 

583 # interleave(filler, chunks) -> [e], [x_0, x_1], [e], [x_2, x_3]... 

584 # islice(..., 1, None) -> [x_0, x_1], [e], [x_2, x_3]... 

585 # flatten(...) -> x_0, x_1, e, x_2, x_3... 

586 filler = repeat([e]) 

587 chunks = chunked(iterable, n) 

588 return flatten(islice(interleave(filler, chunks), 1, None)) 

589 

590 

591def unique_to_each(*iterables): 

592 """Return the elements from each of the input iterables that aren't in the 

593 other input iterables. 

594 

595 For example, suppose you have a set of packages, each with a set of 

596 dependencies:: 

597 

598 {'pkg_1': {'A', 'B'}, 'pkg_2': {'B', 'C'}, 'pkg_3': {'B', 'D'}} 

599 

600 If you remove one package, which dependencies can also be removed? 

601 

602 If ``pkg_1`` is removed, then ``A`` is no longer necessary - it is not 

603 associated with ``pkg_2`` or ``pkg_3``. Similarly, ``C`` is only needed for 

604 ``pkg_2``, and ``D`` is only needed for ``pkg_3``:: 

605 

606 >>> unique_to_each({'A', 'B'}, {'B', 'C'}, {'B', 'D'}) 

607 [['A'], ['C'], ['D']] 

608 

609 If there are duplicates in one input iterable that aren't in the others 

610 they will be duplicated in the output. Input order is preserved:: 

611 

612 >>> unique_to_each("mississippi", "missouri") 

613 [['p', 'p'], ['o', 'u', 'r']] 

614 

615 It is assumed that the elements of each iterable are hashable. 

616 

617 """ 

618 pool = [list(it) for it in iterables] 

619 counts = Counter(chain.from_iterable(map(set, pool))) 

620 uniques = {element for element in counts if counts[element] == 1} 

621 return [list(filter(uniques.__contains__, it)) for it in pool] 

622 

623 

624def windowed(seq, n, fillvalue=None, step=1): 

625 """Return a sliding window of width *n* over the given iterable. 

626 

627 >>> all_windows = windowed([1, 2, 3, 4, 5], 3) 

628 >>> list(all_windows) 

629 [(1, 2, 3), (2, 3, 4), (3, 4, 5)] 

630 

631 When the window is larger than the iterable, *fillvalue* is used in place 

632 of missing values:: 

633 

634 >>> list(windowed([1, 2, 3], 4)) 

635 [(1, 2, 3, None)] 

636 

637 Each window will advance in increments of *step*: 

638 

639 >>> list(windowed([1, 2, 3, 4, 5, 6], 3, fillvalue='!', step=2)) 

640 [(1, 2, 3), (3, 4, 5), (5, 6, '!')] 

641 

642 To slide into the iterable's items, use :func:`chain` to add filler items 

643 to the left: 

644 

645 >>> iterable = [1, 2, 3, 4] 

646 >>> n = 3 

647 >>> padding = [None] * (n - 1) 

648 >>> list(windowed(chain(padding, iterable), 3)) 

649 [(None, None, 1), (None, 1, 2), (1, 2, 3), (2, 3, 4)] 

650 

651 """ 

652 if n < 0: 

653 raise ValueError('n must be >= 0') 

654 if n == 0: 

655 yield tuple() 

656 return 

657 if step < 1: 

658 raise ValueError('step must be >= 1') 

659 

660 it = iter(seq) 

661 window = deque([], n) 

662 append = window.append 

663 

664 # Initial deque fill 

665 for _ in range(n): 

666 append(next(it, fillvalue)) 

667 yield tuple(window) 

668 

669 # Appending new items to the right causes old items to fall off the left 

670 i = 0 

671 for item in it: 

672 append(item) 

673 i = (i + 1) % step 

674 if i % step == 0: 

675 yield tuple(window) 

676 

677 # If there are items from the iterable in the window, pad with the given 

678 # value and emit them. 

679 if (i % step) and (step - i < n): 

680 for _ in range(step - i): 

681 append(fillvalue) 

682 yield tuple(window) 

683 

684 

685def substrings(iterable): 

686 """Yield all of the substrings of *iterable*. 

687 

688 >>> [''.join(s) for s in substrings('more')] 

689 ['m', 'o', 'r', 'e', 'mo', 'or', 're', 'mor', 'ore', 'more'] 

690 

691 Note that non-string iterables can also be subdivided. 

692 

693 >>> list(substrings([0, 1, 2])) 

694 [(0,), (1,), (2,), (0, 1), (1, 2), (0, 1, 2)] 

695 

696 """ 

697 # The length-1 substrings 

698 seq = [] 

699 for item in iter(iterable): 

700 seq.append(item) 

701 yield (item,) 

702 seq = tuple(seq) 

703 item_count = len(seq) 

704 

705 # And the rest 

706 for n in range(2, item_count + 1): 

707 for i in range(item_count - n + 1): 

708 yield seq[i:i + n] 

709 

710 

711class bucket: 

712 """Wrap *iterable* and return an object that buckets it iterable into 

713 child iterables based on a *key* function. 

714 

715 >>> iterable = ['a1', 'b1', 'c1', 'a2', 'b2', 'c2', 'b3'] 

716 >>> s = bucket(iterable, key=lambda x: x[0]) 

717 >>> a_iterable = s['a'] 

718 >>> next(a_iterable) 

719 'a1' 

720 >>> next(a_iterable) 

721 'a2' 

722 >>> list(s['b']) 

723 ['b1', 'b2', 'b3'] 

724 

725 The original iterable will be advanced and its items will be cached until 

726 they are used by the child iterables. This may require significant storage. 

727 

728 By default, attempting to select a bucket to which no items belong will 

729 exhaust the iterable and cache all values. 

730 If you specify a *validator* function, selected buckets will instead be 

731 checked against it. 

732 

733 >>> from itertools import count 

734 >>> it = count(1, 2) # Infinite sequence of odd numbers 

735 >>> key = lambda x: x % 10 # Bucket by last digit 

736 >>> validator = lambda x: x in {1, 3, 5, 7, 9} # Odd digits only 

737 >>> s = bucket(it, key=key, validator=validator) 

738 >>> 2 in s 

739 False 

740 >>> list(s[2]) 

741 [] 

742 

743 """ 

744 def __init__(self, iterable, key, validator=None): 

745 self._it = iter(iterable) 

746 self._key = key 

747 self._cache = defaultdict(deque) 

748 self._validator = validator or (lambda x: True) 

749 

750 def __contains__(self, value): 

751 if not self._validator(value): 

752 return False 

753 

754 try: 

755 item = next(self[value]) 

756 except StopIteration: 

757 return False 

758 else: 

759 self._cache[value].appendleft(item) 

760 

761 return True 

762 

763 def _get_values(self, value): 

764 """ 

765 Helper to yield items from the parent iterator that match *value*. 

766 Items that don't match are stored in the local cache as they 

767 are encountered. 

768 """ 

769 while True: 

770 # If we've cached some items that match the target value, emit 

771 # the first one and evict it from the cache. 

772 if self._cache[value]: 

773 yield self._cache[value].popleft() 

774 # Otherwise we need to advance the parent iterator to search for 

775 # a matching item, caching the rest. 

776 else: 

777 while True: 

778 try: 

779 item = next(self._it) 

780 except StopIteration: 

781 return 

782 item_value = self._key(item) 

783 if item_value == value: 

784 yield item 

785 break 

786 elif self._validator(item_value): 

787 self._cache[item_value].append(item) 

788 

789 def __getitem__(self, value): 

790 if not self._validator(value): 

791 return iter(()) 

792 

793 return self._get_values(value) 

794 

795 

796def spy(iterable, n=1): 

797 """Return a 2-tuple with a list containing the first *n* elements of 

798 *iterable*, and an iterator with the same items as *iterable*. 

799 This allows you to "look ahead" at the items in the iterable without 

800 advancing it. 

801 

802 There is one item in the list by default: 

803 

804 >>> iterable = 'abcdefg' 

805 >>> head, iterable = spy(iterable) 

806 >>> head 

807 ['a'] 

808 >>> list(iterable) 

809 ['a', 'b', 'c', 'd', 'e', 'f', 'g'] 

810 

811 You may use unpacking to retrieve items instead of lists: 

812 

813 >>> (head,), iterable = spy('abcdefg') 

814 >>> head 

815 'a' 

816 >>> (first, second), iterable = spy('abcdefg', 2) 

817 >>> first 

818 'a' 

819 >>> second 

820 'b' 

821 

822 The number of items requested can be larger than the number of items in 

823 the iterable: 

824 

825 >>> iterable = [1, 2, 3, 4, 5] 

826 >>> head, iterable = spy(iterable, 10) 

827 >>> head 

828 [1, 2, 3, 4, 5] 

829 >>> list(iterable) 

830 [1, 2, 3, 4, 5] 

831 

832 """ 

833 it = iter(iterable) 

834 head = take(n, it) 

835 

836 return head, chain(head, it) 

837 

838 

839def interleave(*iterables): 

840 """Return a new iterable yielding from each iterable in turn, 

841 until the shortest is exhausted. 

842 

843 >>> list(interleave([1, 2, 3], [4, 5], [6, 7, 8])) 

844 [1, 4, 6, 2, 5, 7] 

845 

846 For a version that doesn't terminate after the shortest iterable is 

847 exhausted, see :func:`interleave_longest`. 

848 

849 """ 

850 return chain.from_iterable(zip(*iterables)) 

851 

852 

853def interleave_longest(*iterables): 

854 """Return a new iterable yielding from each iterable in turn, 

855 skipping any that are exhausted. 

856 

857 >>> list(interleave_longest([1, 2, 3], [4, 5], [6, 7, 8])) 

858 [1, 4, 6, 2, 5, 7, 3, 8] 

859 

860 This function produces the same output as :func:`roundrobin`, but may 

861 perform better for some inputs (in particular when the number of iterables 

862 is large). 

863 

864 """ 

865 i = chain.from_iterable(zip_longest(*iterables, fillvalue=_marker)) 

866 return (x for x in i if x is not _marker) 

867 

868 

869def collapse(iterable, base_type=None, levels=None): 

870 """Flatten an iterable with multiple levels of nesting (e.g., a list of 

871 lists of tuples) into non-iterable types. 

872 

873 >>> iterable = [(1, 2), ([3, 4], [[5], [6]])] 

874 >>> list(collapse(iterable)) 

875 [1, 2, 3, 4, 5, 6] 

876 

877 String types are not considered iterable and will not be collapsed. 

878 To avoid collapsing other types, specify *base_type*: 

879 

880 >>> iterable = ['ab', ('cd', 'ef'), ['gh', 'ij']] 

881 >>> list(collapse(iterable, base_type=tuple)) 

882 ['ab', ('cd', 'ef'), 'gh', 'ij'] 

883 

884 Specify *levels* to stop flattening after a certain level: 

885 

886 >>> iterable = [('a', ['b']), ('c', ['d'])] 

887 >>> list(collapse(iterable)) # Fully flattened 

888 ['a', 'b', 'c', 'd'] 

889 >>> list(collapse(iterable, levels=1)) # Only one level flattened 

890 ['a', ['b'], 'c', ['d']] 

891 

892 """ 

893 def walk(node, level): 

894 if ( 

895 ((levels is not None) and (level > levels)) or 

896 isinstance(node, str) or 

897 ((base_type is not None) and isinstance(node, base_type)) 

898 ): 

899 yield node 

900 return 

901 

902 try: 

903 tree = iter(node) 

904 except TypeError: 

905 yield node 

906 return 

907 else: 

908 for child in tree: 

909 for x in walk(child, level + 1): 

910 yield x 

911 

912 yield from walk(iterable, 0) 

913 

914 

915def side_effect(func, iterable, chunk_size=None, before=None, after=None): 

916 """Invoke *func* on each item in *iterable* (or on each *chunk_size* group 

917 of items) before yielding the item. 

918 

919 `func` must be a function that takes a single argument. Its return value 

920 will be discarded. 

921 

922 *before* and *after* are optional functions that take no arguments. They 

923 will be executed before iteration starts and after it ends, respectively. 

924 

925 `side_effect` can be used for logging, updating progress bars, or anything 

926 that is not functionally "pure." 

927 

928 Emitting a status message: 

929 

930 >>> from more_itertools import consume 

931 >>> func = lambda item: print('Received {}'.format(item)) 

932 >>> consume(side_effect(func, range(2))) 

933 Received 0 

934 Received 1 

935 

936 Operating on chunks of items: 

937 

938 >>> pair_sums = [] 

939 >>> func = lambda chunk: pair_sums.append(sum(chunk)) 

940 >>> list(side_effect(func, [0, 1, 2, 3, 4, 5], 2)) 

941 [0, 1, 2, 3, 4, 5] 

942 >>> list(pair_sums) 

943 [1, 5, 9] 

944 

945 Writing to a file-like object: 

946 

947 >>> from io import StringIO 

948 >>> from more_itertools import consume 

949 >>> f = StringIO() 

950 >>> func = lambda x: print(x, file=f) 

951 >>> before = lambda: print(u'HEADER', file=f) 

952 >>> after = f.close 

953 >>> it = [u'a', u'b', u'c'] 

954 >>> consume(side_effect(func, it, before=before, after=after)) 

955 >>> f.closed 

956 True 

957 

958 """ 

959 try: 

960 if before is not None: 

961 before() 

962 

963 if chunk_size is None: 

964 for item in iterable: 

965 func(item) 

966 yield item 

967 else: 

968 for chunk in chunked(iterable, chunk_size): 

969 func(chunk) 

970 yield from chunk 

971 finally: 

972 if after is not None: 

973 after() 

974 

975 

976def sliced(seq, n): 

977 """Yield slices of length *n* from the sequence *seq*. 

978 

979 >>> list(sliced((1, 2, 3, 4, 5, 6), 3)) 

980 [(1, 2, 3), (4, 5, 6)] 

981 

982 If the length of the sequence is not divisible by the requested slice 

983 length, the last slice will be shorter. 

984 

985 >>> list(sliced((1, 2, 3, 4, 5, 6, 7, 8), 3)) 

986 [(1, 2, 3), (4, 5, 6), (7, 8)] 

987 

988 This function will only work for iterables that support slicing. 

989 For non-sliceable iterables, see :func:`chunked`. 

990 

991 """ 

992 return takewhile(bool, (seq[i: i + n] for i in count(0, n))) 

993 

994 

995def split_at(iterable, pred): 

996 """Yield lists of items from *iterable*, where each list is delimited by 

997 an item where callable *pred* returns ``True``. The lists do not include 

998 the delimiting items. 

999 

1000 >>> list(split_at('abcdcba', lambda x: x == 'b')) 

1001 [['a'], ['c', 'd', 'c'], ['a']] 

1002 

1003 >>> list(split_at(range(10), lambda n: n % 2 == 1)) 

1004 [[0], [2], [4], [6], [8], []] 

1005 """ 

1006 buf = [] 

1007 for item in iterable: 

1008 if pred(item): 

1009 yield buf 

1010 buf = [] 

1011 else: 

1012 buf.append(item) 

1013 yield buf 

1014 

1015 

1016def split_before(iterable, pred): 

1017 """Yield lists of items from *iterable*, where each list starts with an 

1018 item where callable *pred* returns ``True``: 

1019 

1020 >>> list(split_before('OneTwo', lambda s: s.isupper())) 

1021 [['O', 'n', 'e'], ['T', 'w', 'o']] 

1022 

1023 >>> list(split_before(range(10), lambda n: n % 3 == 0)) 

1024 [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]] 

1025 

1026 """ 

1027 buf = [] 

1028 for item in iterable: 

1029 if pred(item) and buf: 

1030 yield buf 

1031 buf = [] 

1032 buf.append(item) 

1033 yield buf 

1034 

1035 

1036def split_after(iterable, pred): 

1037 """Yield lists of items from *iterable*, where each list ends with an 

1038 item where callable *pred* returns ``True``: 

1039 

1040 >>> list(split_after('one1two2', lambda s: s.isdigit())) 

1041 [['o', 'n', 'e', '1'], ['t', 'w', 'o', '2']] 

1042 

1043 >>> list(split_after(range(10), lambda n: n % 3 == 0)) 

1044 [[0], [1, 2, 3], [4, 5, 6], [7, 8, 9]] 

1045 

1046 """ 

1047 buf = [] 

1048 for item in iterable: 

1049 buf.append(item) 

1050 if pred(item) and buf: 

1051 yield buf 

1052 buf = [] 

1053 if buf: 

1054 yield buf 

1055 

1056 

1057def split_into(iterable, sizes): 

1058 """Yield a list of sequential items from *iterable* of length 'n' for each 

1059 integer 'n' in *sizes*. 

1060 

1061 >>> list(split_into([1,2,3,4,5,6], [1,2,3])) 

1062 [[1], [2, 3], [4, 5, 6]] 

1063 

1064 If the sum of *sizes* is smaller than the length of *iterable*, then the 

1065 remaining items of *iterable* will not be returned. 

1066 

1067 >>> list(split_into([1,2,3,4,5,6], [2,3])) 

1068 [[1, 2], [3, 4, 5]] 

1069 

1070 If the sum of *sizes* is larger than the length of *iterable*, fewer items 

1071 will be returned in the iteration that overruns *iterable* and further 

1072 lists will be empty: 

1073 

1074 >>> list(split_into([1,2,3,4], [1,2,3,4])) 

1075 [[1], [2, 3], [4], []] 

1076 

1077 When a ``None`` object is encountered in *sizes*, the returned list will 

1078 contain items up to the end of *iterable* the same way that itertools.slice 

1079 does: 

1080 

1081 >>> list(split_into([1,2,3,4,5,6,7,8,9,0], [2,3,None])) 

1082 [[1, 2], [3, 4, 5], [6, 7, 8, 9, 0]] 

1083 

1084 :func:`split_into` can be useful for grouping a series of items where the 

1085 sizes of the groups are not uniform. An example would be where in a row 

1086 from a table, multiple columns represent elements of the same feature 

1087 (e.g. a point represented by x,y,z) but, the format is not the same for 

1088 all columns. 

1089 """ 

1090 # convert the iterable argument into an iterator so its contents can 

1091 # be consumed by islice in case it is a generator 

1092 it = iter(iterable) 

1093 

1094 for size in sizes: 

1095 if size is None: 

1096 yield list(it) 

1097 return 

1098 else: 

1099 yield list(islice(it, size)) 

1100 

1101 

1102def padded(iterable, fillvalue=None, n=None, next_multiple=False): 

1103 """Yield the elements from *iterable*, followed by *fillvalue*, such that 

1104 at least *n* items are emitted. 

1105 

1106 >>> list(padded([1, 2, 3], '?', 5)) 

1107 [1, 2, 3, '?', '?'] 

1108 

1109 If *next_multiple* is ``True``, *fillvalue* will be emitted until the 

1110 number of items emitted is a multiple of *n*:: 

1111 

1112 >>> list(padded([1, 2, 3, 4], n=3, next_multiple=True)) 

1113 [1, 2, 3, 4, None, None] 

1114 

1115 If *n* is ``None``, *fillvalue* will be emitted indefinitely. 

1116 

1117 """ 

1118 it = iter(iterable) 

1119 if n is None: 

1120 yield from chain(it, repeat(fillvalue)) 

1121 elif n < 1: 

1122 raise ValueError('n must be at least 1') 

1123 else: 

1124 item_count = 0 

1125 for item in it: 

1126 yield item 

1127 item_count += 1 

1128 

1129 remaining = (n - item_count) % n if next_multiple else n - item_count 

1130 for _ in range(remaining): 

1131 yield fillvalue 

1132 

1133 

1134def distribute(n, iterable): 

1135 """Distribute the items from *iterable* among *n* smaller iterables. 

1136 

1137 >>> group_1, group_2 = distribute(2, [1, 2, 3, 4, 5, 6]) 

1138 >>> list(group_1) 

1139 [1, 3, 5] 

1140 >>> list(group_2) 

1141 [2, 4, 6] 

1142 

1143 If the length of *iterable* is not evenly divisible by *n*, then the 

1144 length of the returned iterables will not be identical: 

1145 

1146 >>> children = distribute(3, [1, 2, 3, 4, 5, 6, 7]) 

1147 >>> [list(c) for c in children] 

1148 [[1, 4, 7], [2, 5], [3, 6]] 

1149 

1150 If the length of *iterable* is smaller than *n*, then the last returned 

1151 iterables will be empty: 

1152 

1153 >>> children = distribute(5, [1, 2, 3]) 

1154 >>> [list(c) for c in children] 

1155 [[1], [2], [3], [], []] 

1156 

1157 This function uses :func:`itertools.tee` and may require significant 

1158 storage. If you need the order items in the smaller iterables to match the 

1159 original iterable, see :func:`divide`. 

1160 

1161 """ 

1162 if n < 1: 

1163 raise ValueError('n must be at least 1') 

1164 

1165 children = tee(iterable, n) 

1166 return [islice(it, index, None, n) for index, it in enumerate(children)] 

1167 

1168 

1169def stagger(iterable, offsets=(-1, 0, 1), longest=False, fillvalue=None): 

1170 """Yield tuples whose elements are offset from *iterable*. 

1171 The amount by which the `i`-th item in each tuple is offset is given by 

1172 the `i`-th item in *offsets*. 

1173 

1174 >>> list(stagger([0, 1, 2, 3])) 

1175 [(None, 0, 1), (0, 1, 2), (1, 2, 3)] 

1176 >>> list(stagger(range(8), offsets=(0, 2, 4))) 

1177 [(0, 2, 4), (1, 3, 5), (2, 4, 6), (3, 5, 7)] 

1178 

1179 By default, the sequence will end when the final element of a tuple is the 

1180 last item in the iterable. To continue until the first element of a tuple 

1181 is the last item in the iterable, set *longest* to ``True``:: 

1182 

1183 >>> list(stagger([0, 1, 2, 3], longest=True)) 

1184 [(None, 0, 1), (0, 1, 2), (1, 2, 3), (2, 3, None), (3, None, None)] 

1185 

1186 By default, ``None`` will be used to replace offsets beyond the end of the 

1187 sequence. Specify *fillvalue* to use some other value. 

1188 

1189 """ 

1190 children = tee(iterable, len(offsets)) 

1191 

1192 return zip_offset( 

1193 *children, offsets=offsets, longest=longest, fillvalue=fillvalue 

1194 ) 

1195 

1196 

1197def zip_offset(*iterables, offsets, longest=False, fillvalue=None): 

1198 """``zip`` the input *iterables* together, but offset the `i`-th iterable 

1199 by the `i`-th item in *offsets*. 

1200 

1201 >>> list(zip_offset('0123', 'abcdef', offsets=(0, 1))) 

1202 [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e')] 

1203 

1204 This can be used as a lightweight alternative to SciPy or pandas to analyze 

1205 data sets in which some series have a lead or lag relationship. 

1206 

1207 By default, the sequence will end when the shortest iterable is exhausted. 

1208 To continue until the longest iterable is exhausted, set *longest* to 

1209 ``True``. 

1210 

1211 >>> list(zip_offset('0123', 'abcdef', offsets=(0, 1), longest=True)) 

1212 [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e'), (None, 'f')] 

1213 

1214 By default, ``None`` will be used to replace offsets beyond the end of the 

1215 sequence. Specify *fillvalue* to use some other value. 

1216 

1217 """ 

1218 if len(iterables) != len(offsets): 

1219 raise ValueError("Number of iterables and offsets didn't match") 

1220 

1221 staggered = [] 

1222 for it, n in zip(iterables, offsets): 

1223 if n < 0: 

1224 staggered.append(chain(repeat(fillvalue, -n), it)) 

1225 elif n > 0: 

1226 staggered.append(islice(it, n, None)) 

1227 else: 

1228 staggered.append(it) 

1229 

1230 if longest: 

1231 return zip_longest(*staggered, fillvalue=fillvalue) 

1232 

1233 return zip(*staggered) 

1234 

1235 

1236def sort_together(iterables, key_list=(0,), reverse=False): 

1237 """Return the input iterables sorted together, with *key_list* as the 

1238 priority for sorting. All iterables are trimmed to the length of the 

1239 shortest one. 

1240 

1241 This can be used like the sorting function in a spreadsheet. If each 

1242 iterable represents a column of data, the key list determines which 

1243 columns are used for sorting. 

1244 

1245 By default, all iterables are sorted using the ``0``-th iterable:: 

1246 

1247 >>> iterables = [(4, 3, 2, 1), ('a', 'b', 'c', 'd')] 

1248 >>> sort_together(iterables) 

1249 [(1, 2, 3, 4), ('d', 'c', 'b', 'a')] 

1250 

1251 Set a different key list to sort according to another iterable. 

1252 Specifying multiple keys dictates how ties are broken:: 

1253 

1254 >>> iterables = [(3, 1, 2), (0, 1, 0), ('c', 'b', 'a')] 

1255 >>> sort_together(iterables, key_list=(1, 2)) 

1256 [(2, 3, 1), (0, 0, 1), ('a', 'c', 'b')] 

1257 

1258 Set *reverse* to ``True`` to sort in descending order. 

1259 

1260 >>> sort_together([(1, 2, 3), ('c', 'b', 'a')], reverse=True) 

1261 [(3, 2, 1), ('a', 'b', 'c')] 

1262 

1263 """ 

1264 return list(zip(*sorted(zip(*iterables), 

1265 key=itemgetter(*key_list), 

1266 reverse=reverse))) 

1267 

1268 

1269def unzip(iterable): 

1270 """The inverse of :func:`zip`, this function disaggregates the elements 

1271 of the zipped *iterable*. 

1272 

1273 The ``i``-th iterable contains the ``i``-th element from each element 

1274 of the zipped iterable. The first element is used to to determine the 

1275 length of the remaining elements. 

1276 

1277 >>> iterable = [('a', 1), ('b', 2), ('c', 3), ('d', 4)] 

1278 >>> letters, numbers = unzip(iterable) 

1279 >>> list(letters) 

1280 ['a', 'b', 'c', 'd'] 

1281 >>> list(numbers) 

1282 [1, 2, 3, 4] 

1283 

1284 This is similar to using ``zip(*iterable)``, but it avoids reading 

1285 *iterable* into memory. Note, however, that this function uses 

1286 :func:`itertools.tee` and thus may require significant storage. 

1287 

1288 """ 

1289 head, iterable = spy(iter(iterable)) 

1290 if not head: 

1291 # empty iterable, e.g. zip([], [], []) 

1292 return () 

1293 # spy returns a one-length iterable as head 

1294 head = head[0] 

1295 iterables = tee(iterable, len(head)) 

1296 

1297 def itemgetter(i): 

1298 def getter(obj): 

1299 try: 

1300 return obj[i] 

1301 except IndexError: 

1302 # basically if we have an iterable like 

1303 # iter([(1, 2, 3), (4, 5), (6,)]) 

1304 # the second unzipped iterable would fail at the third tuple 

1305 # since it would try to access tup[1] 

1306 # same with the third unzipped iterable and the second tuple 

1307 # to support these "improperly zipped" iterables, 

1308 # we create a custom itemgetter 

1309 # which just stops the unzipped iterables 

1310 # at first length mismatch 

1311 raise StopIteration 

1312 return getter 

1313 

1314 return tuple(map(itemgetter(i), it) for i, it in enumerate(iterables)) 

1315 

1316 

1317def divide(n, iterable): 

1318 """Divide the elements from *iterable* into *n* parts, maintaining 

1319 order. 

1320 

1321 >>> group_1, group_2 = divide(2, [1, 2, 3, 4, 5, 6]) 

1322 >>> list(group_1) 

1323 [1, 2, 3] 

1324 >>> list(group_2) 

1325 [4, 5, 6] 

1326 

1327 If the length of *iterable* is not evenly divisible by *n*, then the 

1328 length of the returned iterables will not be identical: 

1329 

1330 >>> children = divide(3, [1, 2, 3, 4, 5, 6, 7]) 

1331 >>> [list(c) for c in children] 

1332 [[1, 2, 3], [4, 5], [6, 7]] 

1333 

1334 If the length of the iterable is smaller than n, then the last returned 

1335 iterables will be empty: 

1336 

1337 >>> children = divide(5, [1, 2, 3]) 

1338 >>> [list(c) for c in children] 

1339 [[1], [2], [3], [], []] 

1340 

1341 This function will exhaust the iterable before returning and may require 

1342 significant storage. If order is not important, see :func:`distribute`, 

1343 which does not first pull the iterable into memory. 

1344 

1345 """ 

1346 if n < 1: 

1347 raise ValueError('n must be at least 1') 

1348 

1349 seq = tuple(iterable) 

1350 q, r = divmod(len(seq), n) 

1351 

1352 ret = [] 

1353 for i in range(n): 

1354 start = (i * q) + (i if i < r else r) 

1355 stop = ((i + 1) * q) + (i + 1 if i + 1 < r else r) 

1356 ret.append(iter(seq[start:stop])) 

1357 

1358 return ret 

1359 

1360 

1361def always_iterable(obj, base_type=(str, bytes)): 

1362 """If *obj* is iterable, return an iterator over its items:: 

1363 

1364 >>> obj = (1, 2, 3) 

1365 >>> list(always_iterable(obj)) 

1366 [1, 2, 3] 

1367 

1368 If *obj* is not iterable, return a one-item iterable containing *obj*:: 

1369 

1370 >>> obj = 1 

1371 >>> list(always_iterable(obj)) 

1372 [1] 

1373 

1374 If *obj* is ``None``, return an empty iterable: 

1375 

1376 >>> obj = None 

1377 >>> list(always_iterable(None)) 

1378 [] 

1379 

1380 By default, binary and text strings are not considered iterable:: 

1381 

1382 >>> obj = 'foo' 

1383 >>> list(always_iterable(obj)) 

1384 ['foo'] 

1385 

1386 If *base_type* is set, objects for which ``isinstance(obj, base_type)`` 

1387 returns ``True`` won't be considered iterable. 

1388 

1389 >>> obj = {'a': 1} 

1390 >>> list(always_iterable(obj)) # Iterate over the dict's keys 

1391 ['a'] 

1392 >>> list(always_iterable(obj, base_type=dict)) # Treat dicts as a unit 

1393 [{'a': 1}] 

1394 

1395 Set *base_type* to ``None`` to avoid any special handling and treat objects 

1396 Python considers iterable as iterable: 

1397 

1398 >>> obj = 'foo' 

1399 >>> list(always_iterable(obj, base_type=None)) 

1400 ['f', 'o', 'o'] 

1401 """ 

1402 if obj is None: 

1403 return iter(()) 

1404 

1405 if (base_type is not None) and isinstance(obj, base_type): 

1406 return iter((obj,)) 

1407 

1408 try: 

1409 return iter(obj) 

1410 except TypeError: 

1411 return iter((obj,)) 

1412 

1413 

1414def adjacent(predicate, iterable, distance=1): 

1415 """Return an iterable over `(bool, item)` tuples where the `item` is 

1416 drawn from *iterable* and the `bool` indicates whether 

1417 that item satisfies the *predicate* or is adjacent to an item that does. 

1418 

1419 For example, to find whether items are adjacent to a ``3``:: 

1420 

1421 >>> list(adjacent(lambda x: x == 3, range(6))) 

1422 [(False, 0), (False, 1), (True, 2), (True, 3), (True, 4), (False, 5)] 

1423 

1424 Set *distance* to change what counts as adjacent. For example, to find 

1425 whether items are two places away from a ``3``: 

1426 

1427 >>> list(adjacent(lambda x: x == 3, range(6), distance=2)) 

1428 [(False, 0), (True, 1), (True, 2), (True, 3), (True, 4), (True, 5)] 

1429 

1430 This is useful for contextualizing the results of a search function. 

1431 For example, a code comparison tool might want to identify lines that 

1432 have changed, but also surrounding lines to give the viewer of the diff 

1433 context. 

1434 

1435 The predicate function will only be called once for each item in the 

1436 iterable. 

1437 

1438 See also :func:`groupby_transform`, which can be used with this function 

1439 to group ranges of items with the same `bool` value. 

1440 

1441 """ 

1442 # Allow distance=0 mainly for testing that it reproduces results with map() 

1443 if distance < 0: 

1444 raise ValueError('distance must be at least 0') 

1445 

1446 i1, i2 = tee(iterable) 

1447 padding = [False] * distance 

1448 selected = chain(padding, map(predicate, i1), padding) 

1449 adjacent_to_selected = map(any, windowed(selected, 2 * distance + 1)) 

1450 return zip(adjacent_to_selected, i2) 

1451 

1452 

1453def groupby_transform(iterable, keyfunc=None, valuefunc=None): 

1454 """An extension of :func:`itertools.groupby` that transforms the values of 

1455 *iterable* after grouping them. 

1456 *keyfunc* is a function used to compute a grouping key for each item. 

1457 *valuefunc* is a function for transforming the items after grouping. 

1458 

1459 >>> iterable = 'AaaABbBCcA' 

1460 >>> keyfunc = lambda x: x.upper() 

1461 >>> valuefunc = lambda x: x.lower() 

1462 >>> grouper = groupby_transform(iterable, keyfunc, valuefunc) 

1463 >>> [(k, ''.join(g)) for k, g in grouper] 

1464 [('A', 'aaaa'), ('B', 'bbb'), ('C', 'cc'), ('A', 'a')] 

1465 

1466 *keyfunc* and *valuefunc* default to identity functions if they are not 

1467 specified. 

1468 

1469 :func:`groupby_transform` is useful when grouping elements of an iterable 

1470 using a separate iterable as the key. To do this, :func:`zip` the iterables 

1471 and pass a *keyfunc* that extracts the first element and a *valuefunc* 

1472 that extracts the second element:: 

1473 

1474 >>> from operator import itemgetter 

1475 >>> keys = [0, 0, 1, 1, 1, 2, 2, 2, 3] 

1476 >>> values = 'abcdefghi' 

1477 >>> iterable = zip(keys, values) 

1478 >>> grouper = groupby_transform(iterable, itemgetter(0), itemgetter(1)) 

1479 >>> [(k, ''.join(g)) for k, g in grouper] 

1480 [(0, 'ab'), (1, 'cde'), (2, 'fgh'), (3, 'i')] 

1481 

1482 Note that the order of items in the iterable is significant. 

1483 Only adjacent items are grouped together, so if you don't want any 

1484 duplicate groups, you should sort the iterable by the key function. 

1485 

1486 """ 

1487 valuefunc = (lambda x: x) if valuefunc is None else valuefunc 

1488 return ((k, map(valuefunc, g)) for k, g in groupby(iterable, keyfunc)) 

1489 

1490 

1491def numeric_range(*args): 

1492 """An extension of the built-in ``range()`` function whose arguments can 

1493 be any orderable numeric type. 

1494 

1495 With only *stop* specified, *start* defaults to ``0`` and *step* 

1496 defaults to ``1``. The output items will match the type of *stop*: 

1497 

1498 >>> list(numeric_range(3.5)) 

1499 [0.0, 1.0, 2.0, 3.0] 

1500 

1501 With only *start* and *stop* specified, *step* defaults to ``1``. The 

1502 output items will match the type of *start*: 

1503 

1504 >>> from decimal import Decimal 

1505 >>> start = Decimal('2.1') 

1506 >>> stop = Decimal('5.1') 

1507 >>> list(numeric_range(start, stop)) 

1508 [Decimal('2.1'), Decimal('3.1'), Decimal('4.1')] 

1509 

1510 With *start*, *stop*, and *step* specified the output items will match 

1511 the type of ``start + step``: 

1512 

1513 >>> from fractions import Fraction 

1514 >>> start = Fraction(1, 2) # Start at 1/2 

1515 >>> stop = Fraction(5, 2) # End at 5/2 

1516 >>> step = Fraction(1, 2) # Count by 1/2 

1517 >>> list(numeric_range(start, stop, step)) 

1518 [Fraction(1, 2), Fraction(1, 1), Fraction(3, 2), Fraction(2, 1)] 

1519 

1520 If *step* is zero, ``ValueError`` is raised. Negative steps are supported: 

1521 

1522 >>> list(numeric_range(3, -1, -1.0)) 

1523 [3.0, 2.0, 1.0, 0.0] 

1524 

1525 Be aware of the limitations of floating point numbers; the representation 

1526 of the yielded numbers may be surprising. 

1527 

1528 """ 

1529 argc = len(args) 

1530 if argc == 1: 

1531 stop, = args 

1532 start = type(stop)(0) 

1533 step = 1 

1534 elif argc == 2: 

1535 start, stop = args 

1536 step = 1 

1537 elif argc == 3: 

1538 start, stop, step = args 

1539 else: 

1540 err_msg = 'numeric_range takes at most 3 arguments, got {}' 

1541 raise TypeError(err_msg.format(argc)) 

1542 

1543 values = (start + (step * n) for n in count()) 

1544 if step > 0: 

1545 return takewhile(partial(gt, stop), values) 

1546 elif step < 0: 

1547 return takewhile(partial(lt, stop), values) 

1548 else: 

1549 raise ValueError('numeric_range arg 3 must not be zero') 

1550 

1551 

1552def count_cycle(iterable, n=None): 

1553 """Cycle through the items from *iterable* up to *n* times, yielding 

1554 the number of completed cycles along with each item. If *n* is omitted the 

1555 process repeats indefinitely. 

1556 

1557 >>> list(count_cycle('AB', 3)) 

1558 [(0, 'A'), (0, 'B'), (1, 'A'), (1, 'B'), (2, 'A'), (2, 'B')] 

1559 

1560 """ 

1561 iterable = tuple(iterable) 

1562 if not iterable: 

1563 return iter(()) 

1564 counter = count() if n is None else range(n) 

1565 return ((i, item) for i in counter for item in iterable) 

1566 

1567 

1568def locate(iterable, pred=bool, window_size=None): 

1569 """Yield the index of each item in *iterable* for which *pred* returns 

1570 ``True``. 

1571 

1572 *pred* defaults to :func:`bool`, which will select truthy items: 

1573 

1574 >>> list(locate([0, 1, 1, 0, 1, 0, 0])) 

1575 [1, 2, 4] 

1576 

1577 Set *pred* to a custom function to, e.g., find the indexes for a particular 

1578 item. 

1579 

1580 >>> list(locate(['a', 'b', 'c', 'b'], lambda x: x == 'b')) 

1581 [1, 3] 

1582 

1583 If *window_size* is given, then the *pred* function will be called with 

1584 that many items. This enables searching for sub-sequences: 

1585 

1586 >>> iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3] 

1587 >>> pred = lambda *args: args == (1, 2, 3) 

1588 >>> list(locate(iterable, pred=pred, window_size=3)) 

1589 [1, 5, 9] 

1590 

1591 Use with :func:`seekable` to find indexes and then retrieve the associated 

1592 items: 

1593 

1594 >>> from itertools import count 

1595 >>> from more_itertools import seekable 

1596 >>> source = (3 * n + 1 if (n % 2) else n // 2 for n in count()) 

1597 >>> it = seekable(source) 

1598 >>> pred = lambda x: x > 100 

1599 >>> indexes = locate(it, pred=pred) 

1600 >>> i = next(indexes) 

1601 >>> it.seek(i) 

1602 >>> next(it) 

1603 106 

1604 

1605 """ 

1606 if window_size is None: 

1607 return compress(count(), map(pred, iterable)) 

1608 

1609 if window_size < 1: 

1610 raise ValueError('window size must be at least 1') 

1611 

1612 it = windowed(iterable, window_size, fillvalue=_marker) 

1613 return compress(count(), starmap(pred, it)) 

1614 

1615 

1616def lstrip(iterable, pred): 

1617 """Yield the items from *iterable*, but strip any from the beginning 

1618 for which *pred* returns ``True``. 

1619 

1620 For example, to remove a set of items from the start of an iterable: 

1621 

1622 >>> iterable = (None, False, None, 1, 2, None, 3, False, None) 

1623 >>> pred = lambda x: x in {None, False, ''} 

1624 >>> list(lstrip(iterable, pred)) 

1625 [1, 2, None, 3, False, None] 

1626 

1627 This function is analogous to to :func:`str.lstrip`, and is essentially 

1628 an wrapper for :func:`itertools.dropwhile`. 

1629 

1630 """ 

1631 return dropwhile(pred, iterable) 

1632 

1633 

1634def rstrip(iterable, pred): 

1635 """Yield the items from *iterable*, but strip any from the end 

1636 for which *pred* returns ``True``. 

1637 

1638 For example, to remove a set of items from the end of an iterable: 

1639 

1640 >>> iterable = (None, False, None, 1, 2, None, 3, False, None) 

1641 >>> pred = lambda x: x in {None, False, ''} 

1642 >>> list(rstrip(iterable, pred)) 

1643 [None, False, None, 1, 2, None, 3] 

1644 

1645 This function is analogous to :func:`str.rstrip`. 

1646 

1647 """ 

1648 cache = [] 

1649 cache_append = cache.append 

1650 cache_clear = cache.clear 

1651 for x in iterable: 

1652 if pred(x): 

1653 cache_append(x) 

1654 else: 

1655 yield from cache 

1656 cache_clear() 

1657 yield x 

1658 

1659 

1660def strip(iterable, pred): 

1661 """Yield the items from *iterable*, but strip any from the 

1662 beginning and end for which *pred* returns ``True``. 

1663 

1664 For example, to remove a set of items from both ends of an iterable: 

1665 

1666 >>> iterable = (None, False, None, 1, 2, None, 3, False, None) 

1667 >>> pred = lambda x: x in {None, False, ''} 

1668 >>> list(strip(iterable, pred)) 

1669 [1, 2, None, 3] 

1670 

1671 This function is analogous to :func:`str.strip`. 

1672 

1673 """ 

1674 return rstrip(lstrip(iterable, pred), pred) 

1675 

1676 

1677def islice_extended(iterable, *args): 

1678 """An extension of :func:`itertools.islice` that supports negative values 

1679 for *stop*, *start*, and *step*. 

1680 

1681 >>> iterable = iter('abcdefgh') 

1682 >>> list(islice_extended(iterable, -4, -1)) 

1683 ['e', 'f', 'g'] 

1684 

1685 Slices with negative values require some caching of *iterable*, but this 

1686 function takes care to minimize the amount of memory required. 

1687 

1688 For example, you can use a negative step with an infinite iterator: 

1689 

1690 >>> from itertools import count 

1691 >>> list(islice_extended(count(), 110, 99, -2)) 

1692 [110, 108, 106, 104, 102, 100] 

1693 

1694 """ 

1695 s = slice(*args) 

1696 start = s.start 

1697 stop = s.stop 

1698 if s.step == 0: 

1699 raise ValueError('step argument must be a non-zero integer or None.') 

1700 step = s.step or 1 

1701 

1702 it = iter(iterable) 

1703 

1704 if step > 0: 

1705 start = 0 if (start is None) else start 

1706 

1707 if (start < 0): 

1708 # Consume all but the last -start items 

1709 cache = deque(enumerate(it, 1), maxlen=-start) 

1710 len_iter = cache[-1][0] if cache else 0 

1711 

1712 # Adjust start to be positive 

1713 i = max(len_iter + start, 0) 

1714 

1715 # Adjust stop to be positive 

1716 if stop is None: 

1717 j = len_iter 

1718 elif stop >= 0: 

1719 j = min(stop, len_iter) 

1720 else: 

1721 j = max(len_iter + stop, 0) 

1722 

1723 # Slice the cache 

1724 n = j - i 

1725 if n <= 0: 

1726 return 

1727 

1728 for index, item in islice(cache, 0, n, step): 

1729 yield item 

1730 elif (stop is not None) and (stop < 0): 

1731 # Advance to the start position 

1732 next(islice(it, start, start), None) 

1733 

1734 # When stop is negative, we have to carry -stop items while 

1735 # iterating 

1736 cache = deque(islice(it, -stop), maxlen=-stop) 

1737 

1738 for index, item in enumerate(it): 

1739 cached_item = cache.popleft() 

1740 if index % step == 0: 

1741 yield cached_item 

1742 cache.append(item) 

1743 else: 

1744 # When both start and stop are positive we have the normal case 

1745 yield from islice(it, start, stop, step) 

1746 else: 

1747 start = -1 if (start is None) else start 

1748 

1749 if (stop is not None) and (stop < 0): 

1750 # Consume all but the last items 

1751 n = -stop - 1 

1752 cache = deque(enumerate(it, 1), maxlen=n) 

1753 len_iter = cache[-1][0] if cache else 0 

1754 

1755 # If start and stop are both negative they are comparable and 

1756 # we can just slice. Otherwise we can adjust start to be negative 

1757 # and then slice. 

1758 if start < 0: 

1759 i, j = start, stop 

1760 else: 

1761 i, j = min(start - len_iter, -1), None 

1762 

1763 for index, item in list(cache)[i:j:step]: 

1764 yield item 

1765 else: 

1766 # Advance to the stop position 

1767 if stop is not None: 

1768 m = stop + 1 

1769 next(islice(it, m, m), None) 

1770 

1771 # stop is positive, so if start is negative they are not comparable 

1772 # and we need the rest of the items. 

1773 if start < 0: 

1774 i = start 

1775 n = None 

1776 # stop is None and start is positive, so we just need items up to 

1777 # the start index. 

1778 elif stop is None: 

1779 i = None 

1780 n = start + 1 

1781 # Both stop and start are positive, so they are comparable. 

1782 else: 

1783 i = None 

1784 n = start - stop 

1785 if n <= 0: 

1786 return 

1787 

1788 cache = list(islice(it, n)) 

1789 

1790 yield from cache[i::step] 

1791 

1792 

1793def always_reversible(iterable): 

1794 """An extension of :func:`reversed` that supports all iterables, not 

1795 just those which implement the ``Reversible`` or ``Sequence`` protocols. 

1796 

1797 >>> print(*always_reversible(x for x in range(3))) 

1798 2 1 0 

1799 

1800 If the iterable is already reversible, this function returns the 

1801 result of :func:`reversed()`. If the iterable is not reversible, 

1802 this function will cache the remaining items in the iterable and 

1803 yield them in reverse order, which may require significant storage. 

1804 """ 

1805 try: 

1806 return reversed(iterable) 

1807 except TypeError: 

1808 return reversed(list(iterable)) 

1809 

1810 

1811def consecutive_groups(iterable, ordering=lambda x: x): 

1812 """Yield groups of consecutive items using :func:`itertools.groupby`. 

1813 The *ordering* function determines whether two items are adjacent by 

1814 returning their position. 

1815 

1816 By default, the ordering function is the identity function. This is 

1817 suitable for finding runs of numbers: 

1818 

1819 >>> iterable = [1, 10, 11, 12, 20, 30, 31, 32, 33, 40] 

1820 >>> for group in consecutive_groups(iterable): 

1821 ... print(list(group)) 

1822 [1] 

1823 [10, 11, 12] 

1824 [20] 

1825 [30, 31, 32, 33] 

1826 [40] 

1827 

1828 For finding runs of adjacent letters, try using the :meth:`index` method 

1829 of a string of letters: 

1830 

1831 >>> from string import ascii_lowercase 

1832 >>> iterable = 'abcdfgilmnop' 

1833 >>> ordering = ascii_lowercase.index 

1834 >>> for group in consecutive_groups(iterable, ordering): 

1835 ... print(list(group)) 

1836 ['a', 'b', 'c', 'd'] 

1837 ['f', 'g'] 

1838 ['i'] 

1839 ['l', 'm', 'n', 'o', 'p'] 

1840 

1841 """ 

1842 for k, g in groupby( 

1843 enumerate(iterable), key=lambda x: x[0] - ordering(x[1]) 

1844 ): 

1845 yield map(itemgetter(1), g) 

1846 

1847 

1848def difference(iterable, func=sub): 

1849 """By default, compute the first difference of *iterable* using 

1850 :func:`operator.sub`. 

1851 

1852 >>> iterable = [0, 1, 3, 6, 10] 

1853 >>> list(difference(iterable)) 

1854 [0, 1, 2, 3, 4] 

1855 

1856 This is the opposite of :func:`accumulate`'s default behavior: 

1857 

1858 >>> from itertools import accumulate 

1859 >>> iterable = [0, 1, 2, 3, 4] 

1860 >>> list(accumulate(iterable)) 

1861 [0, 1, 3, 6, 10] 

1862 >>> list(difference(accumulate(iterable))) 

1863 [0, 1, 2, 3, 4] 

1864 

1865 By default *func* is :func:`operator.sub`, but other functions can be 

1866 specified. They will be applied as follows:: 

1867 

1868 A, B, C, D, ... --> A, func(B, A), func(C, B), func(D, C), ... 

1869 

1870 For example, to do progressive division: 

1871 

1872 >>> iterable = [1, 2, 6, 24, 120] # Factorial sequence 

1873 >>> func = lambda x, y: x // y 

1874 >>> list(difference(iterable, func)) 

1875 [1, 2, 3, 4, 5] 

1876 

1877 """ 

1878 a, b = tee(iterable) 

1879 try: 

1880 item = next(b) 

1881 except StopIteration: 

1882 return iter([]) 

1883 return chain([item], map(lambda x: func(x[1], x[0]), zip(a, b))) 

1884 

1885 

1886class SequenceView(Sequence): 

1887 """Return a read-only view of the sequence object *target*. 

1888 

1889 :class:`SequenceView` objects are analogous to Python's built-in 

1890 "dictionary view" types. They provide a dynamic view of a sequence's items, 

1891 meaning that when the sequence updates, so does the view. 

1892 

1893 >>> seq = ['0', '1', '2'] 

1894 >>> view = SequenceView(seq) 

1895 >>> view 

1896 SequenceView(['0', '1', '2']) 

1897 >>> seq.append('3') 

1898 >>> view 

1899 SequenceView(['0', '1', '2', '3']) 

1900 

1901 Sequence views support indexing, slicing, and length queries. They act 

1902 like the underlying sequence, except they don't allow assignment: 

1903 

1904 >>> view[1] 

1905 '1' 

1906 >>> view[1:-1] 

1907 ['1', '2'] 

1908 >>> len(view) 

1909 4 

1910 

1911 Sequence views are useful as an alternative to copying, as they don't 

1912 require (much) extra storage. 

1913 

1914 """ 

1915 def __init__(self, target): 

1916 if not isinstance(target, Sequence): 

1917 raise TypeError 

1918 self._target = target 

1919 

1920 def __getitem__(self, index): 

1921 return self._target[index] 

1922 

1923 def __len__(self): 

1924 return len(self._target) 

1925 

1926 def __repr__(self): 

1927 return '{}({})'.format(self.__class__.__name__, repr(self._target)) 

1928 

1929 

1930class seekable: 

1931 """Wrap an iterator to allow for seeking backward and forward. This 

1932 progressively caches the items in the source iterable so they can be 

1933 re-visited. 

1934 

1935 Call :meth:`seek` with an index to seek to that position in the source 

1936 iterable. 

1937 

1938 To "reset" an iterator, seek to ``0``: 

1939 

1940 >>> from itertools import count 

1941 >>> it = seekable((str(n) for n in count())) 

1942 >>> next(it), next(it), next(it) 

1943 ('0', '1', '2') 

1944 >>> it.seek(0) 

1945 >>> next(it), next(it), next(it) 

1946 ('0', '1', '2') 

1947 >>> next(it) 

1948 '3' 

1949 

1950 You can also seek forward: 

1951 

1952 >>> it = seekable((str(n) for n in range(20))) 

1953 >>> it.seek(10) 

1954 >>> next(it) 

1955 '10' 

1956 >>> it.seek(20) # Seeking past the end of the source isn't a problem 

1957 >>> list(it) 

1958 [] 

1959 >>> it.seek(0) # Resetting works even after hitting the end 

1960 >>> next(it), next(it), next(it) 

1961 ('0', '1', '2') 

1962 

1963 The cache grows as the source iterable progresses, so beware of wrapping 

1964 very large or infinite iterables. 

1965 

1966 You may view the contents of the cache with the :meth:`elements` method. 

1967 That returns a :class:`SequenceView`, a view that updates automatically: 

1968 

1969 >>> it = seekable((str(n) for n in range(10))) 

1970 >>> next(it), next(it), next(it) 

1971 ('0', '1', '2') 

1972 >>> elements = it.elements() 

1973 >>> elements 

1974 SequenceView(['0', '1', '2']) 

1975 >>> next(it) 

1976 '3' 

1977 >>> elements 

1978 SequenceView(['0', '1', '2', '3']) 

1979 

1980 """ 

1981 

1982 def __init__(self, iterable): 

1983 self._source = iter(iterable) 

1984 self._cache = [] 

1985 self._index = None 

1986 

1987 def __iter__(self): 

1988 return self 

1989 

1990 def __next__(self): 

1991 if self._index is not None: 

1992 try: 

1993 item = self._cache[self._index] 

1994 except IndexError: 

1995 self._index = None 

1996 else: 

1997 self._index += 1 

1998 return item 

1999 

2000 item = next(self._source) 

2001 self._cache.append(item) 

2002 return item 

2003 

2004 def elements(self): 

2005 return SequenceView(self._cache) 

2006 

2007 def seek(self, index): 

2008 self._index = index 

2009 remainder = index - len(self._cache) 

2010 if remainder > 0: 

2011 consume(self, remainder) 

2012 

2013 

2014class run_length: 

2015 """ 

2016 :func:`run_length.encode` compresses an iterable with run-length encoding. 

2017 It yields groups of repeated items with the count of how many times they 

2018 were repeated: 

2019 

2020 >>> uncompressed = 'abbcccdddd' 

2021 >>> list(run_length.encode(uncompressed)) 

2022 [('a', 1), ('b', 2), ('c', 3), ('d', 4)] 

2023 

2024 :func:`run_length.decode` decompresses an iterable that was previously 

2025 compressed with run-length encoding. It yields the items of the 

2026 decompressed iterable: 

2027 

2028 >>> compressed = [('a', 1), ('b', 2), ('c', 3), ('d', 4)] 

2029 >>> list(run_length.decode(compressed)) 

2030 ['a', 'b', 'b', 'c', 'c', 'c', 'd', 'd', 'd', 'd'] 

2031 

2032 """ 

2033 

2034 @staticmethod 

2035 def encode(iterable): 

2036 return ((k, ilen(g)) for k, g in groupby(iterable)) 

2037 

2038 @staticmethod 

2039 def decode(iterable): 

2040 return chain.from_iterable(repeat(k, n) for k, n in iterable) 

2041 

2042 

2043def exactly_n(iterable, n, predicate=bool): 

2044 """Return ``True`` if exactly ``n`` items in the iterable are ``True`` 

2045 according to the *predicate* function. 

2046 

2047 >>> exactly_n([True, True, False], 2) 

2048 True 

2049 >>> exactly_n([True, True, False], 1) 

2050 False 

2051 >>> exactly_n([0, 1, 2, 3, 4, 5], 3, lambda x: x < 3) 

2052 True 

2053 

2054 The iterable will be advanced until ``n + 1`` truthy items are encountered, 

2055 so avoid calling it on infinite iterables. 

2056 

2057 """ 

2058 return len(take(n + 1, filter(predicate, iterable))) == n 

2059 

2060 

2061def circular_shifts(iterable): 

2062 """Return a list of circular shifts of *iterable*. 

2063 

2064 >>> circular_shifts(range(4)) 

2065 [(0, 1, 2, 3), (1, 2, 3, 0), (2, 3, 0, 1), (3, 0, 1, 2)] 

2066 """ 

2067 lst = list(iterable) 

2068 return take(len(lst), windowed(cycle(lst), len(lst))) 

2069 

2070 

2071def make_decorator(wrapping_func, result_index=0): 

2072 """Return a decorator version of *wrapping_func*, which is a function that 

2073 modifies an iterable. *result_index* is the position in that function's 

2074 signature where the iterable goes. 

2075 

2076 This lets you use itertools on the "production end," i.e. at function 

2077 definition. This can augment what the function returns without changing the 

2078 function's code. 

2079 

2080 For example, to produce a decorator version of :func:`chunked`: 

2081 

2082 >>> from more_itertools import chunked 

2083 >>> chunker = make_decorator(chunked, result_index=0) 

2084 >>> @chunker(3) 

2085 ... def iter_range(n): 

2086 ... return iter(range(n)) 

2087 ... 

2088 >>> list(iter_range(9)) 

2089 [[0, 1, 2], [3, 4, 5], [6, 7, 8]] 

2090 

2091 To only allow truthy items to be returned: 

2092 

2093 >>> truth_serum = make_decorator(filter, result_index=1) 

2094 >>> @truth_serum(bool) 

2095 ... def boolean_test(): 

2096 ... return [0, 1, '', ' ', False, True] 

2097 ... 

2098 >>> list(boolean_test()) 

2099 [1, ' ', True] 

2100 

2101 The :func:`peekable` and :func:`seekable` wrappers make for practical 

2102 decorators: 

2103 

2104 >>> from more_itertools import peekable 

2105 >>> peekable_function = make_decorator(peekable) 

2106 >>> @peekable_function() 

2107 ... def str_range(*args): 

2108 ... return (str(x) for x in range(*args)) 

2109 ... 

2110 >>> it = str_range(1, 20, 2) 

2111 >>> next(it), next(it), next(it) 

2112 ('1', '3', '5') 

2113 >>> it.peek() 

2114 '7' 

2115 >>> next(it) 

2116 '7' 

2117 

2118 """ 

2119 # See https://sites.google.com/site/bbayles/index/decorator_factory for 

2120 # notes on how this works. 

2121 def decorator(*wrapping_args, **wrapping_kwargs): 

2122 def outer_wrapper(f): 

2123 def inner_wrapper(*args, **kwargs): 

2124 result = f(*args, **kwargs) 

2125 wrapping_args_ = list(wrapping_args) 

2126 wrapping_args_.insert(result_index, result) 

2127 return wrapping_func(*wrapping_args_, **wrapping_kwargs) 

2128 

2129 return inner_wrapper 

2130 

2131 return outer_wrapper 

2132 

2133 return decorator 

2134 

2135 

2136def map_reduce(iterable, keyfunc, valuefunc=None, reducefunc=None): 

2137 """Return a dictionary that maps the items in *iterable* to categories 

2138 defined by *keyfunc*, transforms them with *valuefunc*, and 

2139 then summarizes them by category with *reducefunc*. 

2140 

2141 *valuefunc* defaults to the identity function if it is unspecified. 

2142 If *reducefunc* is unspecified, no summarization takes place: 

2143 

2144 >>> keyfunc = lambda x: x.upper() 

2145 >>> result = map_reduce('abbccc', keyfunc) 

2146 >>> sorted(result.items()) 

2147 [('A', ['a']), ('B', ['b', 'b']), ('C', ['c', 'c', 'c'])] 

2148 

2149 Specifying *valuefunc* transforms the categorized items: 

2150 

2151 >>> keyfunc = lambda x: x.upper() 

2152 >>> valuefunc = lambda x: 1 

2153 >>> result = map_reduce('abbccc', keyfunc, valuefunc) 

2154 >>> sorted(result.items()) 

2155 [('A', [1]), ('B', [1, 1]), ('C', [1, 1, 1])] 

2156 

2157 Specifying *reducefunc* summarizes the categorized items: 

2158 

2159 >>> keyfunc = lambda x: x.upper() 

2160 >>> valuefunc = lambda x: 1 

2161 >>> reducefunc = sum 

2162 >>> result = map_reduce('abbccc', keyfunc, valuefunc, reducefunc) 

2163 >>> sorted(result.items()) 

2164 [('A', 1), ('B', 2), ('C', 3)] 

2165 

2166 You may want to filter the input iterable before applying the map/reduce 

2167 procedure: 

2168 

2169 >>> all_items = range(30) 

2170 >>> items = [x for x in all_items if 10 <= x <= 20] # Filter 

2171 >>> keyfunc = lambda x: x % 2 # Evens map to 0; odds to 1 

2172 >>> categories = map_reduce(items, keyfunc=keyfunc) 

2173 >>> sorted(categories.items()) 

2174 [(0, [10, 12, 14, 16, 18, 20]), (1, [11, 13, 15, 17, 19])] 

2175 >>> summaries = map_reduce(items, keyfunc=keyfunc, reducefunc=sum) 

2176 >>> sorted(summaries.items()) 

2177 [(0, 90), (1, 75)] 

2178 

2179 Note that all items in the iterable are gathered into a list before the 

2180 summarization step, which may require significant storage. 

2181 

2182 The returned object is a :obj:`collections.defaultdict` with the 

2183 ``default_factory`` set to ``None``, such that it behaves like a normal 

2184 dictionary. 

2185 

2186 """ 

2187 valuefunc = (lambda x: x) if (valuefunc is None) else valuefunc 

2188 

2189 ret = defaultdict(list) 

2190 for item in iterable: 

2191 key = keyfunc(item) 

2192 value = valuefunc(item) 

2193 ret[key].append(value) 

2194 

2195 if reducefunc is not None: 

2196 for key, value_list in ret.items(): 

2197 ret[key] = reducefunc(value_list) 

2198 

2199 ret.default_factory = None 

2200 return ret 

2201 

2202 

2203def rlocate(iterable, pred=bool, window_size=None): 

2204 """Yield the index of each item in *iterable* for which *pred* returns 

2205 ``True``, starting from the right and moving left. 

2206 

2207 *pred* defaults to :func:`bool`, which will select truthy items: 

2208 

2209 >>> list(rlocate([0, 1, 1, 0, 1, 0, 0])) # Truthy at 1, 2, and 4 

2210 [4, 2, 1] 

2211 

2212 Set *pred* to a custom function to, e.g., find the indexes for a particular 

2213 item: 

2214 

2215 >>> iterable = iter('abcb') 

2216 >>> pred = lambda x: x == 'b' 

2217 >>> list(rlocate(iterable, pred)) 

2218 [3, 1] 

2219 

2220 If *window_size* is given, then the *pred* function will be called with 

2221 that many items. This enables searching for sub-sequences: 

2222 

2223 >>> iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3] 

2224 >>> pred = lambda *args: args == (1, 2, 3) 

2225 >>> list(rlocate(iterable, pred=pred, window_size=3)) 

2226 [9, 5, 1] 

2227 

2228 Beware, this function won't return anything for infinite iterables. 

2229 If *iterable* is reversible, ``rlocate`` will reverse it and search from 

2230 the right. Otherwise, it will search from the left and return the results 

2231 in reverse order. 

2232 

2233 See :func:`locate` to for other example applications. 

2234 

2235 """ 

2236 if window_size is None: 

2237 try: 

2238 len_iter = len(iterable) 

2239 return ( 

2240 len_iter - i - 1 for i in locate(reversed(iterable), pred) 

2241 ) 

2242 except TypeError: 

2243 pass 

2244 

2245 return reversed(list(locate(iterable, pred, window_size))) 

2246 

2247 

2248def replace(iterable, pred, substitutes, count=None, window_size=1): 

2249 """Yield the items from *iterable*, replacing the items for which *pred* 

2250 returns ``True`` with the items from the iterable *substitutes*. 

2251 

2252 >>> iterable = [1, 1, 0, 1, 1, 0, 1, 1] 

2253 >>> pred = lambda x: x == 0 

2254 >>> substitutes = (2, 3) 

2255 >>> list(replace(iterable, pred, substitutes)) 

2256 [1, 1, 2, 3, 1, 1, 2, 3, 1, 1] 

2257 

2258 If *count* is given, the number of replacements will be limited: 

2259 

2260 >>> iterable = [1, 1, 0, 1, 1, 0, 1, 1, 0] 

2261 >>> pred = lambda x: x == 0 

2262 >>> substitutes = [None] 

2263 >>> list(replace(iterable, pred, substitutes, count=2)) 

2264 [1, 1, None, 1, 1, None, 1, 1, 0] 

2265 

2266 Use *window_size* to control the number of items passed as arguments to 

2267 *pred*. This allows for locating and replacing subsequences. 

2268 

2269 >>> iterable = [0, 1, 2, 5, 0, 1, 2, 5] 

2270 >>> window_size = 3 

2271 >>> pred = lambda *args: args == (0, 1, 2) # 3 items passed to pred 

2272 >>> substitutes = [3, 4] # Splice in these items 

2273 >>> list(replace(iterable, pred, substitutes, window_size=window_size)) 

2274 [3, 4, 5, 3, 4, 5] 

2275 

2276 """ 

2277 if window_size < 1: 

2278 raise ValueError('window_size must be at least 1') 

2279 

2280 # Save the substitutes iterable, since it's used more than once 

2281 substitutes = tuple(substitutes) 

2282 

2283 # Add padding such that the number of windows matches the length of the 

2284 # iterable 

2285 it = chain(iterable, [_marker] * (window_size - 1)) 

2286 windows = windowed(it, window_size) 

2287 

2288 n = 0 

2289 for w in windows: 

2290 # If the current window matches our predicate (and we haven't hit 

2291 # our maximum number of replacements), splice in the substitutes 

2292 # and then consume the following windows that overlap with this one. 

2293 # For example, if the iterable is (0, 1, 2, 3, 4...) 

2294 # and the window size is 2, we have (0, 1), (1, 2), (2, 3)... 

2295 # If the predicate matches on (0, 1), we need to zap (0, 1) and (1, 2) 

2296 if pred(*w): 

2297 if (count is None) or (n < count): 

2298 n += 1 

2299 yield from substitutes 

2300 consume(windows, window_size - 1) 

2301 continue 

2302 

2303 # If there was no match (or we've reached the replacement limit), 

2304 # yield the first item from the window. 

2305 if w and (w[0] is not _marker): 

2306 yield w[0]