Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Helper functions to generate range-like data for DatetimeArray
3(and possibly TimedeltaArray/PeriodArray)
4"""
6from __future__ import annotations
8from typing import TYPE_CHECKING
10import numpy as np
12from pandas._libs.lib import i8max
13from pandas._libs.tslibs import (
14 BaseOffset,
15 Day,
16 OutOfBoundsDatetime,
17 Timedelta,
18 Timestamp,
19 iNaT,
20)
22from pandas.core.construction import range_to_ndarray
24if TYPE_CHECKING:
25 from pandas._typing import (
26 TimeUnit,
27 npt,
28 )
31def generate_regular_range(
32 start: Timestamp | Timedelta | None,
33 end: Timestamp | Timedelta | None,
34 periods: int | None,
35 freq: BaseOffset,
36 unit: TimeUnit = "ns",
37) -> npt.NDArray[np.intp]:
38 """
39 Generate a range of dates or timestamps with the spans between dates
40 described by the given `freq` DateOffset.
42 Parameters
43 ----------
44 start : Timedelta, Timestamp or None
45 First point of produced date range.
46 end : Timedelta, Timestamp or None
47 Last point of produced date range.
48 periods : int or None
49 Number of periods in produced date range.
50 freq : Tick
51 Describes space between dates in produced date range.
52 unit : {'s', 'ms', 'us', 'ns'}, default "ns"
53 The resolution the output is meant to represent.
55 Returns
56 -------
57 ndarray[np.int64]
58 Representing the given resolution.
59 """
60 istart = start._value if start is not None else None
61 iend = end._value if end is not None else None
62 if isinstance(freq, Day):
63 # In contexts without a timezone, a Day offset is unambiguously
64 # interpretable as Timedelta-like.
65 td = Timedelta(days=freq.n)
66 else:
67 freq.nanos # raises if non-fixed frequency
68 td = Timedelta(freq)
69 b: int
70 e: int
71 try:
72 td = td.as_unit(unit, round_ok=False)
73 except ValueError as err:
74 raise ValueError(
75 f"freq={freq} is incompatible with unit={unit}. "
76 "Use a lower freq or a higher unit instead."
77 ) from err
78 stride = int(td._value)
80 if periods is None and istart is not None and iend is not None:
81 b = istart
82 # cannot just use e = Timestamp(end) + 1 because arange breaks when
83 # stride is too large, see GH10887
84 e = b + (iend - b) // stride * stride + stride // 2 + 1
85 elif istart is not None and periods is not None:
86 b = istart
87 e = _generate_range_overflow_safe(b, periods, stride, side="start")
88 elif iend is not None and periods is not None:
89 e = iend + stride
90 b = _generate_range_overflow_safe(e, periods, stride, side="end")
91 else:
92 raise ValueError(
93 "at least 'start' or 'end' should be specified if a 'period' is given."
94 )
96 return range_to_ndarray(range(b, e, stride))
99def _generate_range_overflow_safe(
100 endpoint: int, periods: int, stride: int, side: str = "start"
101) -> int:
102 """
103 Calculate the second endpoint for passing to np.arange, checking
104 to avoid an integer overflow. Catch OverflowError and re-raise
105 as OutOfBoundsDatetime.
107 Parameters
108 ----------
109 endpoint : int
110 nanosecond timestamp of the known endpoint of the desired range
111 periods : int
112 number of periods in the desired range
113 stride : int
114 nanoseconds between periods in the desired range
115 side : {'start', 'end'}
116 which end of the range `endpoint` refers to
118 Returns
119 -------
120 other_end : int
122 Raises
123 ------
124 OutOfBoundsDatetime
125 """
126 # GH#14187 raise instead of incorrectly wrapping around
127 assert side in ["start", "end"]
129 i64max = np.uint64(i8max)
130 msg = f"Cannot generate range with {side}={endpoint} and periods={periods}"
132 with np.errstate(over="raise"):
133 # if periods * strides cannot be multiplied within the *uint64* bounds,
134 # we cannot salvage the operation by recursing, so raise
135 try:
136 addend = np.uint64(periods) * np.uint64(np.abs(stride))
137 except FloatingPointError as err:
138 raise OutOfBoundsDatetime(msg) from err
140 if np.abs(addend) <= i64max:
141 # relatively easy case without casting concerns
142 return _generate_range_overflow_safe_signed(endpoint, periods, stride, side)
144 elif (endpoint > 0 and side == "start" and stride > 0) or (
145 endpoint < 0 < stride and side == "end"
146 ):
147 # no chance of not-overflowing
148 raise OutOfBoundsDatetime(msg)
150 elif side == "end" and endpoint - stride <= i64max < endpoint:
151 # in _generate_regular_range we added `stride` thereby overflowing
152 # the bounds. Adjust to fix this.
153 return _generate_range_overflow_safe(
154 endpoint - stride, periods - 1, stride, side
155 )
157 # split into smaller pieces
158 mid_periods = periods // 2
159 remaining = periods - mid_periods
160 assert 0 < remaining < periods, (remaining, periods, endpoint, stride)
162 midpoint = int(_generate_range_overflow_safe(endpoint, mid_periods, stride, side))
163 return _generate_range_overflow_safe(midpoint, remaining, stride, side)
166def _generate_range_overflow_safe_signed(
167 endpoint: int, periods: int, stride: int, side: str
168) -> int:
169 """
170 A special case for _generate_range_overflow_safe where `periods * stride`
171 can be calculated without overflowing int64 bounds.
172 """
173 assert side in ["start", "end"]
174 if side == "end":
175 stride *= -1
177 with np.errstate(over="raise"):
178 addend = np.int64(periods) * np.int64(stride)
179 try:
180 # easy case with no overflows
181 result = np.int64(endpoint) + addend
182 if result == iNaT:
183 # Putting this into a DatetimeArray/TimedeltaArray
184 # would incorrectly be interpreted as NaT
185 raise OverflowError
186 return int(result)
187 except (FloatingPointError, OverflowError):
188 # with endpoint negative and addend positive we risk
189 # FloatingPointError; with reversed signed we risk OverflowError
190 pass
192 # if stride and endpoint had opposite signs, then endpoint + addend
193 # should never overflow. so they must have the same signs
194 assert (stride > 0 and endpoint >= 0) or (stride < 0 and endpoint <= 0)
196 if stride > 0:
197 # watch out for very special case in which we just slightly
198 # exceed implementation bounds, but when passing the result to
199 # np.arange will get a result slightly within the bounds
201 uresult = np.uint64(endpoint) + np.uint64(addend)
202 i64max = np.uint64(i8max)
203 assert uresult > i64max
204 if uresult <= i64max + np.uint64(stride):
205 return int(uresult)
207 raise OutOfBoundsDatetime(
208 f"Cannot generate range with {side}={endpoint} and periods={periods}"
209 )