python out_of_process.py

Posted 2021-05-09

def out_of_process(func):
    from types import GeneratorType
    GENERATOR = b'#GENERATOR#'
    DONE = b'#DONE#'

    def runner(conn, *args, **kwargs):
        return_value = func(*args, **kwargs)
        if isinstance(return_value, GeneratorType):
            conn.send(GENERATOR)
            for val in return_value:
                conn.put(val)
            conn.send(DONE)
        else:
            conn.put(return_value)
        conn.close()

    def multi_value(proc, return_queue):
        while True:
            value = return_queue.get()
            if isinstance(value, bytes) and value == DONE:
                break
            yield value

    def wrapper(*args, **kwargs):
        return_queue = Queue()
        args = (return_queue,) + args

        proc = Process(target=runner, args=args, kwargs=kwargs)
        proc.start()

        return_value = return_queue.get()

        if isinstance(return_value, bytes) and return_value == GENERATOR:
            return multi_value(proc, return_queue)
        else:
            proc.join()
            return return_value

    return wrapper