diff --git a/doc/source/whatsnew/v2.0.0.rst b/doc/source/whatsnew/v2.0.0.rst index 0e6d1029d352b..0f6029b00d9e2 100644 --- a/doc/source/whatsnew/v2.0.0.rst +++ b/doc/source/whatsnew/v2.0.0.rst @@ -62,6 +62,7 @@ Other enhancements - Fix ``test`` optional_extra by adding missing test package ``pytest-asyncio`` (:issue:`48361`) - :func:`DataFrame.astype` exception message thrown improved to include column name when type conversion is not possible. (:issue:`47571`) - :func:`date_range` now supports a ``unit`` keyword ("s", "ms", "us", or "ns") to specify the desired resolution of the output index (:issue:`49106`) +- :func:`timedelta_range` now supports a ``unit`` keyword ("s", "ms", "us", or "ns") to specify the desired resolution of the output index (:issue:`49824`) - :meth:`DataFrame.to_json` now supports a ``mode`` keyword with supported inputs 'w' and 'a'. Defaulting to 'w', 'a' can be used when lines=True and orient='records' to append record oriented json lines to an existing json file. (:issue:`35849`) - Added ``name`` parameter to :meth:`IntervalIndex.from_breaks`, :meth:`IntervalIndex.from_arrays` and :meth:`IntervalIndex.from_tuples` (:issue:`48911`) - diff --git a/pandas/core/arrays/timedeltas.py b/pandas/core/arrays/timedeltas.py index fe7cade1711d0..2517ecfc47877 100644 --- a/pandas/core/arrays/timedeltas.py +++ b/pandas/core/arrays/timedeltas.py @@ -254,8 +254,12 @@ def _from_sequence_not_strict( return result + # Signature of "_generate_range" incompatible with supertype + # "DatetimeLikeArrayMixin" @classmethod - def _generate_range(cls, start, end, periods, freq, closed=None): + def _generate_range( # type: ignore[override] + cls, start, end, periods, freq, closed=None, *, unit: str | None = None + ): periods = dtl.validate_periods(periods) if freq is None and any(x is None for x in [periods, start, end]): @@ -273,10 +277,21 @@ def _generate_range(cls, start, end, periods, freq, closed=None): if end is not None: end = Timedelta(end).as_unit("ns") + if unit is not None: + if unit not in ["s", "ms", "us", "ns"]: + raise ValueError("'unit' must be one of 's', 'ms', 'us', 'ns'") + else: + unit = "ns" + + if start is not None and unit is not None: + start = start.as_unit(unit, round_ok=False) + if end is not None and unit is not None: + end = end.as_unit(unit, round_ok=False) + left_closed, right_closed = validate_endpoints(closed) if freq is not None: - index = generate_regular_range(start, end, periods, freq) + index = generate_regular_range(start, end, periods, freq, unit=unit) else: index = np.linspace(start.value, end.value, periods).astype("i8") @@ -285,7 +300,7 @@ def _generate_range(cls, start, end, periods, freq, closed=None): if not right_closed: index = index[:-1] - td64values = index.view("m8[ns]") + td64values = index.view(f"m8[{unit}]") return cls._simple_new(td64values, dtype=td64values.dtype, freq=freq) # ---------------------------------------------------------------- diff --git a/pandas/core/indexes/timedeltas.py b/pandas/core/indexes/timedeltas.py index 113f76a35e13f..851ec5d54711d 100644 --- a/pandas/core/indexes/timedeltas.py +++ b/pandas/core/indexes/timedeltas.py @@ -207,6 +207,8 @@ def timedelta_range( freq=None, name=None, closed=None, + *, + unit: str | None = None, ) -> TimedeltaIndex: """ Return a fixed frequency TimedeltaIndex with day as the default. @@ -226,6 +228,10 @@ def timedelta_range( closed : str, default None Make the interval closed with respect to the given frequency to the 'left', 'right', or both sides (None). + unit : str, default None + Specify the desired resolution of the result. + + .. versionadded:: 2.0.0 Returns ------- @@ -270,10 +276,19 @@ def timedelta_range( TimedeltaIndex(['1 days 00:00:00', '2 days 08:00:00', '3 days 16:00:00', '5 days 00:00:00'], dtype='timedelta64[ns]', freq=None) + + **Specify a unit** + + >>> pd.timedelta_range("1 Day", periods=3, freq="100000D", unit="s") + TimedeltaIndex(['1 days 00:00:00', '100001 days 00:00:00', + '200001 days 00:00:00'], + dtype='timedelta64[s]', freq='100000D') """ if freq is None and com.any_none(periods, start, end): freq = "D" freq, _ = dtl.maybe_infer_freq(freq) - tdarr = TimedeltaArray._generate_range(start, end, periods, freq, closed=closed) + tdarr = TimedeltaArray._generate_range( + start, end, periods, freq, closed=closed, unit=unit + ) return TimedeltaIndex._simple_new(tdarr, name=name) diff --git a/pandas/tests/indexes/timedeltas/test_timedelta_range.py b/pandas/tests/indexes/timedeltas/test_timedelta_range.py index 7277595f1d631..4f0cab2a433f7 100644 --- a/pandas/tests/indexes/timedeltas/test_timedelta_range.py +++ b/pandas/tests/indexes/timedeltas/test_timedelta_range.py @@ -15,6 +15,12 @@ class TestTimedeltas: + def test_timedelta_range_unit(self): + # GH#49824 + tdi = timedelta_range("0 Days", periods=10, freq="100000D", unit="s") + exp_arr = (np.arange(10, dtype="i8") * 100_000).view("m8[D]").astype("m8[s]") + tm.assert_numpy_array_equal(tdi.to_numpy(), exp_arr) + def test_timedelta_range(self): expected = to_timedelta(np.arange(5), unit="D")