Guest User

Untitled

a guest
May 26th, 2018
78
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 1.04 KB | None | 0 0
  1. import collections
  2. import typing
  3.  
  4. T = typing.TypeVar('T')
  5.  
  6.  
  7. class ChainableIterator(typing.Iterator[T], collections.Iterator):
  8. def __init__(self, *iterators: typing.Iterator[T]):
  9. self._index = -1
  10. self._iterators = list((a for a in b) for b in iterators)
  11.  
  12. def __next__(self):
  13. try:
  14. if self._index < 0:
  15. raise StopIteration
  16. return next(self._iterators[self._index])
  17. except StopIteration as e:
  18. if self._index < len(self._iterators) - 1:
  19. self._index += 1
  20. return next(self)
  21. raise e
  22.  
  23. def chain(self, *iterators: typing.Iterator[T]) -> 'ChainableIterator[T]':
  24. self._iterators.extend((a for a in b) for b in iterators)
  25. return self
  26.  
  27.  
  28. if __name__ == '__main__':
  29. iterator = ChainableIterator((1, 2, 3), (4, 5, 6))
  30. iterator.chain((7, 8, 9), (10, 11, 12))
  31.  
  32. for i in range(1, 13):
  33. assert i == next(iterator)
  34.  
  35. try:
  36. next(iterator)
  37. assert False
  38. except StopIteration:
  39. pass
Add Comment
Please, Sign In to add comment