pandas: add progress_map, fix tests

mocitanefelibata · Feb 5, 2024 · a95cd88 · a95cd88
1 parent 123cb1c
commit a95cd88
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 5 deletions.
diff --git a/tests/tests_pandas.py b/tests/tests_pandas.py
@@ -4,6 +4,7 @@
 
 pytestmark = mark.slow
 
+np = importorskip('numpy')
 random = importorskip('numpy.random')
 rand = random.rand
 randint = random.randint
@@ -66,6 +67,7 @@ def test_pandas_series():
                     f"\nExpected:\n{exres} at least twice.\nIn:\n{our_file.read()}\n")
 
 
+@mark.filterwarnings("ignore:DataFrame.applymap has been deprecated:FutureWarning")
 def test_pandas_data_frame():
     """Test pandas.DataFrame.progress_apply and .progress_applymap"""
     with closing(StringIO()) as our_file:
@@ -80,6 +82,12 @@ def task_func(x):
         res2 = df.applymap(task_func)
         assert res1.equals(res2)
 
+        # map
+        if hasattr(df, 'map'):  # pandas>=2.1.0
+            res1 = df.progress_map(task_func)
+            res2 = df.map(task_func)
+            assert res1.equals(res2)
+
         # apply unhashable
         res1 = []
         df.progress_apply(res1.extend)
@@ -119,8 +127,8 @@ def test_pandas_groupby_apply():
         dfs.groupby(['a']).progress_apply(lambda x: None)
 
         df2 = df = pd.DataFrame({'a': randint(1, 8, 10000), 'b': rand(10000)})
-        res1 = df2.groupby("a").apply(max)
-        res2 = df2.groupby("a").progress_apply(max)
+        res1 = df2.groupby("a").apply(np.maximum.reduce)
+        res2 = df2.groupby("a").progress_apply(np.maximum.reduce)
         assert res1.equals(res2)
 
         our_file.seek(0)
@@ -139,11 +147,11 @@ def test_pandas_groupby_apply():
         dfs.loc[0] = [2, 1, 1]
         dfs['d'] = 100
 
-        expects = ['500/500', '1/1', '4/4', '2/2']
+        expects = ['500/500', '1/1', '4/4', '4/4']
         dfs.groupby(dfs.index).progress_apply(lambda x: None)
         dfs.groupby('d').progress_apply(lambda x: None)
-        dfs.groupby(dfs.columns, axis=1).progress_apply(lambda x: None)
-        dfs.groupby([2, 2, 1, 1], axis=1).progress_apply(lambda x: None)
+        dfs.T.groupby(dfs.columns).progress_apply(lambda x: None)
+        dfs.T.groupby([2, 2, 1, 1]).progress_apply(lambda x: None)
 
         our_file.seek(0)
         if our_file.read().count('100%') < 4:

diff --git a/tqdm/std.py b/tqdm/std.py
@@ -930,6 +930,8 @@ def wrapper(*args, **kwargs):
         DataFrame.progress_apply = inner_generator()
         DataFrameGroupBy.progress_apply = inner_generator()
         DataFrame.progress_applymap = inner_generator('applymap')
+        DataFrame.progress_map = inner_generator('map')
+        DataFrameGroupBy.progress_map = inner_generator('map')
 
         if Panel is not None:
             Panel.progress_apply = inner_generator()