# future division is important to divide integers and get as
# a result precise floating numbers (instead of truncated int)
from __future__ import absolute_import
__author__ = "github.com/casperdcl"
__all__ = ['tqdm_pandas']
[docs]def tqdm_pandas(t): # pragma: no cover
"""
Registers the given `tqdm` instance with
`pandas.core.groupby.DataFrameGroupBy.progress_apply`.
It will even close() the `tqdm` instance upon completion.
Examples
--------
>>> import pandas as pd
>>> import numpy as np
>>> from tqdm import tqdm, tqdm_pandas
>>>
>>> df = pd.DataFrame(np.random.randint(0, 100, (100000, 6)))
>>> tqdm_pandas(tqdm()) # can use tqdm_gui, optional kwargs, etc
>>> # Now you can use `progress_apply` instead of `apply`
>>> df.groupby(0).progress_apply(lambda x: x**2)
References
----------
https://stackoverflow.com/questions/18603270/
progress-indicator-during-pandas-operations-python
"""
from pandas.core.groupby import DataFrameGroupBy
def inner(groups, func, *args, **kwargs):
"""
Parameters
----------
groups : DataFrameGroupBy
Grouped data.
func : function
To be applied on the grouped data.
*args and *kwargs are transmitted to DataFrameGroupBy.apply()
"""
t.total = len(groups) + 1 # pandas calls update once too many
def wrapper(*args, **kwargs):
t.update()
return func(*args, **kwargs)
result = groups.apply(wrapper, *args, **kwargs)
t.close()
return result
# Enable custom tqdm progress in pandas!
DataFrameGroupBy.progress_apply = inner