Python generator pipelines demo

Pipeline.py

import types


class Pipeline(object):
    def __init__(self, source):
        if type(source) == types.FunctionType:
            self.pipe = source()
        else:
            self.pipe = source

    def __iter__(self):
        for droplet in self.pipe:
            yield droplet

    def add_link(self, link_fn):
        self.pipe = link_fn(self.pipe)


def main():
    def add(amount):
        def internal(source):
            for droplet in source:
                yield droplet + amount
        return internal

    def square(source):
        for droplet in source:
            yield droplet ** 2

    p = Pipeline(range(10))
    p.add_link(add(5))
    p.add_link(square)

    # now to use suction to draw the source through
    for droplet in p:
        print(droplet)


if __name__ == '__main__':
    main()