utils/multiprogress.py

# -*- coding:utf-8 –*-
import os
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from multiprocessing import Pool

from tqdm import tqdm


class MultiProgress(object):
    def __init__(self, info_lst, workers=6):
        self.info_lst = info_lst
        self.num_items = len(info_lst)
        self.workers = workers
        self.pool = Pool(workers)

    def run(self, func):
        with tqdm(total=self.num_items) as bar:
            for _ in self.pool.imap(func, self.info_lst):
                bar.update(1)

    def __repr__(self):
        return f'MultiProgress(info_lst={type(self.info_lst)}, workers={self.workers})'


class MultiThreading(object):
    def __init__(self, info_lst, workers=6, pbar=True, deception='Running...'):
        self.info_lst = list(info_lst)
        self.num_items = len(info_lst)
        self.process_bar = pbar
        self.deception = deception
        self.workers = workers
        self.exe = ThreadPoolExecutor(workers)

    def run(self, func):
        if self.process_bar:
            res = list(tqdm(self.exe.map(func, self.info_lst), total=self.num_items, desc=self.deception))
        else:
            res = list(self.exe.map(func, self.info_lst))
        return list(res)

    def __repr__(self):
        return f'MultiThreading(info_lst={type(self.info_lst)}, workers={self.workers})'


class MultiProcess(object):
    def __init__(self, info_lst, pbar=True):
        self.info_lst = info_lst
        self.num_items = len(info_lst)
        self.workers = os.cpu_count()
        self.process_bar = pbar

    def run(self, func):
        with ProcessPoolExecutor() as exe:
            if self.process_bar:
                res = list(tqdm(exe.map(func, self.info_lst), total=self.num_items))
            else:
                res = list(exe.map(func, self.info_lst))
        return res

    def __repr__(self):
        return f'MultiProcess(info_lst={self.num_items}, workers={self.workers})'