from __future__ import print_function

from six.moves import queue as Queue

from six.moves import range

from functools import partial

import threading

import time

from tqdm import tqdm

DEFAULT_THREADS = 20

class ThreadedQueue(object):

"""Grant threaded task processing to any derived class."""

def __init__(self, n_threads, queue_size=0, progress=None):

self._n_threads = n_threads

self._queue = Queue.Queue(maxsize=queue_size) # 0 = infinite size

self._error_queue = Queue.Queue(maxsize=queue_size)

self._threads = ()

self._terminate = threading.Event()

self._processed_lock = threading.Lock()

self.processed = 0

self._inserted = 0

self.with_progress = progress

self.start_threads(n_threads)

@property

def pending(self):

return self._queue.qsize()

def put(self, fn):

"""

Enqueue a task function for processing.

Requires:

fn: a function object that takes one argument

that is the interface associated with each

thread.

e.g. def download(api):

results.append(api.download())

self.put(download)

Returns: self

"""

self._inserted += 1

self._queue.put(fn, block=True)

return self

def start_threads(self, n_threads):

"""

Terminate existing threads and create a

new set if the thread number doesn't match

the desired number.

Required:

n_threads: (int) number of threads to spawn

Returns: self

"""

if n_threads == len(self._threads):

return self

# Terminate all previous tasks with the existing

# event object, then create a new one for the next

# generation of threads. The old object will hang

# around in memory until the threads actually terminate

# after another iteration.

self._terminate.set()

self._terminate = threading.Event()

threads = []

for _ in range(n_threads):

worker = threading.Thread(

target=self._consume_queue,

args=(self._terminate,)

)

worker.daemon = True

worker.start()

threads.append(worker)

self._threads = tuple(threads)

return self

def are_threads_alive(self):

"""Returns: boolean indicating if any threads are alive"""

return any(map(lambda t: t.is_alive(), self._threads))

def kill_threads(self):

"""Kill all threads."""

self._terminate.set()

while self.are_threads_alive():

time.sleep(0.001)

self._threads = ()

return self

def _initialize_interface(self):

"""

This is used to initialize the interfaces used in each thread.

You should reimplement it in subclasses. For example, return

an API object, file handle, or network connection. The functions

you pass into the self._queue will get it as the first parameter.

e.g. an implementation in a subclass.

def _initialize_interface(self):

return HTTPConnection()

def other_function(self):

def threaded_file_read(connection):

# do stuff

self._queue.put(threaded_file_handle)

Returns: Interface object used in threads

"""

return None

def _close_interface(self, interface):

"""Allows derived classes to clean up after a thread finishes."""

pass

def _consume_queue(self, terminate_evt):

"""

This is the main thread function that consumes functions that are

inside the _queue object. To use, execute self._queue(fn), where fn

is a function that performs some kind of network IO or otherwise

benefits from threading and is independent.

terminate_evt is automatically passed in on thread creation and

is a common event for this generation of threads. The threads

will terminate when the event is set and the queue burns down.

Returns: void

"""

interface = self._initialize_interface()

while not terminate_evt.is_set():

try:

fn = self._queue.get(block=True, timeout=0.01)

except Queue.Empty:

continue # periodically check if the thread is supposed to die

fn = partial(fn, interface)

try:

self._consume_queue_execution(fn)

except Exception as err:

self._error_queue.put(err)

self._close_interface(interface)

def _consume_queue_execution(self, fn):

"""

The actual task execution in each thread. This

is broken out so that exceptions can be caught

in derived classes and allow them to manipulate

the errant task, e.g. putting it back in the queue

for a retry.

Every task processed will automatically be marked complete.

Required:

[0] fn: A curried function that includes the interface

as its first argument.

Returns: void

"""

# `finally` fires after all success or exceptions

# exceptions are handled in derived classes

# and uncaught ones are caught as a last resort

# in _consume_queue to be raised on the main thread.

try:

fn()

finally:

with self._processed_lock:

self.processed += 1

self._queue.task_done()

def _check_errors(self):

try:

err = self._error_queue.get(block=False)

self._error_queue.task_done()

self.kill_threads()

raise err

except Queue.Empty:

pass

def wait(self, progress=None):

"""

Allow background threads to process until the

task queue is empty. If there are no threads,

in theory the queue should always be empty

as processing happens immediately on the main thread.

Optional:

progress: (bool or str) show a tqdm progress bar optionally

with a description if a string is provided

Returns: self (for chaining)

Raises: The first exception recieved from threads

"""

if not len(self._threads):

return self

desc = None

if type(progress) is str:

desc = progress

last = self._inserted

with tqdm(total=self._inserted, disable=(not progress), desc=desc) as pbar:

# Allow queue to consume, but check up on

# progress and errors every tenth of a second

while not self._queue.empty():

size = self._queue.qsize()

delta = last - size

if delta != 0: # We should crash on negative numbers

pbar.update(delta)

last = size

self._check_errors()

time.sleep(0.1)

# Wait until all tasks in the queue are

# fully processed. queue.task_done must be

# called for each task.

self._queue.join()

self._check_errors()

final = self._inserted - last

if final:

pbar.update(final)

if self._queue.empty():

self._inserted = 0

return self

def __del__(self):

self.wait() # if no threads were set the queue is always empty

self.kill_threads()

def __enter__(self):

if self.__class__ is ThreadedQueue and self._n_threads == 0:

raise ValueError("Using 0 threads in base class ThreadedQueue with statement will never exit.")

self.start_threads(self._n_threads)

return self

def __exit__(self, exception_type, exception_value, traceback):

self.wait(progress=self.with_progress)

self.kill_threads()