[@boydfd] add tqdm multi-thread version.

rlin-tw · rlin-tw · commit 35b87b13694e · 2019-08-10T15:33:59.000+08:00
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,2 @@
+# Created by .ignore support plugin (hsz.mobi)
+__pycache__
diff --git a/test.py b/test.py
@@ -0,0 +1,20 @@
+import threading
+from concurrent.futures import ThreadPoolExecutor
+import time
+
+from tqdm_multithread_factory import TqdmMultiThreadFactory
+
+
+def demo(factory, position, total):
+    with factory.create(position, total) as progress:
+        for _ in range(0, total, 5):
+            progress.update(5)
+            time.sleep(0.001 * (position % 5 + 1))
+
+
+with ThreadPoolExecutor(max_workers=20) as executor:
+    tasks = range(100)
+    lock = threading.Lock()
+    multi_thread_factory = TqdmMultiThreadFactory()
+    for i, url in enumerate(tasks, 1):
+        executor.submit(demo, multi_thread_factory, i, 100)
diff --git a/tqdm_multithread.py b/tqdm_multithread.py
@@ -0,0 +1,84 @@
+import io
+import math
+import sys
+
+from tqdm import tqdm
+
+
+class TqdmMultiThread(io.StringIO):
+    def __init__(self, texts, id, total=100, lock=None):
+        super().__init__()
+        self.id = id
+        self.lock = lock
+        text = "progress #{}".format(id)
+        with self.lock:
+            self.texts = texts
+            self.texts[id] = ''
+            self.down()
+        self.tqdm = tqdm(
+            total=total,
+            position=id,
+            desc=text,
+            file=self
+        )
+
+    def update(self, *params):
+        self.tqdm.update(*params)
+
+    def write(self, buf):
+        self.with_lock_call(self._write, buf)
+
+    def _write(self, buf):
+        buf = self.strip(buf)
+        if buf:
+            self.texts[self.id] = buf
+
+    @classmethod
+    def strip(cls, buf):
+        return buf.strip('\r\n\t\x1b[1A')
+
+    def flush(self):
+        self.with_lock_call(self._flush)
+
+    def _flush(self):
+        self.top()
+
+        for key in sorted(self.texts):
+            sys.stdout.write(self.texts[key])
+            self.down()
+
+    def up(self):
+        self.print('\x1b[1A')
+
+    def down(self):
+        self.print('\n')
+
+    def with_lock_call(self, func, *params):
+        if self.lock:
+            with self.lock:
+                func(*params)
+        else:
+            func(*params)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.tqdm.close()
+        with self.lock:
+            self.top()
+            self.print(self.texts[self.id])
+            self.bottom()
+            del self.texts[self.id]
+
+    def print(self, text):
+        sys.stdout.write(text)
+        sys.stdout.flush()
+
+    def bottom(self):
+        for i in range(self.get_length()):
+            self.down()
+
+    def top(self):
+        for i in range(self.get_length()):
+            self.up()
+
+    def get_length(self):
+        return len(self.texts)
diff --git a/tqdm_multithread_factory.py b/tqdm_multithread_factory.py
@@ -0,0 +1,15 @@
+import os
+import threading
+
+from tqdm import tqdm
+
+from tqdm_multithread import TqdmMultiThread
+
+
+class TqdmMultiThreadFactory:
+    def __init__(self):
+        self.texts = {}
+        self.lock = threading.Lock()
+
+    def create(self, id, total):
+        return TqdmMultiThread(self.texts, id, total, self.lock)

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+# Created by .ignore support plugin (hsz.mobi)`
	`2`	`+__pycache__`