TYP: timedeltas.pyi (pandas-dev#40766)

Author: jbrockmendel

pandas/core/dtypes/cast.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+from contextlib import suppress
 from datetime import (
     date,
     datetime,
@@ -28,6 +29,7 @@
     NaT,
     OutOfBoundsDatetime,
     OutOfBoundsTimedelta,
+    Period,
     Timedelta,
     Timestamp,
     conversion,
@@ -55,6 +57,7 @@
     ensure_str,
     is_bool,
     is_bool_dtype,
+    is_categorical_dtype,
     is_complex,
     is_complex_dtype,
     is_datetime64_dtype,
@@ -78,13 +81,13 @@
     pandas_dtype,
 )
 from pandas.core.dtypes.dtypes import (
-    CategoricalDtype,
     DatetimeTZDtype,
     ExtensionDtype,
     IntervalDtype,
     PeriodDtype,
 )
 from pandas.core.dtypes.generic import (
+    ABCDataFrame,
     ABCExtensionArray,
     ABCSeries,
 )
@@ -189,13 +192,13 @@ def maybe_box_native(value: Scalar) -> Scalar:
         value = maybe_box_datetimelike(value)
     elif is_float(value):
         # error: Argument 1 to "float" has incompatible type
-        # "Union[Union[str, int, float, bool], Union[Any, Timestamp, Timedelta, Any]]";
+        # "Union[Union[str, int, float, bool], Union[Any, Any, Timedelta, Any]]";
         # expected "Union[SupportsFloat, _SupportsIndex, str]"
         value = float(value)  # type: ignore[arg-type]
     elif is_integer(value):
         # error: Argument 1 to "int" has incompatible type
-        # "Union[Union[str, int, float, bool], Union[Any, Timestamp, Timedelta, Any]]";
-        # expected "Union[str, SupportsInt, _SupportsIndex, _SupportsTrunc]"
+        # "Union[Union[str, int, float, bool], Union[Any, Any, Timedelta, Any]]";
+        # pected "Union[str, SupportsInt, _SupportsIndex, _SupportsTrunc]"
         value = int(value)  # type: ignore[arg-type]
     elif is_bool(value):
         value = bool(value)
@@ -246,6 +249,9 @@ def maybe_downcast_to_dtype(result: ArrayLike, dtype: str | np.dtype) -> ArrayLi
     try to cast to the specified dtype (e.g. convert back to bool/int
     or could be an astype of float64->float32
     """
+    if isinstance(result, ABCDataFrame):
+        # see test_pivot_table_doctest_case
+        return result
     do_round = False
 
     if isinstance(dtype, str):
@@ -272,9 +278,15 @@ def maybe_downcast_to_dtype(result: ArrayLike, dtype: str | np.dtype) -> ArrayLi
 
         dtype = np.dtype(dtype)
 
-    if not isinstance(dtype, np.dtype):
-        # enforce our signature annotation
-        raise TypeError(dtype)  # pragma: no cover
+    elif dtype.type is Period:
+        from pandas.core.arrays import PeriodArray
+
+        with suppress(TypeError):
+            # e.g. TypeError: int() argument must be a string, a
+            #  bytes-like object or a number, not 'Period
+
+            # error: "dtype[Any]" has no attribute "freq"
+            return PeriodArray(result, freq=dtype.freq)  # type: ignore[attr-defined]
 
     converted = maybe_downcast_numeric(result, dtype, do_round)
     if converted is not result:
@@ -283,7 +295,15 @@ def maybe_downcast_to_dtype(result: ArrayLike, dtype: str | np.dtype) -> ArrayLi
     # a datetimelike
     # GH12821, iNaT is cast to float
     if dtype.kind in ["M", "m"] and result.dtype.kind in ["i", "f"]:
-        result = result.astype(dtype)
+        if isinstance(dtype, DatetimeTZDtype):
+            # convert to datetime and change timezone
+            i8values = result.astype("i8", copy=False)
+            cls = dtype.construct_array_type()
+            # equiv: DatetimeArray(i8values).tz_localize("UTC").tz_convert(dtype.tz)
+            dt64values = i8values.view("M8[ns]")
+            result = cls._simple_new(dt64values, dtype=dtype)
+        else:
+            result = result.astype(dtype)
 
     return result
 
@@ -359,15 +379,15 @@ def trans(x):
     return result
 
 
-def maybe_cast_pointwise_result(
+def maybe_cast_result(
     result: ArrayLike,
     dtype: DtypeObj,
     numeric_only: bool = False,
+    how: str = "",
     same_dtype: bool = True,
 ) -> ArrayLike:
     """
-    Try casting result of a pointwise operation back to the original dtype if
-    appropriate.
+    Try casting result to a different type if appropriate
 
     Parameters
     ----------
@@ -377,6 +397,8 @@ def maybe_cast_pointwise_result(
         Input Series from which result was calculated.
     numeric_only : bool, default False
         Whether to cast only numerics or datetimes as well.
+    how : str, default ""
+        How the result was computed.
     same_dtype : bool, default True
         Specify dtype when calling _from_sequence
 
@@ -385,12 +407,12 @@ def maybe_cast_pointwise_result(
     result : array-like
         result maybe casted to the dtype.
     """
+    dtype = maybe_cast_result_dtype(dtype, how)
 
     assert not is_scalar(result)
 
     if isinstance(dtype, ExtensionDtype):
-        if not isinstance(dtype, (CategoricalDtype, DatetimeTZDtype)):
-            # TODO: avoid this special-casing
+        if not is_categorical_dtype(dtype) and dtype.kind != "M":
             # We have to special case categorical so as not to upcast
             # things like counts back to categorical
 
@@ -406,6 +428,42 @@ def maybe_cast_pointwise_result(
     return result
 
 
+def maybe_cast_result_dtype(dtype: DtypeObj, how: str) -> DtypeObj:
+    """
+    Get the desired dtype of a result based on the
+    input dtype and how it was computed.
+
+    Parameters
+    ----------
+    dtype : DtypeObj
+        Input dtype.
+    how : str
+        How the result was computed.
+
+    Returns
+    -------
+    DtypeObj
+        The desired dtype of the result.
+    """
+    from pandas.core.arrays.boolean import BooleanDtype
+    from pandas.core.arrays.floating import Float64Dtype
+    from pandas.core.arrays.integer import (
+        Int64Dtype,
+        _IntegerDtype,
+    )
+
+    if how in ["add", "cumsum", "sum", "prod"]:
+        if dtype == np.dtype(bool):
+            return np.dtype(np.int64)
+        elif isinstance(dtype, (BooleanDtype, _IntegerDtype)):
+            return Int64Dtype()
+    elif how in ["mean", "median", "var"] and isinstance(
+        dtype, (BooleanDtype, _IntegerDtype)
+    ):
+        return Float64Dtype()
+    return dtype
+
+
 def maybe_cast_to_extension_array(
     cls: type[ExtensionArray], obj: ArrayLike, dtype: ExtensionDtype | None = None
 ) -> ArrayLike:
@@ -729,9 +787,7 @@ def infer_dtype_from_scalar(val, pandas_dtype: bool = False) -> tuple[DtypeObj,
         except OutOfBoundsDatetime:
             return np.dtype(object), val
 
-        # error: Non-overlapping identity check (left operand type: "Timestamp",
-        # right operand type: "NaTType")
-        if val is NaT or val.tz is None:  # type: ignore[comparison-overlap]
+        if val is NaT or val.tz is None:
             dtype = np.dtype("M8[ns]")
             val = val.to_datetime64()
         else:
@@ -2058,7 +2114,7 @@ def validate_numeric_casting(dtype: np.dtype, value: Scalar) -> None:
     ValueError
     """
     # error: Argument 1 to "__call__" of "ufunc" has incompatible type
-    # "Union[Union[str, int, float, bool], Union[Any, Timestamp, Timedelta, Any]]";
+    # "Union[Union[str, int, float, bool], Union[Any, Any, Timedelta, Any]]";
     # expected "Union[Union[int, float, complex, str, bytes, generic],
     # Sequence[Union[int, float, complex, str, bytes, generic]],
     # Sequence[Sequence[Any]], _SupportsArray]"