diff --git a/pandas/core/arrays/_ranges.py b/pandas/core/arrays/_ranges.py index f2582cd6f2e58..5f0409188e5e3 100644 --- a/pandas/core/arrays/_ranges.py +++ b/pandas/core/arrays/_ranges.py @@ -54,6 +54,8 @@ def generate_regular_range( iend = end._value if end is not None else None freq.nanos # raises if non-fixed frequency td = Timedelta(freq) + b: int | np.int64 | np.uint64 + e: int | np.int64 | np.uint64 try: td = td.as_unit( # pyright: ignore[reportGeneralTypeIssues] unit, round_ok=False @@ -96,7 +98,7 @@ def generate_regular_range( def _generate_range_overflow_safe( endpoint: int, periods: int, stride: int, side: str = "start" -) -> int: +) -> np.int64 | np.uint64: """ Calculate the second endpoint for passing to np.arange, checking to avoid an integer overflow. Catch OverflowError and re-raise @@ -115,7 +117,7 @@ def _generate_range_overflow_safe( Returns ------- - other_end : int + other_end : np.int64 | np.uint64 Raises ------ @@ -157,13 +159,13 @@ def _generate_range_overflow_safe( remaining = periods - mid_periods assert 0 < remaining < periods, (remaining, periods, endpoint, stride) - midpoint = _generate_range_overflow_safe(endpoint, mid_periods, stride, side) + midpoint = int(_generate_range_overflow_safe(endpoint, mid_periods, stride, side)) return _generate_range_overflow_safe(midpoint, remaining, stride, side) def _generate_range_overflow_safe_signed( endpoint: int, periods: int, stride: int, side: str -) -> int: +) -> np.int64 | np.uint64: """ A special case for _generate_range_overflow_safe where `periods * stride` can be calculated without overflowing int64 bounds. @@ -181,9 +183,7 @@ def _generate_range_overflow_safe_signed( # Putting this into a DatetimeArray/TimedeltaArray # would incorrectly be interpreted as NaT raise OverflowError - # error: Incompatible return value type (got "signedinteger[_64Bit]", - # expected "int") - return result # type: ignore[return-value] + return result except (FloatingPointError, OverflowError): # with endpoint negative and addend positive we risk # FloatingPointError; with reversed signed we risk OverflowError @@ -198,15 +198,11 @@ def _generate_range_overflow_safe_signed( # exceed implementation bounds, but when passing the result to # np.arange will get a result slightly within the bounds - # error: Incompatible types in assignment (expression has type - # "unsignedinteger[_64Bit]", variable has type "signedinteger[_64Bit]") - result = np.uint64(endpoint) + np.uint64(addend) # type: ignore[assignment] + uresult = np.uint64(endpoint) + np.uint64(addend) i64max = np.uint64(i8max) - assert result > i64max - if result <= i64max + np.uint64(stride): - # error: Incompatible return value type (got "unsignedinteger", expected - # "int") - return result # type: ignore[return-value] + assert uresult > i64max + if uresult <= i64max + np.uint64(stride): + return uresult raise OutOfBoundsDatetime( f"Cannot generate range with {side}={endpoint} and periods={periods}"