From 935b1f96d4991b9ad923e06278a24f2ce02c6b5c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Fri, 10 Jun 2022 09:43:17 -0400 Subject: [PATCH 1/2] TYP: Series.quantile --- pandas/_typing.py | 3 +++ pandas/core/common.py | 6 +++--- pandas/core/series.py | 30 +++++++++++++++++++++++++++- pandas/io/formats/style.py | 13 ++++++++---- pyright_reportGeneralTypeIssues.json | 1 - 5 files changed, 44 insertions(+), 9 deletions(-) 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..a86999699ca64 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) 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 a6b5715723a24..8b79ea5f3f453 100644 --- a/pyright_reportGeneralTypeIssues.json +++ b/pyright_reportGeneralTypeIssues.json @@ -38,7 +38,6 @@ "pandas/core/common.py", "pandas/core/computation/align.py", "pandas/core/construction.py", - "pandas/core/describe.py", "pandas/core/dtypes/base.py", "pandas/core/dtypes/cast.py", "pandas/core/dtypes/common.py", From 3853e6084aa35cf92afb28ad5f9aa453d0c76f03 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Fri, 10 Jun 2022 16:22:40 -0400 Subject: [PATCH 2/2] common.py --- pandas/core/common.py | 16 ++++++---------- pyright_reportGeneralTypeIssues.json | 1 - 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/pandas/core/common.py b/pandas/core/common.py index a86999699ca64..707201153e44a 100644 --- a/pandas/core/common.py +++ b/pandas/core/common.py @@ -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/pyright_reportGeneralTypeIssues.json b/pyright_reportGeneralTypeIssues.json index 1aedd4b5c6bc6..248e4d1c13f23 100644 --- a/pyright_reportGeneralTypeIssues.json +++ b/pyright_reportGeneralTypeIssues.json @@ -35,7 +35,6 @@ "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/dtypes/cast.py",