william-silversmith · May 10, 2018 · May 10, 2018
diff --git a/AUTHORS b/AUTHORS
@@ -1 +1,2 @@
+Jingpeng Wu <jingpeng.wu@gmail.com>
 William Silversmith <william.silversmith@gmail.com>
diff --git a/ChangeLog b/ChangeLog
@@ -1,6 +1,12 @@
 CHANGES
 =======
 
+0.4.1
+-----
+
+* fix: updated ptq to use updated cloud-volume semantics
+* Update README.md
+
 0.4.0
 -----
 

diff --git a/requirements.txt b/requirements.txt
@@ -2,8 +2,8 @@ cloud-volume>=0.20.2
 google-api-python-client
 google-cloud>=0.32
 google-auth-httplib2==0.0.2
-numpy>=1.14.3
+numpy>=1.13.1
 pytest>=3.3.1
 pbr
-tqdm==4.23.2
+-e git+https://github.com/seung-lab/tqdm.git#egg=tqdm
 requests>2,<3
diff --git a/taskqueue/__init__.py b/taskqueue/__init__.py
@@ -1,5 +1,5 @@
 from .registered_task import RegisteredTask, MockTask, PrintTask
-from .taskqueue import TaskQueue, MockTaskQueue
+from .taskqueue import TaskQueue, MockTaskQueue, LocalTaskQueue
 from .secrets import (
   QUEUE_NAME, TEST_QUEUE_NAME, QUEUE_TYPE, 
   PROJECT_NAME, APPENGINE_QUEUE_URL

diff --git a/taskqueue/registered_task.py b/taskqueue/registered_task.py
@@ -77,13 +77,17 @@ def __repr__(self):
         return string + ")"  
 
 class PrintTask(RegisteredTask):
-    def __init__(self):
+    def __init__(self, txt=''):
         super(PrintTask, self).__init__()
+        self.txt = txt
     def execute(self):
-        print(self)
+        if self.txt:
+            print(str(self) + ": " + str(self.txt))
+        else:
+            print(self)
 
 class MockTask(RegisteredTask):
-    def __init__(self):
+    def __init__(self, *args):
         super(MockTask, self).__init__()
     def execute(self):
         pass
diff --git a/taskqueue/taskqueue.py b/taskqueue/taskqueue.py
@@ -4,8 +4,11 @@
 import json
 from functools import partial
 
+import concurrent.futures 
+import multiprocessing as mp
 import googleapiclient.errors
 import numpy as np
+from tqdm import tqdm
 
 from cloudvolume.threaded_queue import ThreadedQueue
 
@@ -16,7 +19,6 @@
 from .secrets import PROJECT_NAME, QUEUE_NAME, QUEUE_TYPE
 
 def totask(task):
-    print(task)
     taskobj = payloadBase64Decode(task['payloadBase64'])
     taskobj._id = task['id']
     return taskobj
@@ -231,3 +233,36 @@ def __enter__(self):
 
     def __exit__(self, exception_type, exception_value, traceback):
       pass
+
+class LocalTaskQueue(object):
+    def __init__(self, parallel=1, queue_name='', queue_server='', progress=True):
+        if parallel and type(parallel) == bool:
+            parallel = mp.cpu_count()
+
+        self.parallel = parallel
+        self.queue = []
+        self.progress = progress
+
+    def insert(self, task):
+        self.queue.append(task)
+
+    def wait(self, progress=None):
+      return self
+
+    def kill_threads(self):
+      return self
+
+    def __enter__(self):
+      return self
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        with tqdm(total=len(self.queue), desc="Tasks") as pbar:
+            with concurrent.futures.ProcessPoolExecutor(max_workers=self.parallel) as executor:
+                for _ in executor.map(task_execute, self.queue):
+                    pbar.update()
+        self.queue = []
+
+# Necessary to define here to make the 
+# function picklable
+def task_execute(task):
+    task.execute()
diff --git a/test/test_taskqueue.py b/test/test_taskqueue.py
@@ -4,7 +4,7 @@
 from six.moves import range
 import pytest
 
-from taskqueue import RegisteredTask, TaskQueue, PrintTask
+from taskqueue import RegisteredTask, TaskQueue, MockTask, PrintTask, LocalTaskQueue
 from taskqueue import QUEUE_NAME
 
 TRAVIS_BRANCH = None if 'TRAVIS_BRANCH' not in os.environ else os.environ['TRAVIS_BRANCH']
@@ -98,3 +98,9 @@ def test_400_errors():
     with TaskQueue(n_threads=1, queue_name=QUEUE_NAME, queue_server=qtype, qurl=QURL) as tq:
       tq.delete('nonexistent')
 
+def test_local_taskqueue():
+  with LocalTaskQueue(parallel=True, progress=False) as tq:
+    for i in range(20000):
+      tq.insert(
+        MockTask(i)
+      )
Original file line number	Diff line number	Diff line change
		@@ -1 +1,2 @@
		Jingpeng Wu <jingpeng.wu@gmail.com>
		William Silversmith <william.silversmith@gmail.com>