from math import ceil def check_type(name, cls, *args): """Raises an error if any of args is not an instance of cls. name should be the name of the caller.""" error = '{0} required a {1}, got a {2}: {3}' for arg in args: if not isinstance(arg, cls): error = error.format(name, cls.__name__, type(arg).__name__, arg) raise TypeError(error) def check_nonzero(name, arg, argname): """Raises an error if any of args is 0. name should be the name of the caller.""" error = 'argument {0} of {1} must be nonzero' if arg == 0: error = error.format(argname, name) raise TypeError(error) class Range(object): """An implicit regular sequence of integers. >>> r = Range(3, 12, 2) >>> len(r) 5 >>> r[3] 9 >>> r[-3] 7 """ def __init__(self, start, end=None, step=1): if end is None: start, end = 0, start check_type('Range.__init__', int, start, end, step) check_nonzero('Range.__init__', step, 'step') self.start = start self.end = end self.step = step def __repr__(self): return 'Range({0}, {1}, {2})'.format(self.start, self.end, self.step) def __len__(self): return max(0, ceil((self.end - self.start) / self.step)) def __getitem__(self, k): check_type('Range.__getitem__', int, k) if k < 0: k = len(self) + k if k < 0 or k >= len(self): raise IndexError('index out of range') return self.start + k * self.step def __iter__(self): return RangeIter(self.start, self.end, self.step) # Built-in iterators s = [1, 2, 3, 4] siter = iter(s) def print_all(iterator): """Print all items in an iterator on consecutive lines. >>> print_all(siter) 1 2 3 4 """ try: print(next(iterator)) except StopIteration: return print_all(iterator) # Defining an iterator class RangeIter(object): """An iterator over a Range. >>> ri = RangeIter(3, 7, 1) >>> next(ri) 3 >>> next(ri) 4 >>> sum(ri) 11 """ def __init__(self, start, end, step): check_type('RangeIter.__init__', int, start, end) check_nonzero('RangeIter.__init__', step, 'step') self.current = start self.end = end self.step = step self.sign = 1 if step > 0 else -1 def __next__(self): if self.current * self.sign >= self.end * self.sign: raise StopIteration result = self.current self.current += self.step return result def __iter__(self): return self class LettersIter(object): """An iterator over the letters a through d.""" def __init__(self): self.current = 'a' def __next__(self): if self.current > 'd': raise StopIteration result = self.current self.current = chr(ord(result)+1) return result def __iter__(self): return self class FibIter(object): """An iterator over Fibonacci numbers. >>> fi = FibIter() >>> [(i, n) for i, n in zip(range(1, 9), fi)] [(1, 0), (2, 1), (3, 1), (4, 2), (5, 3), (6, 5), (7, 8), (8, 13)] """ def __init__(self): self.prev = -1 self.current = 1 def __next__(self): self.prev, self.current = self.current, self.prev + self.current return self.current def __iter__(self): return self # Generator expressions double_fibs = (fib * 2 for fib in FibIter()) # Generator functions def range_generator(start, end, step): """A generator function that returns an iterator over a range. >>> for num in range_generator(7, 3, -1): print(num) 7 6 5 4 """ check_type('range_generator', int, start, end) check_nonzero('range_generator', step, 'step') sign = 1 if step > 0 else -1 while True: if start * sign >= end * sign: raise StopIteration yield start start += step def letters_generator(): """A generator function that returns an iterator over letters. >>> for letter in letters_generator(): print(letter) a b c d """ current = 'a' while current <= 'd': yield current current = chr(ord(current)+1) def fib_generator(): """A generator function for Fibonacci numbers. >>> fg = fib_generator() >>> [(i, n) for i, n in zip(range(1, 9), fg)] [(1, 0), (2, 1), (3, 1), (4, 2), (5, 3), (6, 5), (7, 8), (8, 13)] """ yield 0 prev, current = 0, 1 while True: yield current prev, current = current, prev + current # Map and filter def map_gen(fn, iterable): """Lazily map a function across an iterable. >>> squares = map_gen(lambda x: x*x, s) >>> print_all(squares) 1 4 9 16 """ iterator = iter(iterable) while True: yield fn(next(iterator)) def filter_gen(fn, iterable): """Lazily filter an iterable by the given function. >>> evens = filter_gen(lambda x: x % 2 == 0, s) >>> print_all(evens) 2 4 """ iterator = iter(iterable) while True: item = next(iterator) if fn(item): yield item # An infinite sequence of bitstrings from itertools import product def bitstrings(): """Generate bitstrings in order of increasing size. >>> bs = bitstrings() >>> [next(bs) for _ in range(0, 10)] ['', '0', '1', '00', '01', '10', '11', '000', '001', '010'] """ size = 0 while True: tuples = product(('0', '1'), repeat=size) for elem in tuples: yield ''.join(elem) size += 1