CLN: refactor core/arrays/_range to support timedeltas

Author: hasB4K

pandas/core/arrays/_ranges.py

@@ -7,12 +7,12 @@
 
 import numpy as np
 
-from pandas._libs.tslibs import OutOfBoundsDatetime, Timestamp
+from pandas._libs.tslibs import OutOfBoundsDatetime, Timedelta, Timestamp
 
 from pandas.tseries.offsets import DateOffset, Tick, generate_range
 
 
-def generate_regular_range(
+def generate_timestamps_range(
     start: Timestamp, end: Timestamp, periods: int, freq: DateOffset
 ) -> Tuple[np.ndarray, str]:
     """
@@ -32,57 +32,79 @@ def generate_regular_range(
 
     Returns
     -------
-    ndarray[np.int64] representing nanosecond unix timestamps
+    (tuple): containing:
+
+        values : ndarray[np.int64] representing nanosecond unix timestamps
+        tz : the timezone of the range
     """
     if isinstance(freq, Tick):
-        stride = freq.nanos
-        if periods is None:
-            b = Timestamp(start).value
-            # cannot just use e = Timestamp(end) + 1 because arange breaks when
-            # stride is too large, see GH10887
-            e = b + (Timestamp(end).value - b) // stride * stride + stride // 2 + 1
-            # end.tz == start.tz by this point due to _generate implementation
-            tz = start.tz
-        elif start is not None:
-            b = Timestamp(start).value
-            e = _generate_range_overflow_safe(b, periods, stride, side="start")
-            tz = start.tz
-        elif end is not None:
-            e = Timestamp(end).value + stride
-            b = _generate_range_overflow_safe(e, periods, stride, side="end")
-            tz = end.tz
-        else:
-            raise ValueError(
-                "at least 'start' or 'end' should be specified "
-                "if a 'period' is given."
-            )
-
-        with np.errstate(over="raise"):
-            # If the range is sufficiently large, np.arange may overflow
-            #  and incorrectly return an empty array if not caught.
-            try:
-                values = np.arange(b, e, stride, dtype=np.int64)
-            except FloatingPointError:
-                xdr = [b]
-                while xdr[-1] != e:
-                    xdr.append(xdr[-1] + stride)
-                values = np.array(xdr[:-1], dtype=np.int64)
-
+        start_value = Timestamp(start).value if start is not None else None
+        end_value = Timestamp(end).value if end is not None else None
+        values = _generate_regular_range(start_value, end_value, periods, freq.nanos)
     else:
-        tz = None
-        # start and end should have the same timezone by this point
-        if start is not None:
-            tz = start.tz
-        elif end is not None:
-            tz = end.tz
-
         xdr = generate_range(start=start, end=end, periods=periods, offset=freq)
-
         values = np.array([x.value for x in xdr], dtype=np.int64)
 
+    tz = start.tz if start is not None else end.tz
     return values, tz
 
 
+def generate_timedeltas_range(
+    start: Timedelta, end: Timedelta, periods: int, freq: DateOffset
+):
+    """
+    Generate a range of dates with the spans between dates described by
+    the given `freq` DateOffset.
+
+    Parameters
+    ----------
+    start : Timedelta or None
+        first point of produced date range
+    end : Timedelta or None
+        last point of produced date range
+    periods : int
+        number of periods in produced date range
+    freq : DateOffset
+        describes space between dates in produced date range
+
+    Returns
+    -------
+    ndarray[np.int64] representing nanosecond timedeltas
+    """
+    start_value = Timedelta(start).value if start is not None else None
+    end_value = Timedelta(end).value if end is not None else None
+    return _generate_regular_range(start_value, end_value, periods, freq.nanos)
+
+
+def _generate_regular_range(start: int, end: int, periods: int, stride: int):
+    b = start
+    if periods is None:
+        # cannot just use e = Timestamp(end) + 1 because arange breaks when
+        # stride is too large, see GH10887
+        e = b + (end - b) // stride * stride + stride // 2 + 1
+    elif start is not None:
+        e = _generate_range_overflow_safe(b, periods, stride, side="start")
+    elif end is not None:
+        e = end + stride
+        b = _generate_range_overflow_safe(e, periods, stride, side="end")
+    else:
+        raise ValueError(
+            "at least 'start' or 'end' should be specified if a 'period' is given."
+        )
+
+    with np.errstate(over="raise"):
+        # If the range is sufficiently large, np.arange may overflow
+        #  and incorrectly return an empty array if not caught.
+        try:
+            values = np.arange(b, e, stride, dtype=np.int64)
+        except FloatingPointError:
+            xdr = [b]
+            while xdr[-1] != e:
+                xdr.append(xdr[-1] + stride)
+            values = np.array(xdr[:-1], dtype=np.int64)
+    return values
+
+
 def _generate_range_overflow_safe(
     endpoint: int, periods: int, stride: int, side: str = "start"
 ) -> int:

pandas/core/arrays/datetimes.py

@@ -44,7 +44,7 @@
 
 from pandas.core.algorithms import checked_add_with_arr
 from pandas.core.arrays import datetimelike as dtl
-from pandas.core.arrays._ranges import generate_regular_range
+from pandas.core.arrays._ranges import generate_timestamps_range
 import pandas.core.common as com
 
 from pandas.tseries.frequencies import get_period_alias, to_offset
@@ -404,7 +404,7 @@ def _generate_range(
                 if end is not None:
                     end = end.tz_localize(None)
             # TODO: consider re-implementing _cached_range; GH#17914
-            values, _tz = generate_regular_range(start, end, periods, freq)
+            values, _tz = generate_timestamps_range(start, end, periods, freq)
             index = cls._simple_new(values, freq=freq, dtype=tz_to_dtype(_tz))
 
             if tz is not None and index.tz is None:

pandas/core/arrays/timedeltas.py

@@ -38,7 +38,7 @@
 from pandas.core import nanops
 from pandas.core.algorithms import checked_add_with_arr
 from pandas.core.arrays import datetimelike as dtl
-from pandas.core.arrays._ranges import _generate_range_overflow_safe
+from pandas.core.arrays._ranges import generate_timedeltas_range
 import pandas.core.common as com
 from pandas.core.construction import extract_array
 
@@ -256,7 +256,7 @@ def _generate_range(cls, start, end, periods, freq, closed=None):
         left_closed, right_closed = dtl.validate_endpoints(closed)
 
         if freq is not None:
-            index = _generate_regular_range(start, end, periods, freq)
+            index = generate_timedeltas_range(start, end, periods, freq)
         else:
             index = np.linspace(start.value, end.value, periods).astype("i8")
             if len(index) >= 2:
@@ -1056,25 +1056,3 @@ def _validate_td64_dtype(dtype):
         raise ValueError(f"dtype {dtype} cannot be converted to timedelta64[ns]")
 
     return dtype
-
-
-def _generate_regular_range(start, end, periods, offset):
-    stride = offset.nanos
-    if periods is None:
-        b = Timedelta(start).value
-        # cannot just use e = Timestamp(end) + 1 because arange breaks when
-        # stride is too large, see GH 10887 & GH 30353
-        e = b + (Timedelta(end).value - b) // stride * stride + stride // 2 + 1
-    elif start is not None:
-        b = Timedelta(start).value
-        e = _generate_range_overflow_safe(b, periods, stride, side="start")
-    elif end is not None:
-        e = Timedelta(end).value + stride
-        b = _generate_range_overflow_safe(e, periods, stride, side="end")
-    else:
-        raise ValueError(
-            "at least 'start' or 'end' should be specified if a 'period' is given."
-        )
-
-    data = np.arange(b, e, stride, dtype=np.int64)
-    return data