Fix tqdm_pandas(tqdm_notebook) by delaying instanciation in tqdm_pandas

lrq3000 · lrq3000 · commit ee6abc98b8f9 · 2016-07-03T23:14:15.000+02:00
Signed-off-by: Stephen L. &lt;lrq3000@gmail.com&gt;
diff --git a/README.rst b/README.rst
@@ -479,8 +479,8 @@ for ``DataFrame.progress_apply`` and ``DataFrameGroupBy.progress_apply``:
     df = pd.DataFrame(np.random.randint(0, 100, (100000, 6)))
 
     # Create and register a new `tqdm` instance with `pandas`
-    # (can use tqdm_gui, optional kwargs, etc.)
-    tqdm_pandas(tqdm())
+    # (can use tqdm_gui, tqdm_notebook, optional kwargs, etc.)
+    tqdm_pandas(tqdm, leave=True)
 
     # Now you can use `progress_apply` instead of `apply`
     df.progress_apply(lambda x: x**2)
diff --git a/tqdm/_tqdm_pandas.py b/tqdm/_tqdm_pandas.py
@@ -7,20 +7,25 @@
 __all__ = ['tqdm_pandas']
 
 
-def tqdm_pandas(t):
+def tqdm_pandas(tclass, *targs, **tkwargs):
     """
     Registers the given `tqdm` instance with
     `pandas.core.groupby.DataFrameGroupBy.progress_apply`.
     It will even close() the `tqdm` instance upon completion.
 
+    Parameters
+    ----------
+    tclass  : tqdm class you want to use (eg, tqdm, tqdm_notebook, etc)
+    targs and tkwargs  : arguments for the tqdm instance
+
     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
+    >>> tqdm_pandas(tqdm, leave=True)  # 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)
 
@@ -44,23 +49,31 @@ def inner(df, func, *args, **kwargs):
 
         *args and *kwargs are transmitted to DataFrameGroupBy.apply()
         """
-        t.total = getattr(df, 'ngroups', None)
-        if t.total is None:  # not grouped
-            t.total = len(df) if isinstance(df, Series) \
+        # 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:
-            t.total += 1  # pandas calls update once too many
+            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 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
diff --git a/tqdm/tests/tests_pandas.py b/tqdm/tests/tests_pandas.py
@@ -18,9 +18,9 @@ def test_pandas_groupby_apply():
         df = pd.DataFrame(randint(0, 50, (500, 3)))
         dfs = pd.DataFrame(randint(0, 50, (500, 3)),
                            columns=list('abc'))
-        tqdm_pandas(tqdm(file=our_file, leave=False, ascii=True))
+        tqdm_pandas(tqdm, file=our_file, leave=False, ascii=True)
         df.groupby(0).progress_apply(lambda x: None)
-        tqdm_pandas(tqdm(file=our_file, leave=False, ascii=True))
+        tqdm_pandas(tqdm, file=our_file, leave=False, ascii=True)
         dfs.groupby(['a']).progress_apply(lambda x: None)
 
         our_file.seek(0)
@@ -48,9 +48,9 @@ def test_pandas_apply():
         df = pd.DataFrame(randint(0, 50, (500, 3)))
         dfs = pd.DataFrame(randint(0, 50, (500, 3)),
                            columns=list('abc'))
-        tqdm_pandas(tqdm(file=our_file, leave=True, ascii=True))
+        tqdm_pandas(tqdm, file=our_file, leave=True, ascii=True)
         df.progress_apply(lambda x: None)
-        tqdm_pandas(tqdm(file=our_file, leave=True, ascii=True))
+        tqdm_pandas(tqdm, file=our_file, leave=True, ascii=True)
         dfs.a.progress_apply(lambda x: None)
 
         our_file.seek(0)
@@ -73,7 +73,7 @@ def test_pandas_leave():
 
     with closing(StringIO()) as our_file:
         df = pd.DataFrame(randint(0, 100, (1000, 6)))
-        tqdm_pandas(tqdm(file=our_file, leave=True, ascii=True))
+        tqdm_pandas(tqdm, file=our_file, leave=True, ascii=True)
         df.groupby(0).progress_apply(lambda x: None)
 
         our_file.seek(0)
