REF: implement Categorical._box_func, make _box_func a method (#36206)

Author: jbrockmendel

pandas/core/arrays/categorical.py

@@ -1734,6 +1734,11 @@ def _ndarray(self) -> np.ndarray:
     def _from_backing_data(self, arr: np.ndarray) -> "Categorical":
         return self._constructor(arr, dtype=self.dtype, fastpath=True)
 
+    def _box_func(self, i: int):
+        if i == -1:
+            return np.NaN
+        return self.categories[i]
+
     # ------------------------------------------------------------------
 
     def take_nd(self, indexer, allow_fill: bool = False, fill_value=None):
@@ -1874,10 +1879,7 @@ def __getitem__(self, key):
         """
         if isinstance(key, (int, np.integer)):
             i = self._codes[key]
-            if i == -1:
-                return np.nan
-            else:
-                return self.categories[i]
+            return self._box_func(i)
 
         key = check_array_indexer(self, key)
 

pandas/core/arrays/datetimelike.py

@@ -478,8 +478,7 @@ def _from_backing_data(self: _T, arr: np.ndarray) -> _T:
 
     # ------------------------------------------------------------------
 
-    @property
-    def _box_func(self):
+    def _box_func(self, x):
         """
         box function to get object from internal representation
         """

pandas/core/arrays/datetimes.py

@@ -7,6 +7,7 @@
 from pandas._libs import lib, tslib
 from pandas._libs.tslibs import (
     NaT,
+    NaTType,
     Resolution,
     Timestamp,
     conversion,
@@ -475,9 +476,8 @@ def _maybe_clear_freq(self):
     # -----------------------------------------------------------------
     # Descriptive Properties
 
-    @property
-    def _box_func(self):
-        return lambda x: Timestamp(x, freq=self.freq, tz=self.tz)
+    def _box_func(self, x) -> Union[Timestamp, NaTType]:
+        return Timestamp(x, freq=self.freq, tz=self.tz)
 
     @property
     def dtype(self) -> Union[np.dtype, DatetimeTZDtype]:

pandas/core/arrays/period.py

@@ -484,9 +484,8 @@ def _time_shift(self, periods, freq=None):
             values[self._isnan] = iNaT
         return type(self)(values, freq=self.freq)
 
-    @property
-    def _box_func(self):
-        return lambda x: Period._from_ordinal(ordinal=x, freq=self.freq)
+    def _box_func(self, x) -> Union[Period, NaTType]:
+        return Period._from_ordinal(ordinal=x, freq=self.freq)
 
     def asfreq(self, freq=None, how: str = "E") -> "PeriodArray":
         """

pandas/core/arrays/timedeltas.py

@@ -1,10 +1,19 @@
 from datetime import timedelta
-from typing import List
+from typing import List, Union
 
 import numpy as np
 
 from pandas._libs import lib, tslibs
-from pandas._libs.tslibs import NaT, Period, Tick, Timedelta, Timestamp, iNaT, to_offset
+from pandas._libs.tslibs import (
+    NaT,
+    NaTType,
+    Period,
+    Tick,
+    Timedelta,
+    Timestamp,
+    iNaT,
+    to_offset,
+)
 from pandas._libs.tslibs.conversion import precision_from_unit
 from pandas._libs.tslibs.fields import get_timedelta_field
 from pandas._libs.tslibs.timedeltas import array_to_timedelta64, parse_timedelta_unit
@@ -108,9 +117,8 @@ class TimedeltaArray(dtl.DatetimeLikeArrayMixin, dtl.TimelikeOps):
     # Note: ndim must be defined to ensure NaT.__richcmp(TimedeltaArray)
     #  operates pointwise.
 
-    @property
-    def _box_func(self):
-        return lambda x: Timedelta(x, unit="ns")
+    def _box_func(self, x) -> Union[Timedelta, NaTType]:
+        return Timedelta(x, unit="ns")
 
     @property
     def dtype(self):

pandas/core/indexes/datetimelike.py

@@ -81,7 +81,7 @@ def wrapper(left, right):
     DatetimeLikeArrayMixin,
     cache=True,
 )
-@inherit_names(["mean", "asi8", "freq", "freqstr", "_box_func"], DatetimeLikeArrayMixin)
+@inherit_names(["mean", "asi8", "freq", "freqstr"], DatetimeLikeArrayMixin)
 class DatetimeIndexOpsMixin(ExtensionIndex):
     """
     Common ops mixin to support a unified interface datetimelike Index.
@@ -244,7 +244,7 @@ def min(self, axis=None, skipna=True, *args, **kwargs):
             # quick check
             if len(i8) and self.is_monotonic:
                 if i8[0] != iNaT:
-                    return self._box_func(i8[0])
+                    return self._data._box_func(i8[0])
 
             if self.hasnans:
                 if skipna:
@@ -253,7 +253,7 @@ def min(self, axis=None, skipna=True, *args, **kwargs):
                     return self._na_value
             else:
                 min_stamp = i8.min()
-            return self._box_func(min_stamp)
+            return self._data._box_func(min_stamp)
         except ValueError:
             return self._na_value
 
@@ -301,7 +301,7 @@ def max(self, axis=None, skipna=True, *args, **kwargs):
             # quick check
             if len(i8) and self.is_monotonic:
                 if i8[-1] != iNaT:
-                    return self._box_func(i8[-1])
+                    return self._data._box_func(i8[-1])
 
             if self.hasnans:
                 if skipna:
@@ -310,7 +310,7 @@ def max(self, axis=None, skipna=True, *args, **kwargs):
                     return self._na_value
             else:
                 max_stamp = i8.max()
-            return self._box_func(max_stamp)
+            return self._data._box_func(max_stamp)
         except ValueError:
             return self._na_value