Skip to content

REF: tighten types in maybe_downcast_numeric, maybe_downcast_to_dtype #40087

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Feb 27, 2021
16 changes: 7 additions & 9 deletions pandas/core/dtypes/cast.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,6 @@
PeriodDtype,
)
from pandas.core.dtypes.generic import (
ABCDataFrame,
ABCExtensionArray,
ABCIndex,
ABCSeries,
Expand Down Expand Up @@ -233,19 +232,15 @@ def _disallow_mismatched_datetimelike(value, dtype: DtypeObj):
raise TypeError(f"Cannot cast {repr(value)} to {dtype}")


def maybe_downcast_to_dtype(result, dtype: Union[str, np.dtype]):
def maybe_downcast_to_dtype(
result: ArrayLike, dtype: Union[str, np.dtype]
) -> ArrayLike:
"""
try to cast to the specified dtype (e.g. convert back to bool/int
or could be an astype of float64->float32
"""
do_round = False

if is_scalar(result):
return result
elif isinstance(result, ABCDataFrame):
# occurs in pivot_table doctest
return result

if isinstance(dtype, str):
if dtype == "infer":
inferred_type = lib.infer_dtype(ensure_object(result), skipna=False)
Expand All @@ -265,6 +260,7 @@ def maybe_downcast_to_dtype(result, dtype: Union[str, np.dtype]):
do_round = True

else:
# TODO: complex? what if result is already non-object?
dtype = "object"

dtype = np.dtype(dtype)
Expand Down Expand Up @@ -296,7 +292,9 @@ def maybe_downcast_to_dtype(result, dtype: Union[str, np.dtype]):
return result


def maybe_downcast_numeric(result, dtype: DtypeObj, do_round: bool = False):
def maybe_downcast_numeric(
result: ArrayLike, dtype: DtypeObj, do_round: bool = False
) -> ArrayLike:
"""
Subset of maybe_downcast_to_dtype restricted to numeric dtypes.

Expand Down