Skip to content

ENH: timedelta_range unit keyword #49824

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Nov 22, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions doc/source/whatsnew/v2.0.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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`)
-
Expand Down
21 changes: 18 additions & 3 deletions pandas/core/arrays/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]):
Expand All @@ -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")

Expand All @@ -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)

# ----------------------------------------------------------------
Expand Down
17 changes: 16 additions & 1 deletion pandas/core/indexes/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
-------
Expand Down Expand Up @@ -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)
6 changes: 6 additions & 0 deletions pandas/tests/indexes/timedeltas/test_timedelta_range.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down