REF: get_supported_reso->get_supported_dtype (#56439)

Author: jbrockmendel

pandas/_libs/tslibs/__init__.py

@@ -30,20 +30,16 @@
  "get_unit_from_dtype",
  "periods_per_day",
  "periods_per_second",
- "is_supported_unit",
- "npy_unit_to_abbrev",
- "get_supported_reso",
  "guess_datetime_format",
  "add_overflowsafe",
+ "get_supported_dtype",
+ "is_supported_dtype",
 ]
 
 from pandas._libs.tslibs import dtypes # pylint: disable=import-self
 from pandas._libs.tslibs.conversion import localize_pydatetime
 from pandas._libs.tslibs.dtypes import (
  Resolution,
- get_supported_reso,
- is_supported_unit,
- npy_unit_to_abbrev,
  periods_per_day,
  periods_per_second,
 )
@@ -58,6 +54,8 @@
  OutOfBoundsTimedelta,
  add_overflowsafe,
  astype_overflowsafe,
+ get_supported_dtype,
+ is_supported_dtype,
  is_unitless,
  py_get_unit_from_dtype as get_unit_from_dtype,
 )

pandas/_libs/tslibs/dtypes.pxd

@@ -3,13 +3,13 @@ from numpy cimport int64_t
 from pandas._libs.tslibs.np_datetime cimport NPY_DATETIMEUNIT
 
 
-cpdef str npy_unit_to_abbrev(NPY_DATETIMEUNIT unit)
+cdef str npy_unit_to_abbrev(NPY_DATETIMEUNIT unit)
 cpdef NPY_DATETIMEUNIT abbrev_to_npy_unit(str abbrev)
 cdef NPY_DATETIMEUNIT freq_group_code_to_npy_unit(int freq) noexcept nogil
 cpdef int64_t periods_per_day(NPY_DATETIMEUNIT reso=*) except? -1
 cpdef int64_t periods_per_second(NPY_DATETIMEUNIT reso) except? -1
-cpdef NPY_DATETIMEUNIT get_supported_reso(NPY_DATETIMEUNIT reso)
-cpdef bint is_supported_unit(NPY_DATETIMEUNIT reso)
+cdef NPY_DATETIMEUNIT get_supported_reso(NPY_DATETIMEUNIT reso)
+cdef bint is_supported_unit(NPY_DATETIMEUNIT reso)
 
 cpdef freq_to_period_freqstr(freq_n, freq_name)
 cdef dict c_OFFSET_TO_PERIOD_FREQSTR

pandas/_libs/tslibs/dtypes.pyi

@@ -4,9 +4,6 @@ OFFSET_TO_PERIOD_FREQSTR: dict[str, str]
 
 def periods_per_day(reso: int = ...) -> int: ...
 def periods_per_second(reso: int) -> int: ...
-def is_supported_unit(reso: int) -> bool: ...
-def npy_unit_to_abbrev(unit: int) -> str: ...
-def get_supported_reso(reso: int) -> int: ...
 def abbrev_to_npy_unit(abbrev: str) -> int: ...
 def freq_to_period_freqstr(freq_n: int, freq_name: str) -> str: ...
 

pandas/_libs/tslibs/dtypes.pyx

@@ -559,7 +559,7 @@ class NpyDatetimeUnit(Enum):
  NPY_FR_GENERIC = NPY_DATETIMEUNIT.NPY_FR_GENERIC
 
 
-cpdef NPY_DATETIMEUNIT get_supported_reso(NPY_DATETIMEUNIT reso):
+cdef NPY_DATETIMEUNIT get_supported_reso(NPY_DATETIMEUNIT reso):
  # If we have an unsupported reso, return the nearest supported reso.
  if reso == NPY_DATETIMEUNIT.NPY_FR_GENERIC:
  # TODO: or raise ValueError? trying this gives unraisable errors, but
@@ -572,7 +572,7 @@ cpdef NPY_DATETIMEUNIT get_supported_reso(NPY_DATETIMEUNIT reso):
  return reso
 
 
-cpdef bint is_supported_unit(NPY_DATETIMEUNIT reso):
+cdef bint is_supported_unit(NPY_DATETIMEUNIT reso):
  return (
  reso == NPY_DATETIMEUNIT.NPY_FR_ns
  or reso == NPY_DATETIMEUNIT.NPY_FR_us
@@ -581,7 +581,7 @@ cpdef bint is_supported_unit(NPY_DATETIMEUNIT reso):
  )
 
 
-cpdef str npy_unit_to_abbrev(NPY_DATETIMEUNIT unit):
+cdef str npy_unit_to_abbrev(NPY_DATETIMEUNIT unit):
  if unit == NPY_DATETIMEUNIT.NPY_FR_ns or unit == NPY_DATETIMEUNIT.NPY_FR_GENERIC:
  # generic -> default to nanoseconds
  return "ns"

pandas/_libs/tslibs/np_datetime.pyi

@@ -23,3 +23,5 @@ def add_overflowsafe(
  left: npt.NDArray[np.int64],
  right: npt.NDArray[np.int64],
 ) -> npt.NDArray[np.int64]: ...
+def get_supported_dtype(dtype: np.dtype) -> np.dtype: ...
+def is_supported_dtype(dtype: np.dtype) -> bool: ...

pandas/_libs/tslibs/np_datetime.pyx

@@ -39,6 +39,8 @@ from numpy cimport (
 )
 
 from pandas._libs.tslibs.dtypes cimport (
+ get_supported_reso,
+ is_supported_unit,
  npy_unit_to_abbrev,
  npy_unit_to_attrname,
 )
@@ -91,6 +93,28 @@ def py_get_unit_from_dtype(dtype):
  return get_unit_from_dtype(dtype)
 
 
+def get_supported_dtype(dtype: cnp.dtype) -> cnp.dtype:
+ reso = get_unit_from_dtype(dtype)
+ new_reso = get_supported_reso(reso)
+ new_unit = npy_unit_to_abbrev(new_reso)
+
+ # Accessing dtype.kind here incorrectly(?) gives "" instead of "m"/"M",
+ # so we check type_num instead
+ if dtype.type_num == cnp.NPY_DATETIME:
+ new_dtype = np.dtype(f"M8[{new_unit}]")
+ else:
+ new_dtype = np.dtype(f"m8[{new_unit}]")
+ return new_dtype
+
+
+def is_supported_dtype(dtype: cnp.dtype) -> bool:
+ if dtype.type_num not in [cnp.NPY_DATETIME, cnp.NPY_TIMEDELTA]:
+ raise ValueError("is_unitless dtype must be datetime64 or timedelta64")
+ cdef:
+ NPY_DATETIMEUNIT unit = get_unit_from_dtype(dtype)
+ return is_supported_unit(unit)
+
+
 def is_unitless(dtype: cnp.dtype) -> bool:
  """
  Check if a datetime64 or timedelta64 dtype has no attached unit.

pandas/core/arrays/_mixins.py

@@ -13,10 +13,7 @@
 
 from pandas._libs import lib
 from pandas._libs.arrays import NDArrayBacked
-from pandas._libs.tslibs import (
- get_unit_from_dtype,
- is_supported_unit,
-)
+from pandas._libs.tslibs import is_supported_dtype
 from pandas._typing import (
  ArrayLike,
  AxisInt,
@@ -141,16 +138,12 @@ def view(self, dtype: Dtype | None = None) -> ArrayLike:
  cls = dtype.construct_array_type() # type: ignore[assignment]
  dt64_values = arr.view(f"M8[{dtype.unit}]")
  return cls(dt64_values, dtype=dtype)
- elif lib.is_np_dtype(dtype, "M") and is_supported_unit(
- get_unit_from_dtype(dtype)
- ):
+ elif lib.is_np_dtype(dtype, "M") and is_supported_dtype(dtype):
  from pandas.core.arrays import DatetimeArray
 
  dt64_values = arr.view(dtype)
  return DatetimeArray(dt64_values, dtype=dtype)
- elif lib.is_np_dtype(dtype, "m") and is_supported_unit(
- get_unit_from_dtype(dtype)
- ):
+ elif lib.is_np_dtype(dtype, "m") and is_supported_dtype(dtype):
  from pandas.core.arrays import TimedeltaArray
 
  td64_values = arr.view(dtype)

pandas/core/arrays/datetimes.py

@@ -27,14 +27,13 @@
  astype_overflowsafe,
  fields,
  get_resolution,
- get_supported_reso,
+ get_supported_dtype,
  get_unit_from_dtype,
  ints_to_pydatetime,
  is_date_array_normalized,
- is_supported_unit,
+ is_supported_dtype,
  is_unitless,
  normalize_i8_timestamps,
- npy_unit_to_abbrev,
  timezones,
  to_offset,
  tz_convert_from_utc,
@@ -712,7 +711,7 @@ def astype(self, dtype, copy: bool = True):
  self.tz is None
  and lib.is_np_dtype(dtype, "M")
  and not is_unitless(dtype)
- and is_supported_unit(get_unit_from_dtype(dtype))
+ and is_supported_dtype(dtype)
  ):
  # unit conversion e.g. datetime64[s]
  res_values = astype_overflowsafe(self._ndarray, dtype, copy=True)
@@ -2307,7 +2306,7 @@ def _sequence_to_dt64(
  assert isinstance(result, np.ndarray), type(result)
  assert result.dtype.kind == "M"
  assert result.dtype != "M8"
- assert is_supported_unit(get_unit_from_dtype(result.dtype))
+ assert is_supported_dtype(result.dtype)
  return result, tz
 
 
@@ -2321,14 +2320,10 @@ def _construct_from_dt64_naive(
  # lib.is_np_dtype(data.dtype)
 
  new_dtype = data.dtype
- data_unit = get_unit_from_dtype(new_dtype)
- if not is_supported_unit(data_unit):
+ if not is_supported_dtype(new_dtype):
  # Cast to the nearest supported unit, generally "s"
- new_reso = get_supported_reso(data_unit)
- new_unit = npy_unit_to_abbrev(new_reso)
- new_dtype = np.dtype(f"M8[{new_unit}]")
+ new_dtype = get_supported_dtype(new_dtype)
  data = astype_overflowsafe(data, dtype=new_dtype, copy=False)
- data_unit = get_unit_from_dtype(new_dtype)
  copy = False
 
  if data.dtype.byteorder == ">":
@@ -2346,6 +2341,7 @@ def _construct_from_dt64_naive(
  if data.ndim > 1:
  data = data.ravel()
 
+ data_unit = get_unit_from_dtype(new_dtype)
  data = tzconversion.tz_localize_to_utc(
  data.view("i8"), tz, ambiguous=ambiguous, creso=data_unit
  )
@@ -2552,7 +2548,7 @@ def _validate_dt64_dtype(dtype):
 
  if (
  isinstance(dtype, np.dtype)
- and (dtype.kind != "M" or not is_supported_unit(get_unit_from_dtype(dtype)))
+ and (dtype.kind != "M" or not is_supported_dtype(dtype))
  ) or not isinstance(dtype, (np.dtype, DatetimeTZDtype)):
  raise ValueError(
  f"Unexpected value for 'dtype': '{dtype}'. "

pandas/core/arrays/masked.py

@@ -15,10 +15,7 @@
  lib,
  missing as libmissing,
 )
-from pandas._libs.tslibs import (
- get_unit_from_dtype,
- is_supported_unit,
-)
+from pandas._libs.tslibs import is_supported_dtype
 from pandas._typing import (
  ArrayLike,
  AstypeArg,
@@ -876,9 +873,7 @@ def _maybe_mask_result(
 
  return BooleanArray(result, mask, copy=False)
 
- elif lib.is_np_dtype(result.dtype, "m") and is_supported_unit(
- get_unit_from_dtype(result.dtype)
- ):
+ elif lib.is_np_dtype(result.dtype, "m") and is_supported_dtype(result.dtype):
  # e.g. test_numeric_arr_mul_tdscalar_numexpr_path
  from pandas.core.arrays import TimedeltaArray
 

pandas/core/arrays/numpy_.py

@@ -8,10 +8,7 @@
 import numpy as np
 
 from pandas._libs import lib
-from pandas._libs.tslibs import (
- get_unit_from_dtype,
- is_supported_unit,
-)
+from pandas._libs.tslibs import is_supported_dtype
 from pandas.compat.numpy import function as nv
 
 from pandas.core.dtypes.astype import astype_array
@@ -553,9 +550,7 @@ def _cmp_method(self, other, op):
  def _wrap_ndarray_result(self, result: np.ndarray):
  # If we have timedelta64[ns] result, return a TimedeltaArray instead
  # of a NumpyExtensionArray
- if result.dtype.kind == "m" and is_supported_unit(
- get_unit_from_dtype(result.dtype)
- ):
+ if result.dtype.kind == "m" and is_supported_dtype(result.dtype):
  from pandas.core.arrays import TimedeltaArray
 
  return TimedeltaArray._simple_new(result, dtype=result.dtype)