flowit.py

import subprocess
import collections
import itertools
from multiprocessing import Pool, cpu_count
from functools import partial


def foo(input):
    return 1


def check(v):
    return v == 1


def bar(input):
    return 'bar'


def baz(input):
    return 'baz'


class IfTask(object):

    def __init__(self, check, a, b):
        self.check = check
        self.a = a
        self.b = b

    def __call__(self, input):
        if check(input):
            return do_flow(self.a, input)
        return do_flow(self.b, input)



class MapTask(object):

    def __init__(self, args, task):
        self.args = args
        self.task = task
        self.pool = Pool(cpu_count())

    def iter_input(self, input):
        for args in self.args:
            if not isinstance(args, (tuple, list)):
                args = [args]
            yield tuple([input] + args)


    def __call__(self, input):
        return self.pool.map(self.task, self.iter_input(input))



def finish(input):
    print('done: %s' % input)
    return input


def do_flow(flow, result=None):
    for item in flow:
        print(item, result)
        result = item(result)
    return result


def digit(input):
    input, zone = input
    cmd = 'dig +short %s' % zone
    print(cmd)
    ip = subprocess.check_output(cmd, shell=True).strip()
    return (input, ip)


myflow = [
    foo,
    IfTask(check, [bar], [baz]),
    MapTask([
        'ionrock.org',
        'google.com',
        'rackspace.com',
    ], digit),
    finish,
]


def main():
    print(do_flow(myflow))


if __name__ == '__main__':

    main()