make pandas a classmethod. TODO: Readme update, unit tests (incl subclasses)

casperdcl · casperdcl · commit f020419a2bf4 · 2016-07-18T22:52:52.000+01:00
diff --git a/tqdm/_tqdm.py b/tqdm/_tqdm.py
@@ -324,6 +324,86 @@ def write(cls, s, file=sys.stdout, end="\n"):
             inst.refresh()
         # TODO: make list of all instances incl. absolutely positioned ones?
 
+    @classmethod
+    def pandas(tclass, *targs, **tkwargs):
+        """
+        Registers the given `tqdm` class with
+            pandas.core.
+            ( frame.DataFrame
+            | series.Series
+            | groupby.DataFrameGroupBy
+            | groupby.SeriesGroupBy
+            ).progress_apply
+
+        A new instance will be create every time `progress_apply` is called,
+        and each instance will automatically close() upon completion.
+
+        Parameters
+        ----------
+        targs, tkwargs  : arguments for the tqdm instance
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import numpy as np
+        >>> from tqdm import tqdm, tqdm_gui
+        >>>
+        >>> df = pd.DataFrame(np.random.randint(0, 100, (100000, 6)))
+        >>> tqdm.pandas(ncols=50)  # 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.frame import DataFrame
+        from pandas.core.series import Series
+        from pandas.core.groupby import DataFrameGroupBy, SeriesGroupBy
+
+        def inner(df, func, *args, **kwargs):
+            """
+            Parameters
+            ----------
+            df  : (DataFrame|Series)[GroupBy]
+                Data (may be grouped).
+            func  : function
+                To be applied on the (grouped) data.
+            *args, *kwargs  : optional
+                Transmitted to `df.apply()`.
+            """
+            # Precompute total iterations
+            total = getattr(df, 'ngroups', None)
+            if total is None:  # not grouped
+                total = len(df) if isinstance(df, Series) \
+                    else df.size // len(df)
+            else:
+                total += 1  # pandas calls update once too many
+
+            # Init bar
+            t = tclass(*targs, total=total, **tkwargs)
+
+            # Define bar updating wrapper
+            def wrapper(*args, **kwargs):
+                t.update()
+                return func(*args, **kwargs)
+
+            # Apply the provided function (in *args and **kwargs)
+            # on the df using our wrapper (which provides bar updating)
+            result = df.apply(wrapper, *args, **kwargs)
+
+            # Close bar and return pandas calculation result
+            t.close()
+            return result
+
+        # Monkeypatch pandas to provide easy methods
+        # Enable custom tqdm progress in pandas!
+        DataFrame.progress_apply = inner
+        DataFrameGroupBy.progress_apply = inner
+        Series.progress_apply = inner
+        SeriesGroupBy.progress_apply = inner
+
     def __init__(self, iterable=None, desc=None, total=None, leave=True,
                  file=sys.stderr, ncols=None, mininterval=0.1,
                  maxinterval=10.0, miniters=None, ascii=None, disable=False,
diff --git a/tqdm/_tqdm_pandas.py b/tqdm/_tqdm_pandas.py
@@ -34,6 +34,18 @@ def tqdm_pandas(tclass, *targs, **tkwargs):
     https://stackoverflow.com/questions/18603270/
     progress-indicator-during-pandas-operations-python
     """
+
+    if isinstance(tclass, type) or (getattr(tclass, '__name__', '').startswith(
+            'tqdm_')):  # delayed adapter case
+        tclass.write("Warning: `tqdm_pandas(tqdm(...))` is deprecated,"
+                     " please use `tqdm.pandas(...)` instead.")
+        tclass.pandas(*targs, **tkwargs)
+        return
+
+    t = tclass
+    t.write("Warning: tqdm_pandas: using a bar instance is deprecated,"
+            " please provide a bar class instead.", file=t.fp)
+
     from pandas.core.frame import DataFrame
     from pandas.core.series import Series
     from pandas.core.groupby import DataFrameGroupBy, SeriesGroupBy
@@ -49,24 +61,14 @@ def inner(df, func, *args, **kwargs):
 
         *args and *kwargs are transmitted to DataFrameGroupBy.apply()
         """
+
         # Precompute total iterations
         total = getattr(df, 'ngroups', None)
         if total is None:  # not grouped
-            total = len(df) if isinstance(df, Series) \
+            t.total = len(df) if isinstance(df, Series) \
                 else df.size // len(df)
         else:
-            total += 1  # pandas calls update once too many
-
-        # Init bar
-        if isinstance(tclass, type) or \
-            (hasattr(tclass, '__name__') and
-             tclass.__name__.startswith('tqdm_')):  # delayed adapter case
-            t = tclass(*targs, total=total, **tkwargs)
-        else:
-            t = tclass
-            t.total = total
-            t.write("Warning: tqdm_pandas: using a bar instance is deprecated,"
-                    " please provide a bar class instead.", file=t.fp)
+            t.total = total + 1  # pandas calls update once too many
 
         # Define bar updating wrapper
         def wrapper(*args, **kwargs):
