tqdm Function decorator

tqdm_function_decorator.py

from tqdm import tqdm

def tqdm_function_decorator(*args, **kwargs):
    """
    Decorate a function by adding a progress bar

    Parameters
    ----------
    *args, **kwargs
        Parameters passed to tqdm.
    
    Returns
    -------
    Decorated function.
    """
    class PbarFuncDecorator(object):
        """
        The decorator class. It takes a function as an input
        and decorates it, so every call invokes tqdm an update
        in progress bar.
        """
        def __init__(self, func):
            self.func = func
            self.pbar = tqdm(*args, **kwargs)

        def __call__(self, *args, **kwargs):
            tmp = self.func(*args, **kwargs)
            self.pbar.update()
            return tmp
    
    return PbarFuncDecorator


if __name__ == "__main__":
    
    from time import sleep

    def func_iterator(func, iters):
        for _ in range(iters):                                                  
            func()

    N_ITERS = 5

    @tqdm_function_decorator(total=N_ITERS)
    def sleeper():
        sleep(1)
    
    func_iterator(sleeper, N_ITERS)