REF: simplify .map method for datetime-likes

Author: topper-123

doc/source/whatsnew/v2.1.0.rst

@@ -133,12 +133,14 @@ Categorical
 Datetimelike
 ^^^^^^^^^^^^
 - Bug in :meth:`Timestamp.round` with values close to the implementation bounds returning incorrect results instead of raising ``OutOfBoundsDatetime`` (:issue:`51494`)
+- :meth:`arrays.DatetimeArray.map` can now take a ``na_action`` argument. :meth:`DatetimeIndex.map` with ``na_action="ignore"`` now works as expected. (:issue:`51644`)
 -
 
 Timedelta
 ^^^^^^^^^
 - Bug in :meth:`Timedelta.round` with values close to the implementation bounds returning incorrect results instead of raising ``OutOfBoundsTimedelta`` (:issue:`51494`)
 - Bug in :class:`TimedeltaIndex` division or multiplication leading to ``.freq`` of "0 Days" instead of ``None`` (:issue:`51575`)
+- :meth:`arrays.TimedeltaArray.map` can now take a ``na_action`` argument. :meth:`TimedeltaIndex.map` with ``na_action="ignore"`` now works as expected. (:issue:`51644`)
 -
 
 Timezones
@@ -190,6 +192,7 @@ Period
 ^^^^^^
 - Bug in :class:`PeriodDtype` constructor failing to raise ``TypeError`` when no argument is passed or when ``None`` is passed (:issue:`27388`)
 - Bug in :class:`PeriodDtype` constructor raising ``ValueError`` instead of ``TypeError`` when an invalid type is passed (:issue:`51790`)
+- :meth:`arrays.PeriodArray.map` can now take a ``na_action`` argument. :meth:`PeriodIndex.map` with ``na_action="ignore"`` now works as expected. (:issue:`51644`)
 -
 
 Plotting

pandas/core/algorithms.py

@@ -1750,4 +1750,4 @@ def map_array(
     if na_action is None:
         return lib.map_infer(values, mapper)
     else:
-        return lib.map_infer_mask(values, mapper, isna(values).view(np.uint8))
+        return lib.map_infer_mask(values, mapper, mask=isna(values).view(np.uint8))

pandas/core/arrays/datetimelike.py

@@ -119,6 +119,7 @@
 from pandas.core.algorithms import (
     checked_add_with_arr,
     isin,
+    map_array,
     unique1d,
 )
 from pandas.core.array_algos import datetimelike_accumulations
@@ -751,17 +752,34 @@ def _unbox(self, other) -> np.int64 | np.datetime64 | np.timedelta64 | np.ndarra
 
     @ravel_compat
     def map(self, mapper, na_action=None):
-        if na_action is not None:
-            raise NotImplementedError
-
         # TODO(GH-23179): Add ExtensionArray.map
         # Need to figure out if we want ExtensionArray.map first.
         # If so, then we can refactor IndexOpsMixin._map_values to
         # a standalone function and call from here..
         # Else, just rewrite _map_infer_values to do the right thing.
         from pandas import Index
 
-        return Index(self).map(mapper).array
+        idx = Index(self)
+
+        try:
+            if na_action is not None:
+                raise ValueError("calling mapper directly only na_action=None")
+            result = mapper(idx)
+
+            # Try to use this result if we can
+            if isinstance(result, np.ndarray):
+                result = Index(result)
+
+            if not isinstance(result, Index):
+                raise TypeError("The map function must return an Index object")
+        except Exception:
+            result = map_array(self, mapper, na_action=na_action)
+            result = Index(result)
+
+        if isinstance(result, ABCMultiIndex):
+            return result.to_numpy()
+        else:
+            return result.array
 
     def isin(self, values) -> npt.NDArray[np.bool_]:
         """

pandas/core/base.py

@@ -881,14 +881,9 @@ def _map_values(self, mapper, na_action=None):
         """
         arr = extract_array(self, extract_numpy=True, extract_range=True)
 
-        if is_extension_array_dtype(arr.dtype):
-            # Item "IndexOpsMixin" of "Union[IndexOpsMixin, ExtensionArray,
-            # ndarray[Any, Any]]" has no attribute "map"
-            return arr.map(mapper, na_action=na_action)  # type: ignore[union-attr]
-
-        # Argument 1 to "map_array" has incompatible type
-        # "Union[IndexOpsMixin, ExtensionArray, ndarray[Any, Any]]";
-        # expected "Union[ExtensionArray, ndarray[Any, Any]]"
+        if isinstance(arr, ExtensionArray):
+            return arr.map(mapper, na_action=na_action)
+
         return algorithms.map_array(
             arr, mapper, na_action=na_action  # type: ignore[arg-type]
         )

pandas/core/indexes/extension.py

@@ -9,18 +9,15 @@
     TypeVar,
 )
 
-import numpy as np
-
-from pandas.util._decorators import (
-    cache_readonly,
-    doc,
-)
+from pandas.util._decorators import cache_readonly
 
 from pandas.core.dtypes.generic import ABCDataFrame
 
 from pandas.core.indexes.base import Index
 
 if TYPE_CHECKING:
+    import numpy as np
+
     from pandas._typing import (
         ArrayLike,
         npt,
@@ -154,23 +151,6 @@ def _validate_fill_value(self, value):
         """
         return self._data._validate_setitem_value(value)
 
-    @doc(Index.map)
-    def map(self, mapper, na_action=None):
-        # Try to run function on index first, and then on elements of index
-        # Especially important for group-by functionality
-        try:
-            result = mapper(self)
-
-            # Try to use this result if we can
-            if isinstance(result, np.ndarray):
-                result = Index(result)
-
-            if not isinstance(result, Index):
-                raise TypeError("The map function must return an Index object")
-            return result
-        except Exception:
-            return self.astype(object).map(mapper)
-
     @cache_readonly
     def _isnan(self) -> npt.NDArray[np.bool_]:
         # error: Incompatible return value type (got "ExtensionArray", expected

pandas/tests/apply/test_invalid_arg.py

@@ -83,13 +83,6 @@ def test_map_categorical_na_action():
         s.map(lambda x: x, na_action="ignore")
 
 
-def test_map_datetimetz_na_action():
-    values = date_range("2011-01-01", "2011-01-02", freq="H").tz_localize("Asia/Tokyo")
-    s = Series(values, name="XX")
-    with pytest.raises(NotImplementedError, match=tm.EMPTY_STRING_PATTERN):
-        s.map(lambda x: x, na_action="ignore")
-
-
 @pytest.mark.parametrize("method", ["apply", "agg", "transform"])
 @pytest.mark.parametrize("func", [{"A": {"B": "sum"}}, {"A": {"B": ["sum"]}}])
 def test_nested_renamer(frame_or_series, method, func):

pandas/tests/extension/test_datetime.py

@@ -118,12 +118,8 @@ def test_combine_add(self, data_repeated):
 
     @pytest.mark.parametrize("na_action", [None, "ignore"])
     def test_map(self, data, na_action):
-        if na_action is not None:
-            with pytest.raises(NotImplementedError, match=""):
-                data.map(lambda x: x, na_action=na_action)
-        else:
-            result = data.map(lambda x: x, na_action=na_action)
-            self.assert_extension_array_equal(result, data)
+        result = data.map(lambda x: x, na_action=na_action)
+        self.assert_extension_array_equal(result, data)
 
 
 class TestInterface(BaseDatetimeTests, base.BaseInterfaceTests):

pandas/tests/extension/test_period.py

@@ -107,12 +107,8 @@ def test_diff(self, data, periods):
 
     @pytest.mark.parametrize("na_action", [None, "ignore"])
     def test_map(self, data, na_action):
-        if na_action is not None:
-            with pytest.raises(NotImplementedError, match=""):
-                data.map(lambda x: x, na_action=na_action)
-        else:
-            result = data.map(lambda x: x, na_action=na_action)
-            self.assert_extension_array_equal(result, data)
+        result = data.map(lambda x: x, na_action=na_action)
+        self.assert_extension_array_equal(result, data)
 
 
 class TestInterface(BasePeriodTests, base.BaseInterfaceTests):