leancheck.py

from itertools import (chain,   # catenation for enumerations
                       product, # cartesian product
                       )

from collections import deque
import string

def take(n, t):
    it = iter(t)
    for _ in range(n):
        yield next(t)

def to_tiers(t):
    yield from map(lambda x: [x], t)

def tiers(f):
    def copy(t):
        if isinstance(t, Tiers):
            return t.tee()
        return t

    class Tiers:
        def __init__(self, gen, parent=None):
            self.iterators = [self]
            self.gen = gen
            self.parent = parent
            self._d = deque()

        def __iter__(self):
            return self

        def tee(self):
            v = Tiers(self.gen, parent=self)
            v._d = deque(self._d)
            self.ancestor.iterators.append(v)
            return v

        def __next__(self):
            if not self._d:
                next_tier = next(self.gen)
                for it in self.ancestor.iterators:
                    it._d.append(deque(map(copy, next_tier)))

            return list(self._d.popleft())

        @property
        def ancestor(self):
            if self.parent is None:
                return self
            return self.parent.ancestor


    return lambda *a, **kw: Tiers(f(*a, **kw))

@tiers
def nats():
    n = 0
    while True:
        n += 1
        yield [n]

@tiers
def ints():
    n = 0
    yield [n]

    while True:
        n += 1
        yield [n]
        yield [-n]

@tiers
def mapT(f, t):
    yield from map(lambda ts: map(f, ts), t)

@tiers
def delay(t):
    yield []
    yield from t

@tiers
def concatT(xss, yss):
    # tiers t -> tiers t -> tiers t
    # [a, b, c] `concatT` [d, e, f] => [a+d, b+e, c+f]
    
    while True:
        try:
            xs = next(xss)
        except StopIteration:
            yield from yss
            return
        try:
            ys = next(yss)
        except StopIteration:
            yield from chain([xs], xss)
            return

        yield list(chain(xs, ys))
        yield from concatT(xss, yss)

@tiers
def productT(s, t):
    # tiers s -> tiers t -> tiers (s, t)
    try:
        s_head = next(s)
    except StopIteration:
        yield []
        return

    t_ = t.tee()
    prods = (product(s_head, ys) for ys in t)
    remaining = delay(productT(s, t_))
    yield from concatT(prods, remaining)


def interleave(s, *t):
    s = iter(s)
    try:
        yield next(s)
    except StopIteration:
        if t:
            yield from interleave(*t)
        return
        
    yield from interleave(*(t + (s,)))

@tiers
def interleaveT(*t):
    for tiers in zip(*t):
        yield list(interleave(*tiers))

@tiers
def cons0(v):
    yield [v]

@tiers
def cons1(f, arg):
    return delay(mapT(f, arg))

@tiers
def cons2(f, arg1, arg2):
    return delay(mapT(lambda t: f(*t), productT(arg1, arg2)))

def cons(x, xs):
    return [x] + xs

@tiers
def lists(t):
    yield from concatT(cons0([]), cons2(cons, t.tee(), lists(t.tee())))


@tiers
def chars():
    yield from interleaveT(to_tiers(string.ascii_lowercase),
                           to_tiers(string.whitespace),
                           to_tiers(string.ascii_uppercase),
                           to_tiers(string.digits),
                           to_tiers(string.punctuation),
                           )

@tiers
def bools():
    yield [False, True]


for lst_tier in take(5, lists(ints())):
    print(lst_tier)

"""
[[]]
[[0]]
[[0, 0], [1]]
[[0, 0, 0], [0, 1], [1, 0], [-1]]
[[0, 0, 0, 0], [0, 0, 1], [0, 1, 0], [0, -1], [1, 0, 0], [1, 1], [-1, 0], [2]]
"""