Hugh's Blog

Python 简单实现多线程任务队列

 

使用队列实现一个简单的多线程任务。

from queue import Queue
from threading import Thread

class TaskQueue(Queue):
    def __init__(self, num_workers=1):
        Queue.__init__(self)
        self.num_workers = num_workers
        self.start_workers()

    def start_workers(self):
        for _ in range(self.num_workers):
            t = Thread(target=self.do_task)
            t.daemon = True
            t.start()

    def add_task(self, func, *args, **kwargs):
        args = args or ()
        kwargs = kwargs or {}
        self.put((func, args, kwargs))

    def do_task(self):
        while True:
            func, args, kwargs = self.get()
            func(*args, **kwargs)
            self.task_done()

q = TaskQueue(num_workers=5)

if __name__ == '__main__':
    import time

    def test_func(task_id:int):
        time.sleep(task_id)
        print('Task done: %d' % task_id)

    for i in range(10):
        q.add_task(test_func, i)

    q.join()
    print('Task all done')