Generators as a Library

generator.py

"""
A simple library level implementation of yielding generators using threads. Not
efficient, but it works!
"""

import threading

class BoundedConcurrentQueue(object):
    def __init__(self, cap):
        self.cap = cap
        self.xs = []

        self.lock = threading.Lock()
        self.data_ready = threading.Condition(self.lock)
        self.space_free = threading.Condition(self.lock)

    def push(self, x):
        with self.lock:
            while len(self.xs) > self.cap:
                self.space_free.wait()
            self.xs.append(x)
            self.data_ready.notify()

    def pop(self):
        with self.lock:
            while len(self.xs) == 0:
                self.data_ready.wait()
            x = self.xs.pop(0)
            self.space_free.notify()
            return x

class Generator(object):
    def __init__(self):
        self.data = BoundedConcurrentQueue(1)
        self.ready = BoundedConcurrentQueue(1)

        self.t = threading.Thread(target=self.run_f, args=())
        self.t.daemon = True
        self.t.start()

    def next(self):
        self.ready.push(None)
        return self.data.pop()

    def yield_val(self, x):
        self.data.push(x)
        self.ready.pop()

    def run_f(self):
        self.ready.pop()
        self.f()
        while True:
            self.yield_val(None)

class Range(Generator):
    def __init__(self, start, stop):
        Generator.__init__(self)
        self.start = start
        self.stop = stop

    def f(self):
        i = self.start
        while self.stop is None or i < self.stop:
            self.yield_val(i)
            i += 1

class Cycle(Generator):
    def __init__(self, xs):
        Generator.__init__(self)
        self.xs = xs

    def f(self):
        while True:
            for x in self.xs:
                self.yield_val(x)

class Take(Generator):
    def __init__(self, gen, n):
        Generator.__init__(self)
        self.gen = gen
        self.n = n

    def f(self):
        for i in range(self.n):
            x = self.gen.next()
            if x is None:
                return
            self.yield_val(x)

def wrap(gen):
    x = gen.next()
    while x is not None:
        yield x
        x = gen.next()

def main():
    r = Range(0, 10)
    for i in wrap(r):
        print i

    c = Cycle(["a", "b", "c"])
    t = Take(c, 10)
    for i in wrap(t):
        print i

if __name__ == "__main__":
    main()