tqdm/tqdm

Add wrapper for functions that are called repeatedly

Opened this issue · 0 comments

Proposed enhancement

Add a decorator that wraps a function in a progress bar when you expect a specific number of calls to that function.

Use case

As a specific example I'd like to look at creating animations with matplotlib. The workflow is roughly to create a figure and then providing a method to update the figure each frame. When creating the animation matplotlib repeatedly calls the update method. Matplotlib does not support a progress bar as far as I know, however, if we wrap the update method ourselves we could still provide the user with feedback on the progress of their animation.

Proposed implementation

def progress_bar_decorator(*tqdm_args, **tqdm_kwargs):
    progress_bar = tqdm(*tqdm_args, **tqdm_kwargs)

    def decorator(f):
        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            result = f(*args, **kwargs)
            progress_bar.update()
            return result

        return wrapper

    return decorator

In the example of the matplotlib animation it can then be used as such:

figure, axes = plt.subplots(1, 1)

scatter = axes.scatter(X[0], Y[0])

@progress_bar_decorator(total=len(X), desc="Generating animation")
def update(frame):
      nonlocal scatter
      
      scatter.set_offsets(X[frame], Y[frame])

      return scatter

FuncAnimation(figure, update, frames=len(X), ...).save("file.mp4")

Matplotlib will call update exactly len(X) times causing tqdm to update the progress bar accordingly. In this example X and Y are some fictional data arrays.