Add types for pandas.DataFrame.map

Author: andrewpmk

pandas-stubs/core/frame.pyi

@@ -1316,9 +1316,14 @@ class DataFrame(NDFrame, OpsMixin):
     ) -> DataFrame: ...
 
     # Add spacing between apply() overloads and remaining annotations
+    # Deprecated in pandas 2.1.0 and later
     def applymap(
         self, func: Callable, na_action: Literal["ignore"] | None = ..., **kwargs
     ) -> DataFrame: ...
+    # In pandas 2.1.0 and later: applymap is renamed map, applymap is deprecated
+    def map(
+        self, func: Callable, na_action: Literal["ignore"] | None = ..., **kwargs
+    ) -> DataFrame: ...   
     def join(
         self,
         other: DataFrame | Series | list[DataFrame | Series],

tests/test_frame.py

@@ -728,6 +728,7 @@ def gethead(s: pd.Series, y: int) -> pd.Series:
     )
 
 
+# Deprecated in pandas 2.1.0
 def test_types_applymap() -> None:
     with pytest_warns_bounded(
         FutureWarning, "DataFrame.applymap has been deprecated", lower="2.0.99"
@@ -741,6 +742,17 @@ def test_types_applymap() -> None:
         df.applymap(str, na_action=None)
 
 
+# Added in pandas 2.1.0
+def test_types_map() -> None:
+    df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
+    df.map(lambda x: x**2)
+    df.map(np.exp)
+    df.map(str)
+    # na_action parameter was added in 1.2.0 https://pandas.pydata.org/docs/whatsnew/v1.2.0.html
+    df.map(np.exp, na_action="ignore")
+    df.map(str, na_action=None)
+
+
 def test_types_element_wise_arithmetic() -> None:
     df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
     df2 = pd.DataFrame(data={"col1": [10, 20], "col3": [3, 4]})