Skip to content

Commit 6a5eb13

Browse files
authored
ENH: timedelta_range unit keyword (#49824)
* ENH: timedelta_range unit keyword * GH ref * mypy fixup
1 parent e1c38c9 commit 6a5eb13

File tree

4 files changed

+41
-4
lines changed

4 files changed

+41
-4
lines changed

doc/source/whatsnew/v2.0.0.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ Other enhancements
6262
- Fix ``test`` optional_extra by adding missing test package ``pytest-asyncio`` (:issue:`48361`)
6363
- :func:`DataFrame.astype` exception message thrown improved to include column name when type conversion is not possible. (:issue:`47571`)
6464
- :func:`date_range` now supports a ``unit`` keyword ("s", "ms", "us", or "ns") to specify the desired resolution of the output index (:issue:`49106`)
65+
- :func:`timedelta_range` now supports a ``unit`` keyword ("s", "ms", "us", or "ns") to specify the desired resolution of the output index (:issue:`49824`)
6566
- :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`)
6667
- Added ``name`` parameter to :meth:`IntervalIndex.from_breaks`, :meth:`IntervalIndex.from_arrays` and :meth:`IntervalIndex.from_tuples` (:issue:`48911`)
6768
-

pandas/core/arrays/timedeltas.py

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -254,8 +254,12 @@ def _from_sequence_not_strict(
254254

255255
return result
256256

257+
# Signature of "_generate_range" incompatible with supertype
258+
# "DatetimeLikeArrayMixin"
257259
@classmethod
258-
def _generate_range(cls, start, end, periods, freq, closed=None):
260+
def _generate_range( # type: ignore[override]
261+
cls, start, end, periods, freq, closed=None, *, unit: str | None = None
262+
):
259263

260264
periods = dtl.validate_periods(periods)
261265
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):
273277
if end is not None:
274278
end = Timedelta(end).as_unit("ns")
275279

280+
if unit is not None:
281+
if unit not in ["s", "ms", "us", "ns"]:
282+
raise ValueError("'unit' must be one of 's', 'ms', 'us', 'ns'")
283+
else:
284+
unit = "ns"
285+
286+
if start is not None and unit is not None:
287+
start = start.as_unit(unit, round_ok=False)
288+
if end is not None and unit is not None:
289+
end = end.as_unit(unit, round_ok=False)
290+
276291
left_closed, right_closed = validate_endpoints(closed)
277292

278293
if freq is not None:
279-
index = generate_regular_range(start, end, periods, freq)
294+
index = generate_regular_range(start, end, periods, freq, unit=unit)
280295
else:
281296
index = np.linspace(start.value, end.value, periods).astype("i8")
282297

@@ -285,7 +300,7 @@ def _generate_range(cls, start, end, periods, freq, closed=None):
285300
if not right_closed:
286301
index = index[:-1]
287302

288-
td64values = index.view("m8[ns]")
303+
td64values = index.view(f"m8[{unit}]")
289304
return cls._simple_new(td64values, dtype=td64values.dtype, freq=freq)
290305

291306
# ----------------------------------------------------------------

pandas/core/indexes/timedeltas.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,8 @@ def timedelta_range(
207207
freq=None,
208208
name=None,
209209
closed=None,
210+
*,
211+
unit: str | None = None,
210212
) -> TimedeltaIndex:
211213
"""
212214
Return a fixed frequency TimedeltaIndex with day as the default.
@@ -226,6 +228,10 @@ def timedelta_range(
226228
closed : str, default None
227229
Make the interval closed with respect to the given frequency to
228230
the 'left', 'right', or both sides (None).
231+
unit : str, default None
232+
Specify the desired resolution of the result.
233+
234+
.. versionadded:: 2.0.0
229235
230236
Returns
231237
-------
@@ -270,10 +276,19 @@ def timedelta_range(
270276
TimedeltaIndex(['1 days 00:00:00', '2 days 08:00:00', '3 days 16:00:00',
271277
'5 days 00:00:00'],
272278
dtype='timedelta64[ns]', freq=None)
279+
280+
**Specify a unit**
281+
282+
>>> pd.timedelta_range("1 Day", periods=3, freq="100000D", unit="s")
283+
TimedeltaIndex(['1 days 00:00:00', '100001 days 00:00:00',
284+
'200001 days 00:00:00'],
285+
dtype='timedelta64[s]', freq='100000D')
273286
"""
274287
if freq is None and com.any_none(periods, start, end):
275288
freq = "D"
276289

277290
freq, _ = dtl.maybe_infer_freq(freq)
278-
tdarr = TimedeltaArray._generate_range(start, end, periods, freq, closed=closed)
291+
tdarr = TimedeltaArray._generate_range(
292+
start, end, periods, freq, closed=closed, unit=unit
293+
)
279294
return TimedeltaIndex._simple_new(tdarr, name=name)

pandas/tests/indexes/timedeltas/test_timedelta_range.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,12 @@
1515

1616

1717
class TestTimedeltas:
18+
def test_timedelta_range_unit(self):
19+
# GH#49824
20+
tdi = timedelta_range("0 Days", periods=10, freq="100000D", unit="s")
21+
exp_arr = (np.arange(10, dtype="i8") * 100_000).view("m8[D]").astype("m8[s]")
22+
tm.assert_numpy_array_equal(tdi.to_numpy(), exp_arr)
23+
1824
def test_timedelta_range(self):
1925

2026
expected = to_timedelta(np.arange(5), unit="D")

0 commit comments

Comments
 (0)