# Parallel processing

import cPickle

import itertools

import Queue

import rand

import types

import warnings

import collection

import dicts

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 vars_id_dict(locals_, globals_, *misc):

    '''Usage: vars_id_dict(locals(), globals(), misc...)'''

    vars_ = map(lambda v: v.values(), [locals_, globals_]) + list(misc)

    return dicts.id_dict(vars_)

def prepickle(value, vars_id_dict_):

    def filter_(value):

        id_ = id(value)

        if id_ in vars_id_dict_: value = id_

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

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

        else: try_pickle(value)

        return value

    return collection.rmap(filter_, value)

def post_unpickle(value, vars_id_dict_):

    def filter_(value):

        try: return vars_id_dict_[value] # value is an id()

        except KeyError: return value

    return collection.rmap(filter_, 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=(), kw_args={}, callback=None):

        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, locals_=None, globals_=None, *shared):

        '''

        @param processes If 0, uses SyncPool

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

        '''

        if locals_ == None: locals_ = locals()

        if globals_ == None: globals_ = globals()

        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_()

        # Values that may be pickled by id()

        self.vars_id_dict = vars_id_dict(locals_, globals_, *shared)

    def share(self, value):

        '''Call this on every value that that may be pickled by id()'''

        self.vars_id_dict[id(value)] = value

    def main_loop(self):

        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, self.post_unpickle(call.args),

                    self.post_unpickle(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=(), kw_args={}, callback=None):

        assert callback == None, 'Callbacks not supported'

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

        call.callback = callback # store this inside the Runnable

        self.queue.put_nowait(call)

        return self.Result()

    def prepickle(self, value): return prepickle(value, self.vars_id_dict)

    def post_unpickle(self, value):

        return post_unpickle(value, self.vars_id_dict)