REF: ensure _sanitize_column returns ArrayLike (#39962)

Author: jbrockmendel

pandas/core/arrays/categorical.py

@@ -399,7 +399,9 @@ def __init__(
             # sanitize_array coerces np.nan to a string under certain versions
             # of numpy
             values = maybe_infer_to_datetimelike(values)
-            if not isinstance(values, (np.ndarray, ExtensionArray)):
+            if isinstance(values, np.ndarray):
+                values = sanitize_to_nanoseconds(values)
+            elif not isinstance(values, ExtensionArray):
                 values = com.convert_to_list_like(values)
 
                 # By convention, empty lists result in object dtype:
@@ -409,9 +411,6 @@ def __init__(
                     values = [values[idx] for idx in np.where(~null_mask)[0]]
                 values = sanitize_array(values, None, dtype=sanitize_dtype)
 
-            else:
-                values = sanitize_to_nanoseconds(values)
-
         if dtype.categories is None:
             try:
                 codes, categories = factorize(values, sort=True)

pandas/core/dtypes/cast.py

@@ -1432,7 +1432,7 @@ def maybe_infer_to_datetimelike(
     if not len(v):
         return value
 
-    def try_datetime(v):
+    def try_datetime(v: np.ndarray) -> ArrayLike:
         # safe coerce to datetime64
         try:
             # GH19671
@@ -1451,14 +1451,15 @@ def try_datetime(v):
             except (ValueError, TypeError):
                 pass
             else:
-                return DatetimeIndex(values).tz_localize("UTC").tz_convert(tz=tz)
+                dti = DatetimeIndex(values).tz_localize("UTC").tz_convert(tz=tz)
+                return dti._data
         except TypeError:
             # e.g. <class 'numpy.timedelta64'> is not convertible to datetime
             pass
 
         return v.reshape(shape)
 
-    def try_timedelta(v):
+    def try_timedelta(v: np.ndarray) -> np.ndarray:
         # safe coerce to timedelta64
 
         # will try first with a string & object conversion

pandas/core/frame.py

@@ -4000,7 +4000,7 @@ def assign(self, **kwargs) -> DataFrame:
             data[k] = com.apply_if_callable(v, data)
         return data
 
-    def _sanitize_column(self, value):
+    def _sanitize_column(self, value) -> ArrayLike:
         """
         Ensures new columns (which go into the BlockManager as new blocks) are
         always copied and converted into an array.
@@ -4011,7 +4011,7 @@ def _sanitize_column(self, value):
 
         Returns
         -------
-        numpy.ndarray
+        numpy.ndarray or ExtensionArray
         """
         self._ensure_valid_index(value)
 
@@ -4025,7 +4025,7 @@ def _sanitize_column(self, value):
             value = value.copy()
             value = sanitize_index(value, self.index)
 
-        elif isinstance(value, Index) or is_sequence(value):
+        elif is_sequence(value):
 
             # turn me into an ndarray
             value = sanitize_index(value, self.index)
@@ -4035,7 +4035,7 @@ def _sanitize_column(self, value):
                 else:
                     value = com.asarray_tuplesafe(value)
             elif isinstance(value, Index):
-                value = value.copy(deep=True)
+                value = value.copy(deep=True)._values
             else:
                 value = value.copy()