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_daily_offset_range(
100 start: Timestamp | None,
101 end: Timestamp | None,
102 periods: int | None,
103 freq: BaseOffset,
104 unit: TimeUnit = "ns",
105) -> npt.NDArray[np.int64]:
106 """
107 Generate a range for offsets whose on-offset dates are a subset of a
108 daily grid, by generating a daily range and filtering.
110 This is a performance optimization (GH#16463) for offsets like BusinessDay
111 that implement ``_get_daily_offset_mask``. Instead of iteratively applying
112 the offset, we generate a daily range and filter it vectorized.
114 Parameters
115 ----------
116 start : Timestamp or None
117 end : Timestamp or None
118 periods : int or None
119 freq : BaseOffset
120 Must have ``_supports_daily_offset_mask`` and ``n >= 1``
121 (caller ensures both).
122 unit : str, default "ns"
124 Returns
125 -------
126 ndarray[int64]
127 The filtered i8 values.
128 """
129 abs_n = freq.n # caller ensures n >= 1
131 # Resolve the endpoints with offset arithmetic so the on-offset count is
132 # exact regardless of holidays/weekmasks. The old ``needed * 7 + 6``-day
133 # buffer assumed >=1 on-offset day per 7 calendar days, which fails for e.g.
134 # a CustomBusinessDay whose holidays blank out whole weeks -- the result was
135 # then silently short (GH#64648 post-merge). Arithmetic runs on normalized
136 # dates with the time-of-day re-attached, so ``normalize=True`` offsets
137 # don't strip it from the output (GH#44025).
138 if periods is not None and end is not None:
139 # end + periods: anchor at the last on-offset date <= end so that
140 # ``periods`` on-offset dates are returned (GH#64834).
141 tod: Timedelta = end - end.normalize()
142 # pyright can't see through the NaT-returning Timestamp constructor.
143 anchor: Timestamp = Timestamp( # pyright: ignore[reportAssignmentType]
144 freq.rollback(end.normalize())
145 )
146 start = (anchor - (periods - 1) * freq + tod).as_unit(unit)
147 end = (anchor + tod).as_unit(unit)
148 else:
149 # start is guaranteed non-None here by the caller (exactly two of
150 # start/end/periods are given, and this branch is not end+periods).
151 assert start is not None
152 # Roll an off-offset start forward to the first on-offset date (n >= 1).
153 tod = start - start.normalize()
154 anchor = Timestamp( # pyright: ignore[reportAssignmentType]
155 freq.rollforward(start.normalize())
156 )
157 start = (anchor + tod).as_unit(unit)
158 if periods is not None:
159 # start + periods.
160 end = (anchor + (periods - 1) * freq + tod).as_unit(unit)
162 i8values = generate_regular_range(start, end, None, Day(), unit=unit)
163 dt64 = i8values.view(f"datetime64[{unit}]")
164 i8values = i8values[freq._get_daily_offset_mask(dt64)] # type: ignore[attr-defined]
166 if abs_n > 1:
167 # start is the first on-offset date and end the last, so a forward
168 # stride anchored at the front lands on end (GH#64648, GH#65604).
169 i8values = i8values[::abs_n]
171 return i8values
174def _generate_range_overflow_safe(
175 endpoint: int, periods: int, stride: int, side: str = "start"
176) -> int:
177 """
178 Calculate the second endpoint for passing to np.arange, checking
179 to avoid an integer overflow. Catch OverflowError and re-raise
180 as OutOfBoundsDatetime.
182 Parameters
183 ----------
184 endpoint : int
185 nanosecond timestamp of the known endpoint of the desired range
186 periods : int
187 number of periods in the desired range
188 stride : int
189 nanoseconds between periods in the desired range
190 side : {'start', 'end'}
191 which end of the range `endpoint` refers to
193 Returns
194 -------
195 other_end : int
197 Raises
198 ------
199 OutOfBoundsDatetime
200 """
201 # GH#14187 raise instead of incorrectly wrapping around
202 assert side in ["start", "end"]
204 i64max = np.uint64(i8max)
205 msg = f"Cannot generate range with {side}={endpoint} and periods={periods}"
207 with np.errstate(over="raise"):
208 # if periods * strides cannot be multiplied within the *uint64* bounds,
209 # we cannot salvage the operation by recursing, so raise
210 try:
211 addend = np.uint64(periods) * np.uint64(np.abs(stride))
212 except FloatingPointError as err:
213 raise OutOfBoundsDatetime(msg) from err
215 if np.abs(addend) <= i64max:
216 # relatively easy case without casting concerns
217 return _generate_range_overflow_safe_signed(endpoint, periods, stride, side)
219 elif (endpoint > 0 and side == "start" and stride > 0) or (
220 endpoint < 0 < stride and side == "end"
221 ):
222 # no chance of not-overflowing
223 raise OutOfBoundsDatetime(msg)
225 elif side == "end" and endpoint - stride <= i64max < endpoint:
226 # in _generate_regular_range we added `stride` thereby overflowing
227 # the bounds. Adjust to fix this.
228 return _generate_range_overflow_safe(
229 endpoint - stride, periods - 1, stride, side
230 )
232 # split into smaller pieces
233 mid_periods = periods // 2
234 remaining = periods - mid_periods
235 assert 0 < remaining < periods, (remaining, periods, endpoint, stride)
237 midpoint = int(_generate_range_overflow_safe(endpoint, mid_periods, stride, side))
238 return _generate_range_overflow_safe(midpoint, remaining, stride, side)
241def _generate_range_overflow_safe_signed(
242 endpoint: int, periods: int, stride: int, side: str
243) -> int:
244 """
245 A special case for _generate_range_overflow_safe where `periods * stride`
246 can be calculated without overflowing int64 bounds.
247 """
248 assert side in ["start", "end"]
249 if side == "end":
250 stride *= -1
252 with np.errstate(over="raise"):
253 addend = np.int64(periods) * np.int64(stride)
254 try:
255 # easy case with no overflows
256 result = np.int64(endpoint) + addend
257 if result == iNaT:
258 # Putting this into a DatetimeArray/TimedeltaArray
259 # would incorrectly be interpreted as NaT
260 raise OverflowError
261 return int(result)
262 except (FloatingPointError, OverflowError):
263 # with endpoint negative and addend positive we risk
264 # FloatingPointError; with reversed signed we risk OverflowError
265 pass
267 # if stride and endpoint had opposite signs, then endpoint + addend
268 # should never overflow. so they must have the same signs
269 assert (stride > 0 and endpoint >= 0) or (stride < 0 and endpoint <= 0)
271 if stride > 0:
272 # watch out for very special case in which we just slightly
273 # exceed implementation bounds, but when passing the result to
274 # np.arange will get a result slightly within the bounds
276 uresult = np.uint64(endpoint) + np.uint64(addend)
277 i64max = np.uint64(i8max)
278 assert uresult > i64max
279 if uresult <= i64max + np.uint64(stride):
280 return int(uresult)
282 raise OutOfBoundsDatetime(
283 f"Cannot generate range with {side}={endpoint} and periods={periods}"
284 )