ashlr-plugin 1.36.4

package/bench/refs/python-lib/pandas/core/arrays/_ranges.py

@@ -0,0 +1,276 @@
+"""
+Helper functions to generate range-like data for DatetimeArray
+(and possibly TimedeltaArray/PeriodArray)
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from pandas._libs.lib import i8max
+from pandas._libs.tslibs import (
+    BaseOffset,
+    Day,
+    OutOfBoundsDatetime,
+    Timedelta,
+    Timestamp,
+    iNaT,
+)
+
+from pandas.core.construction import range_to_ndarray
+
+if TYPE_CHECKING:
+    from pandas._typing import (
+        TimeUnit,
+        npt,
+    )
+
+
+def generate_regular_range(
+    start: Timestamp | Timedelta | None,
+    end: Timestamp | Timedelta | None,
+    periods: int | None,
+    freq: BaseOffset,
+    unit: TimeUnit = "ns",
+) -> npt.NDArray[np.intp]:
+    """
+    Generate a range of dates or timestamps with the spans between dates
+    described by the given `freq` DateOffset.
+
+    Parameters
+    ----------
+    start : Timedelta, Timestamp or None
+        First point of produced date range.
+    end : Timedelta, Timestamp or None
+        Last point of produced date range.
+    periods : int or None
+        Number of periods in produced date range.
+    freq : Tick
+        Describes space between dates in produced date range.
+    unit : {'s', 'ms', 'us', 'ns'}, default "ns"
+        The resolution the output is meant to represent.
+
+    Returns
+    -------
+    ndarray[np.int64]
+        Representing the given resolution.
+    """
+    istart = start._value if start is not None else None
+    iend = end._value if end is not None else None
+    if isinstance(freq, Day):
+        # In contexts without a timezone, a Day offset is unambiguously
+        #  interpretable as Timedelta-like.
+        td = Timedelta(days=freq.n)
+    else:
+        freq.nanos  # raises if non-fixed frequency
+        td = Timedelta(freq)
+    b: int
+    e: int
+    try:
+        td = td.as_unit(unit, round_ok=False)
+    except ValueError as err:
+        raise ValueError(
+            f"freq={freq} is incompatible with unit={unit}. "
+            "Use a lower freq or a higher unit instead."
+        ) from err
+    stride = int(td._value)
+
+    if periods is None and istart is not None and iend is not None:
+        b = istart
+        # cannot just use e = Timestamp(end) + 1 because arange breaks when
+        # stride is too large, see GH10887
+        e = b + (iend - b) // stride * stride + stride // 2 + 1
+    elif istart is not None and periods is not None:
+        b = istart
+        e = _generate_range_overflow_safe(b, periods, stride, side="start")
+    elif iend is not None and periods is not None:
+        e = iend + 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."
+        )
+
+    return range_to_ndarray(range(b, e, stride))
+
+
+def generate_daily_offset_range(
+    start: Timestamp | None,
+    end: Timestamp | None,
+    periods: int | None,
+    freq: BaseOffset,
+    unit: TimeUnit = "ns",
+) -> npt.NDArray[np.int64]:
+    """
+    Generate a range for offsets whose on-offset dates are a subset of a
+    daily grid, by generating a daily range and filtering.
+
+    This is a performance optimization (GH#16463) for offsets like BusinessDay
+    that implement ``_get_daily_offset_mask``. Instead of iteratively applying
+    the offset, we generate a daily range and filter it vectorized.
+
+    Parameters
+    ----------
+    start : Timestamp or None
+    end : Timestamp or None
+    periods : int or None
+    freq : BaseOffset
+        Must have ``_supports_daily_offset_mask`` and ``n >= 1``
+        (caller ensures both).
+    unit : str, default "ns"
+
+    Returns
+    -------
+    ndarray[int64]
+        The filtered i8 values.
+    """
+    abs_n = freq.n  # caller ensures n >= 1
+
+    # For the periods case, convert to start+end by extending the known
+    # endpoint with a generous Timedelta. The worst-case gap for any
+    # daily-grid offset is 7 calendar days per on-offset day (e.g. Week),
+    # so needed * 7 + 6 is always sufficient.
+    trim_to = None
+    trim_from_end = False
+    if periods is not None:
+        needed_on_offset = (periods - 1) * abs_n + 1
+        buffer_days = needed_on_offset * 7 + 6
+        if end is None:
+            # start is guaranteed non-None by the caller
+            end = (  # type: ignore[assignment]
+                start + Timedelta(days=buffer_days)  # type: ignore[operator]
+            ).as_unit(unit)
+        else:
+            trim_from_end = True
+            start = (end - Timedelta(days=buffer_days)).as_unit(unit)  # pyright: ignore[reportAssignmentType]
+        trim_to = periods
+
+    i8values = generate_regular_range(start, end, None, Day(), unit=unit)
+    dt64 = i8values.view(f"datetime64[{unit}]")
+    i8values = i8values[freq._get_daily_offset_mask(dt64)]  # type: ignore[attr-defined]
+
+    if abs_n > 1:
+        i8values = i8values[::abs_n]
+
+    if trim_to is not None:
+        if trim_from_end:
+            i8values = i8values[-trim_to:]
+        else:
+            i8values = i8values[:trim_to]
+
+    return i8values
+
+
+def _generate_range_overflow_safe(
+    endpoint: int, periods: int, stride: int, side: str = "start"
+) -> int:
+    """
+    Calculate the second endpoint for passing to np.arange, checking
+    to avoid an integer overflow.  Catch OverflowError and re-raise
+    as OutOfBoundsDatetime.
+
+    Parameters
+    ----------
+    endpoint : int
+        nanosecond timestamp of the known endpoint of the desired range
+    periods : int
+        number of periods in the desired range
+    stride : int
+        nanoseconds between periods in the desired range
+    side : {'start', 'end'}
+        which end of the range `endpoint` refers to
+
+    Returns
+    -------
+    other_end : int
+
+    Raises
+    ------
+    OutOfBoundsDatetime
+    """
+    # GH#14187 raise instead of incorrectly wrapping around
+    assert side in ["start", "end"]
+
+    i64max = np.uint64(i8max)
+    msg = f"Cannot generate range with {side}={endpoint} and periods={periods}"
+
+    with np.errstate(over="raise"):
+        # if periods * strides cannot be multiplied within the *uint64* bounds,
+        #  we cannot salvage the operation by recursing, so raise
+        try:
+            addend = np.uint64(periods) * np.uint64(np.abs(stride))
+        except FloatingPointError as err:
+            raise OutOfBoundsDatetime(msg) from err
+
+    if np.abs(addend) <= i64max:
+        # relatively easy case without casting concerns
+        return _generate_range_overflow_safe_signed(endpoint, periods, stride, side)
+
+    elif (endpoint > 0 and side == "start" and stride > 0) or (
+        endpoint < 0 < stride and side == "end"
+    ):
+        # no chance of not-overflowing
+        raise OutOfBoundsDatetime(msg)
+
+    elif side == "end" and endpoint - stride <= i64max < endpoint:
+        # in _generate_regular_range we added `stride` thereby overflowing
+        #  the bounds.  Adjust to fix this.
+        return _generate_range_overflow_safe(
+            endpoint - stride, periods - 1, stride, side
+        )
+
+    # split into smaller pieces
+    mid_periods = periods // 2
+    remaining = periods - mid_periods
+    assert 0 < remaining < periods, (remaining, periods, endpoint, stride)
+
+    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:
+    """
+    A special case for _generate_range_overflow_safe where `periods * stride`
+    can be calculated without overflowing int64 bounds.
+    """
+    assert side in ["start", "end"]
+    if side == "end":
+        stride *= -1
+
+    with np.errstate(over="raise"):
+        addend = np.int64(periods) * np.int64(stride)
+        try:
+            # easy case with no overflows
+            result = np.int64(endpoint) + addend
+            if result == iNaT:
+                # Putting this into a DatetimeArray/TimedeltaArray
+                #  would incorrectly be interpreted as NaT
+                raise OverflowError
+            return int(result)
+        except (FloatingPointError, OverflowError):
+            # with endpoint negative and addend positive we risk
+            #  FloatingPointError; with reversed signed we risk OverflowError
+            pass
+
+        # if stride and endpoint had opposite signs, then endpoint + addend
+        #  should never overflow.  so they must have the same signs
+        assert (stride > 0 and endpoint >= 0) or (stride < 0 and endpoint <= 0)
+
+        if stride > 0:
+            # watch out for very special case in which we just slightly
+            #  exceed implementation bounds, but when passing the result to
+            #  np.arange will get a result slightly within the bounds
+
+            uresult = np.uint64(endpoint) + np.uint64(addend)
+            i64max = np.uint64(i8max)
+            assert uresult > i64max
+            if uresult <= i64max + np.uint64(stride):
+                return int(uresult)
+
+    raise OutOfBoundsDatetime(
+        f"Cannot generate range with {side}={endpoint} and periods={periods}"
+    )

package/bench/refs/python-lib/pandas/core/arrays/_utils.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+import numpy as np
+
+from pandas._config import is_nan_na
+
+from pandas._libs import lib
+from pandas._libs.missing import NA
+from pandas.errors import LossySetitemError
+
+from pandas.core.dtypes.cast import np_can_hold_element
+from pandas.core.dtypes.common import is_numeric_dtype
+
+if TYPE_CHECKING:
+    from pandas._typing import (
+        npt,
+    )
+
+    from pandas.core.arrays.base import ExtensionArray
+
+
+def to_numpy_dtype_inference(
+    arr: ExtensionArray,
+    dtype: npt.DTypeLike | None,
+    na_value,
+    hasna: bool,
+) -> tuple[np.dtype | None, Any]:
+    result_dtype: np.dtype | None
+    inferred_numeric_dtype = False
+    if dtype is None and is_numeric_dtype(arr.dtype):
+        inferred_numeric_dtype = True
+        if hasna:
+            if arr.dtype.kind == "b":
+                result_dtype = np.dtype(np.object_)
+            else:
+                if arr.dtype.kind in "iu":
+                    result_dtype = np.dtype(np.float64)
+                else:
+                    result_dtype = arr.dtype.numpy_dtype  # type: ignore[attr-defined]
+                if na_value is lib.no_default:
+                    if not is_nan_na():
+                        na_value = NA
+                        dtype = np.dtype(object)
+                    else:
+                        na_value = np.nan
+        else:
+            result_dtype = arr.dtype.numpy_dtype  # type: ignore[attr-defined]
+    elif dtype is not None:
+        result_dtype = np.dtype(dtype)
+    else:
+        result_dtype = None
+
+    if na_value is lib.no_default:
+        if result_dtype is None or not hasna:
+            na_value = arr.dtype.na_value
+        elif result_dtype.kind == "f":
+            na_value = np.nan
+        elif result_dtype.kind == "M":
+            unit = np.datetime_data(result_dtype)[0]  # type: ignore[arg-type]
+            na_value = np.datetime64("NaT", unit)  # type: ignore[call-overload]
+        elif result_dtype.kind == "m":
+            unit = np.datetime_data(result_dtype)[0]  # type: ignore[arg-type]
+            na_value = np.timedelta64("NaT", unit)  # type: ignore[call-overload]
+        else:
+            na_value = arr.dtype.na_value
+
+    if inferred_numeric_dtype and hasna:
+        try:
+            np_can_hold_element(result_dtype, na_value)  # type: ignore[arg-type]
+        except LossySetitemError:
+            result_dtype = np.dtype(np.object_)
+    return result_dtype, na_value