nilmdb/nilmdb/utils/iteratorizer.py

import Queue
import threading
import sys
import contextlib

# This file provides a context manager that converts a function
# that takes a callback into a generator that returns an iterable.
# This is done by running the function in a new thread.

# Based partially on http://stackoverflow.com/questions/9968592/

class IteratorizerThread(threading.Thread):
    def __init__(self, queue, function, curl_hack):
        """
        function: function to execute, which takes the
        callback (provided by this class) as an argument
        """
        threading.Thread.__init__(self)
        self.function = function
        self.queue = queue
        self.die = False
        self.curl_hack = curl_hack

    def callback(self, data):
        try:
            if self.die:
                raise Exception() # trigger termination
            self.queue.put((1, data))
        except:
            if self.curl_hack:
                # We can't raise exceptions, because the pycurl
                # extension module will unconditionally print the
                # exception itself, and not pass it up to the caller.
                # Instead, just return a value that tells curl to
                # abort.  (-1 would be best, in case we were given 0
                # bytes, but the extension doesn't support that).
                self.queue.put((2, sys.exc_info()))
                return 0
            raise

    def run(self):
        try:
            result = self.function(self.callback)
        except:
            self.queue.put((2, sys.exc_info()))
        else:
            self.queue.put((0, result))

@contextlib.contextmanager
def Iteratorizer(function, curl_hack = False):
    """
    Context manager that takes a function expecting a callback,
    and provides an iterable that yields the values passed to that
    callback instead.

    function: function to execute, which takes a callback
    (provided by this context manager) as an argument

        with iteratorizer(func) as it:
            for i in it:
                print 'callback was passed:', i
        print 'function returned:', it.retval
    """
    queue = Queue.Queue(maxsize = 1)
    thread = IteratorizerThread(queue, function, curl_hack)
    thread.daemon = True
    thread.start()

    class iteratorizer_gen(object):
        def __init__(self, queue):
            self.queue = queue
            self.retval = None

        def __iter__(self):
            return self

        def next(self):
            (typ, data) = self.queue.get()
            if typ == 0:
                # function has returned
                self.retval = data
                raise StopIteration
            elif typ == 1:
                # data is available
                return data
            else:
                # callback raised an exception
                raise data[0], data[1], data[2]

    try:
        yield iteratorizer_gen(queue)
    finally:
        # Ask the thread to die, if it's still running.
        thread.die = True
        while thread.isAlive():
            try:
                queue.get(True, 0.01)
            except:
                pass