diff --git a/pandas/_typing.py b/pandas/_typing.py index a85820a403fde..85ed2a3b636de 100644 --- a/pandas/_typing.py +++ b/pandas/_typing.py @@ -323,3 +323,6 @@ def closed(self) -> bool: # sort_index SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"] NaPosition = Literal["first", "last"] + +# quantile interpolation +QuantileInterpolation = Literal["linear", "lower", "higher", "midpoint", "nearest"] diff --git a/pandas/core/common.py b/pandas/core/common.py index 7225b26a910dd..707201153e44a 100644 --- a/pandas/core/common.py +++ b/pandas/core/common.py @@ -125,11 +125,11 @@ def is_bool_indexer(key: Any) -> bool: is_array_like(key) and is_extension_array_dtype(key.dtype) ): if key.dtype == np.object_: - key = np.asarray(key) + key_array = np.asarray(key) - if not lib.is_bool_array(key): + if not lib.is_bool_array(key_array): na_msg = "Cannot mask with non-boolean array containing NA / NaN values" - if lib.infer_dtype(key) == "boolean" and isna(key).any(): + if lib.infer_dtype(key_array) == "boolean" and isna(key_array).any(): # Don't raise on e.g. ["A", "B", np.nan], see # test_loc_getitem_list_of_labels_categoricalindex_with_na raise ValueError(na_msg) @@ -508,18 +508,14 @@ def get_rename_function(mapper): Returns a function that will map names/labels, dependent if mapper is a dict, Series or just a function. """ - if isinstance(mapper, (abc.Mapping, ABCSeries)): - def f(x): - if x in mapper: - return mapper[x] - else: - return x - - else: - f = mapper + def f(x): + if x in mapper: + return mapper[x] + else: + return x - return f + return f if isinstance(mapper, (abc.Mapping, ABCSeries)) else mapper def convert_to_list_like( diff --git a/pandas/core/series.py b/pandas/core/series.py index 6ef024f13fbb1..c2260afb2438e 100644 --- a/pandas/core/series.py +++ b/pandas/core/series.py @@ -33,6 +33,7 @@ from pandas._libs.lib import no_default from pandas._typing import ( AggFuncType, + AnyArrayLike, ArrayLike, Axis, Dtype, @@ -42,6 +43,7 @@ IndexKeyFunc, Level, NaPosition, + QuantileInterpolation, Renamer, SingleManager, SortKind, @@ -2478,7 +2480,33 @@ def round(self, decimals=0, *args, **kwargs) -> Series: return result - def quantile(self, q=0.5, interpolation="linear"): + @overload + def quantile( + self, q: float = ..., interpolation: QuantileInterpolation = ... + ) -> float: + ... + + @overload + def quantile( + self, + q: Sequence[float] | AnyArrayLike, + interpolation: QuantileInterpolation = ..., + ) -> Series: + ... + + @overload + def quantile( + self, + q: float | Sequence[float] | AnyArrayLike = ..., + interpolation: QuantileInterpolation = ..., + ) -> float | Series: + ... + + def quantile( + self, + q: float | Sequence[float] | AnyArrayLike = 0.5, + interpolation: QuantileInterpolation = "linear", + ) -> float | Series: """ Return value at the given quantile. diff --git a/pandas/io/formats/style.py b/pandas/io/formats/style.py index 24646da9162b0..24669e84443a6 100644 --- a/pandas/io/formats/style.py +++ b/pandas/io/formats/style.py @@ -25,6 +25,7 @@ FilePath, IndexLabel, Level, + QuantileInterpolation, Scalar, WriteBuffer, ) @@ -3467,7 +3468,7 @@ def highlight_quantile( axis: Axis | None = 0, q_left: float = 0.0, q_right: float = 1.0, - interpolation: str = "linear", + interpolation: QuantileInterpolation = "linear", inclusive: str = "both", props: str | None = None, ) -> Styler: @@ -3539,13 +3540,17 @@ def highlight_quantile( # after quantile is found along axis, e.g. along rows, # applying the calculated quantile to alternate axis, e.g. to each column - kwargs = {"q": [q_left, q_right], "interpolation": interpolation} + quantiles = [q_left, q_right] if axis is None: - q = Series(data.to_numpy().ravel()).quantile(**kwargs) + q = Series(data.to_numpy().ravel()).quantile( + q=quantiles, interpolation=interpolation + ) axis_apply: int | None = None else: axis = self.data._get_axis_number(axis) - q = data.quantile(axis=axis, numeric_only=False, **kwargs) + q = data.quantile( + axis=axis, numeric_only=False, q=quantiles, interpolation=interpolation + ) axis_apply = 1 - axis if props is None: diff --git a/pyright_reportGeneralTypeIssues.json b/pyright_reportGeneralTypeIssues.json index ab1782af51260..248e4d1c13f23 100644 --- a/pyright_reportGeneralTypeIssues.json +++ b/pyright_reportGeneralTypeIssues.json @@ -35,10 +35,8 @@ "pandas/core/arrays/string_.py", "pandas/core/arrays/string_arrow.py", "pandas/core/arrays/timedeltas.py", - "pandas/core/common.py", "pandas/core/computation/align.py", "pandas/core/construction.py", - "pandas/core/describe.py", "pandas/core/dtypes/cast.py", "pandas/core/dtypes/common.py", "pandas/core/dtypes/concat.py",