Project

General

Profile

1 1881 aaronmk
# Iterators and iterables
2 1711 aaronmk
3 2254 aaronmk
import rand
4
5 1881 aaronmk
def is_iterable(value): return getattr(value, '__iter__', False)
6
7 2254 aaronmk
# Sentinel to use in iter() if function return value should be ignored.
8
# Can't use object() because iter() compares sentinel using ==.
9
iter_end = rand.rand_int()
10 1892 aaronmk
11
def func_iter(func):
12
    '''Used to iterate a generating function which raises StopIteration instead
13
    of returning a sentinel value'''
14
    return iter(func, iter_end)
15
16
def consume_iter(iter_):
17
    for elem in iter_: pass
18
19 1711 aaronmk
def flatten(outer_iter):
20 2156 aaronmk
    if is_iterable(outer_iter):
21
        for iter_ in outer_iter:
22
            if is_iterable(iter_):
23
                for el in iter_: yield el
24
            else: yield iter_
25
    else: yield outer_iter
26 1711 aaronmk
27
def flatten_n(outer_iter, depth=2):
28
    if depth < 2: return outer_iter
29
    elif depth == 2: return flatten(outer_iter)
30
    else: return flatten(flatten_n(outer_iter, depth-1))