# Parallel processing

import cPickle

import Queue

import types

import warnings

import exc

from Runnable import Runnable

def try_pickle(value):

    try: cPickle.dumps(value)

    except Exception, e:

        exc.add_msg(e, 'Tried to pickle: '+repr(value))

        raise

def prepickle(value):

    if isinstance(value, types.FunctionType): return Runnable(value)

    elif isinstance(value, MultiProducerPool): return None

    elif isinstance(value, list): return map(prepickle, list_)

    elif isinstance(value, dict):

        return dict(((k, prepickle(v)) for k, v in value.iteritems()))

    else: return value

class SyncPool:

    '''A dummy synchronous Pool to use if multiprocessing is not available'''

    def __init__(self, processes=None): pass

    class Result:

        def __init__(self, value): self.value = value

        def get(timeout=None): return self.value

        def wait(timeout=None): pass

        def ready(): return True

        def successful(): return True # TODO: False if raised exception

    def apply_async(self, func, args=None, kw_args=None, callback=None):

        if args == None: args = ()

        if kw_args == None: kw_args = {}

        if callback == None: callback = lambda v: None

        value = func(*args, **kw_args)

        callback(value)

        return self.Result(value)

class MultiProducerPool:

    '''A multi-producer pool. You must call pool.main_loop() in the thread that

    created this to process new tasks.'''

    def __init__(self, processes=None):

        '''

        @param processes If 0, uses SyncPool

        @post The # processes actually used is made available in self.process_ct

        '''

        try:

            if processes == 0: raise ImportError('turned off')

            import multiprocessing

            import multiprocessing.pool

        except ImportError, e:

            warnings.warn(UserWarning('Not using parallel processing: '+str(e)))

            processes = 1

            Pool_ = SyncPool

            Queue_ = Queue.Queue

        else:

            if processes == None: processes = multiprocessing.cpu_count()

            Pool_ = multiprocessing.pool.Pool

            Queue_ = multiprocessing.Queue

        self.process_ct = processes

        self.pool = Pool_(processes)

        self.queue = Queue_()

    def main_loop(self):

        '''@param pool Must be a pool returned by mk_pool()'''

        try:

            while True:

                # block=False raises Empty immediately if the queue is empty,

                # which indicates that the program is done

                call = self.queue.get(block=False)

                self.pool.apply_async(call.func, call.args, call.kw_args,

                    call.callback)

        except Queue.Empty: pass

    class Result:

        def get(timeout=None): raise NotImplementedError()

        def wait(timeout=None): raise NotImplementedError()

        def ready(): raise NotImplementedError()

        def successful(): raise NotImplementedError()

    def apply_async(self, func, args=None, kw_args=None, callback=None):

        if args == None: args = ()

        if kw_args == None: kw_args = {}

        call = Runnable(func, *prepickle(args), **prepickle(kw_args))

        call.callback = callback # store this inside the Runnable

        # Try pickling the args. If it fails, we'll get a full traceback here,

        # which is not provided with pickling errors in multiprocessing's Pool.

        try_pickle(call)

        self.queue.put_nowait(call)

        return self.Result()