Merge remote-tracking branch 'upstream/main' into docs_bisect

Author: seljaks

pandas/core/apply.py

@@ -506,7 +506,7 @@ def agg_dict_like(self) -> DataFrame | Series:
                 ktu._set_names(selected_obj.columns.names)
                 keys_to_use = ktu
 
-            axis = 0 if isinstance(obj, ABCSeries) else 1
+            axis: AxisInt = 0 if isinstance(obj, ABCSeries) else 1
             # error: Key expression in dictionary comprehension has incompatible type
             # "Hashable"; expected type "NDFrame"  [misc]
             result = concat(
@@ -932,7 +932,7 @@ def apply_str(self) -> DataFrame | Series:
 
 
 class FrameRowApply(FrameApply):
-    axis = 0
+    axis: AxisInt = 0
 
     def apply_broadcast(self, target: DataFrame) -> DataFrame:
         return super().apply_broadcast(target)
@@ -992,7 +992,7 @@ def wrap_results_for_axis(
 
 
 class FrameColumnApply(FrameApply):
-    axis = 1
+    axis: AxisInt = 1
 
     def apply_broadcast(self, target: DataFrame) -> DataFrame:
         result = super().apply_broadcast(target.T)
@@ -1069,7 +1069,7 @@ def infer_to_same_shape(self, results: ResType, res_index: Index) -> DataFrame:
 
 class SeriesApply(NDFrameApply):
     obj: Series
-    axis = 0
+    axis: AxisInt = 0
 
     def __init__(
         self,
@@ -1199,7 +1199,7 @@ def transform(self):
 
 
 class ResamplerWindowApply(Apply):
-    axis = 0
+    axis: AxisInt = 0
     obj: Resampler | BaseWindow
 
     def __init__(

pandas/core/arrays/datetimelike.py

@@ -36,6 +36,7 @@
     Period,
     Resolution,
     Tick,
+    Timedelta,
     Timestamp,
     astype_overflowsafe,
     delta_to_nanoseconds,
@@ -1122,7 +1123,7 @@ def _get_i8_values_and_mask(
         if isinstance(other, Period):
             i8values = other.ordinal
             mask = None
-        elif isinstance(other, Timestamp):
+        elif isinstance(other, (Timestamp, Timedelta)):
             i8values = other.value
             mask = None
         else:
@@ -1203,33 +1204,12 @@ def _sub_datetimelike_scalar(self, other: datetime | np.datetime64):
         self = cast("DatetimeArray", self)
         # subtract a datetime from myself, yielding a ndarray[timedelta64[ns]]
 
-        # error: Non-overlapping identity check (left operand type: "Union[datetime,
-        # datetime64]", right operand type: "NaTType")  [comparison-overlap]
-        assert other is not NaT  # type: ignore[comparison-overlap]
-        other = Timestamp(other)
-        # error: Non-overlapping identity check (left operand type: "Timestamp",
-        # right operand type: "NaTType")
-        if other is NaT:  # type: ignore[comparison-overlap]
+        if isna(other):
+            # i.e. np.datetime64("NaT")
             return self - NaT
 
-        try:
-            self._assert_tzawareness_compat(other)
-        except TypeError as err:
-            new_message = str(err).replace("compare", "subtract")
-            raise type(err)(new_message) from err
-
-        i8 = self.asi8
-        result = checked_add_with_arr(i8, -other.value, arr_mask=self._isnan)
-        res_m8 = result.view(f"timedelta64[{self._unit}]")
-
-        new_freq = None
-        if isinstance(self.freq, Tick):
-            # adding a scalar preserves freq
-            new_freq = self.freq
-
-        from pandas.core.arrays import TimedeltaArray
-
-        return TimedeltaArray._simple_new(res_m8, dtype=res_m8.dtype, freq=new_freq)
+        other = Timestamp(other)
+        return self._sub_datetimelike(other)
 
     @final
     def _sub_datetime_arraylike(self, other):
@@ -1241,19 +1221,28 @@ def _sub_datetime_arraylike(self, other):
 
         self = cast("DatetimeArray", self)
         other = ensure_wrapped_if_datetimelike(other)
+        return self._sub_datetimelike(other)
+
+    @final
+    def _sub_datetimelike(self, other: Timestamp | DatetimeArray) -> TimedeltaArray:
+        self = cast("DatetimeArray", self)
+
+        from pandas.core.arrays import TimedeltaArray
 
         try:
             self._assert_tzawareness_compat(other)
         except TypeError as err:
             new_message = str(err).replace("compare", "subtract")
             raise type(err)(new_message) from err
 
-        self_i8 = self.asi8
-        other_i8 = other.asi8
-        new_values = checked_add_with_arr(
-            self_i8, -other_i8, arr_mask=self._isnan, b_mask=other._isnan
+        other_i8, o_mask = self._get_i8_values_and_mask(other)
+        res_values = checked_add_with_arr(
+            self.asi8, -other_i8, arr_mask=self._isnan, b_mask=o_mask
         )
-        return new_values.view("timedelta64[ns]")
+        res_m8 = res_values.view(f"timedelta64[{self._unit}]")
+
+        new_freq = self._get_arithmetic_result_freq(other)
+        return TimedeltaArray._simple_new(res_m8, dtype=res_m8.dtype, freq=new_freq)
 
     @final
     def _sub_period(self, other: Period) -> npt.NDArray[np.object_]:
@@ -1289,24 +1278,15 @@ def _add_timedeltalike_scalar(self, other):
         Same type as self
         """
         if isna(other):
-            # i.e np.timedelta64("NaT"), not recognized by delta_to_nanoseconds
+            # i.e np.timedelta64("NaT")
             new_values = np.empty(self.shape, dtype="i8").view(self._ndarray.dtype)
             new_values.fill(iNaT)
             return type(self)._simple_new(new_values, dtype=self.dtype)
 
         # PeriodArray overrides, so we only get here with DTA/TDA
         self = cast("DatetimeArray | TimedeltaArray", self)
-        inc = delta_to_nanoseconds(other, reso=self._reso)
-
-        new_values = checked_add_with_arr(self.asi8, inc, arr_mask=self._isnan)
-        new_values = new_values.view(self._ndarray.dtype)
-
-        new_freq = None
-        if isinstance(self.freq, Tick) or is_period_dtype(self.dtype):
-            # adding a scalar preserves freq
-            new_freq = self.freq
-
-        return type(self)._simple_new(new_values, dtype=self.dtype, freq=new_freq)
+        other = Timedelta(other)._as_unit(self._unit)
+        return self._add_timedeltalike(other)
 
     def _add_timedelta_arraylike(
         self, other: TimedeltaArray | npt.NDArray[np.timedelta64]
@@ -1334,13 +1314,21 @@ def _add_timedelta_arraylike(
             else:
                 other = other._as_unit(self._unit)
 
-        self_i8 = self.asi8
-        other_i8 = other.asi8
+        return self._add_timedeltalike(other)
+
+    @final
+    def _add_timedeltalike(self, other: Timedelta | TimedeltaArray):
+        self = cast("DatetimeArray | TimedeltaArray", self)
+
+        other_i8, o_mask = self._get_i8_values_and_mask(other)
         new_values = checked_add_with_arr(
-            self_i8, other_i8, arr_mask=self._isnan, b_mask=other._isnan
+            self.asi8, other_i8, arr_mask=self._isnan, b_mask=o_mask
         )
         res_values = new_values.view(self._ndarray.dtype)
-        return type(self)._simple_new(res_values, dtype=self.dtype)
+
+        new_freq = self._get_arithmetic_result_freq(other)
+
+        return type(self)._simple_new(res_values, dtype=self.dtype, freq=new_freq)
 
     @final
     def _add_nat(self):

pandas/core/construction.py

@@ -49,6 +49,7 @@
 )
 from pandas.core.dtypes.common import (
     is_datetime64_ns_dtype,
+    is_dtype_equal,
     is_extension_array_dtype,
     is_float_dtype,
     is_integer_dtype,
@@ -328,6 +329,14 @@ def array(
     if isinstance(dtype, str):
         dtype = registry.find(dtype) or dtype
 
+    if isinstance(data, ExtensionArray) and (
+        dtype is None or is_dtype_equal(dtype, data.dtype)
+    ):
+        # e.g. TimedeltaArray[s], avoid casting to PandasArray
+        if copy:
+            return data.copy()
+        return data
+
     if is_extension_array_dtype(dtype):
         cls = cast(ExtensionDtype, dtype).construct_array_type()
         return cls._from_sequence(data, dtype=dtype, copy=copy)

pandas/core/frame.py

@@ -7586,7 +7586,7 @@ class      diet
     # Arithmetic Methods
 
     def _cmp_method(self, other, op):
-        axis = 1  # only relevant for Series other case
+        axis: Literal[1] = 1  # only relevant for Series other case
 
         self, other = ops.align_method_FRAME(self, other, axis, flex=False, level=None)
 
@@ -7598,7 +7598,7 @@ def _arith_method(self, other, op):
         if ops.should_reindex_frame_op(self, other, op, 1, 1, None, None):
             return ops.frame_arith_method_with_reindex(self, other, op)
 
-        axis = 1  # only relevant for Series other case
+        axis: Literal[1] = 1  # only relevant for Series other case
         other = ops.maybe_prepare_scalar_for_op(other, (self.shape[axis],))
 
         self, other = ops.align_method_FRAME(self, other, axis, flex=True, level=None)
@@ -10657,7 +10657,7 @@ def c(x):
         if not drop:
             # Find non-matching labels along the given axis
             # and append missing correlations (GH 22375)
-            raxis = 1 if axis == 0 else 0
+            raxis: AxisInt = 1 if axis == 0 else 0
             result_index = this._get_axis(raxis).union(other._get_axis(raxis))
             idx_diff = result_index.difference(correl.index)
 
@@ -11666,7 +11666,7 @@ def isin(self, values: Series | DataFrame | Sequence | Mapping) -> DataFrame:
         "columns": 1,
     }
     _AXIS_LEN = len(_AXIS_ORDERS)
-    _info_axis_number = 1
+    _info_axis_number: Literal[1] = 1
     _info_axis_name: Literal["columns"] = "columns"
 
     index = properties.AxisProperty(

pandas/core/generic.py

@@ -577,7 +577,7 @@ def _get_axis(self, axis: Axis) -> Index:
 
     @final
     @classmethod
-    def _get_block_manager_axis(cls, axis: Axis) -> int:
+    def _get_block_manager_axis(cls, axis: Axis) -> AxisInt:
         """Map the axis to the block_manager axis."""
         axis = cls._get_axis_number(axis)
         ndim = cls._AXIS_LEN
@@ -1705,7 +1705,7 @@ def _is_level_reference(self, key: Level, axis: Axis = 0) -> bool_t:
         )
 
     @final
-    def _is_label_reference(self, key: Level, axis: int = 0) -> bool_t:
+    def _is_label_reference(self, key: Level, axis: Axis = 0) -> bool_t:
         """
         Test whether a key is a label reference for a given axis.
 
@@ -1725,8 +1725,8 @@ def _is_label_reference(self, key: Level, axis: int = 0) -> bool_t:
         -------
         is_label: bool
         """
-        axis = self._get_axis_number(axis)
-        other_axes = (ax for ax in range(self._AXIS_LEN) if ax != axis)
+        axis_int = self._get_axis_number(axis)
+        other_axes = (ax for ax in range(self._AXIS_LEN) if ax != axis_int)
 
         return (
             key is not None

pandas/core/indexing.py

@@ -1111,7 +1111,7 @@ class _LocIndexer(_LocationIndexer):
     # Key Checks
 
     @doc(_LocationIndexer._validate_key)
-    def _validate_key(self, key, axis: AxisInt):
+    def _validate_key(self, key, axis: Axis):
         # valid for a collection of labels (we check their presence later)
         # slice of labels (where start-end in labels)
         # slice of integers (only if in the labels)

pandas/core/reshape/concat.py

@@ -21,6 +21,7 @@
 
 from pandas._typing import (
     Axis,
+    AxisInt,
     HashableT,
 )
 from pandas.util._decorators import (
@@ -636,7 +637,7 @@ def _get_new_axes(self) -> list[Index]:
             for i in range(ndim)
         ]
 
-    def _get_comb_axis(self, i: int) -> Index:
+    def _get_comb_axis(self, i: AxisInt) -> Index:
         data_axis = self.objs[0]._get_block_manager_axis(i)
         return get_objs_combined_axis(
             self.objs,

pandas/core/series.py

@@ -6228,7 +6228,7 @@ def mask(  # type: ignore[override]
     # Add index
     _AXIS_ORDERS: list[Literal["index", "columns"]] = ["index"]
     _AXIS_LEN = len(_AXIS_ORDERS)
-    _info_axis_number = 0
+    _info_axis_number: Literal[0] = 0
     _info_axis_name: Literal["index"] = "index"
 
     index = properties.AxisProperty(

pandas/io/excel/_base.py

@@ -1662,7 +1662,7 @@ def __init__(
                     f"only the xls format is supported. Install openpyxl instead."
                 )
             elif ext and ext != "xls":
-                stacklevel = find_stack_level()
+                stacklevel = find_stack_level(inspect.currentframe())
                 warnings.warn(
                     f"Your version of xlrd is {xlrd_version}. In xlrd >= 2.0, "
                     f"only the xls format is supported. Install "