Progress bar wrapper for pandas apply

[sandstromviktor]

Hi. I really like making nice looking progress bars using rich. Thanks to the creators of this fine piece of art.
Sometimes when using Pandas, you will use the apply function. This can take a long time if you have a large DataFrame. It is possible to use tqdm and their progress_apply monkey patch to get a progress bar for this operation. However, i like rich more, so i created a similar wrapper that i want to share for anyone looking at this in the future.

It is quite simple. First, define this function

from rich.progress import Progress


def inner_generator(df_function='apply'):
    def inner(df, func, **kwargs):

        progress = Progress()
        task = progress.add_task("Processing...", total=len(df))
        def wrapper(*args, **kwargs):
            result = func(*args, **kwargs)
            progress.update(task, advance=1)
            return result

        with progress:
            return getattr(df, df_function)(wrapper, **kwargs)
    return inner

then monkey patch this to Pandas like this

import pandas as pd
pd.DataFrame.apply_rich = inner_generator() 
pd.Series.apply_rich = inner_generator()

and use it like this

df2 = df.apply_rich(my_function, axis=1)

You could also wrap the map function by specifying inner_generator(df_function="map")