From 1be1015d6e5afcdfafcb65cbd53f422653d113f7 Mon Sep 17 00:00:00 2001 From: aram Date: Tue, 25 Jun 2024 20:09:44 +0400 Subject: [PATCH 1/3] Typing improvements for Index --- pandas/core/indexes/base.py | 36 ++++++++++++++++++++++++------------ 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 71dfff520113c..dfc4ae34bf36a 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -45,6 +45,7 @@ ArrayLike, Axes, Axis, + AxisInt, DropKeep, DtypeObj, F, @@ -1087,7 +1088,7 @@ def view(self, cls=None): result._id = self._id return result - def astype(self, dtype, copy: bool = True): + def astype(self, dtype: DtypeObj, copy: bool = True): """ Create an Index with values cast to dtypes. @@ -1298,7 +1299,7 @@ def _maybe_disallow_fill(self, allow_fill: bool, fill_value, indices) -> bool: """ @Appender(_index_shared_docs["repeat"] % _index_doc_kwargs) - def repeat(self, repeats, axis: None = None) -> Self: + def repeat(self, repeats, axis: AxisInt | None = None) -> Self: repeats = ensure_platform_int(repeats) nv.validate_repeat((), {"axis": axis}) res_values = self._values.repeat(repeats) @@ -2957,7 +2958,7 @@ def _dti_setop_align_tzs(self, other: Index, setop: str_t) -> tuple[Index, Index return self, other @final - def union(self, other, sort=None): + def union(self, other, sort: bool | None = None): """ Form the union of two Index objects. @@ -3334,7 +3335,7 @@ def _intersection_via_get_indexer( return result @final - def difference(self, other, sort=None): + def difference(self, other, sort: bool | None = None): """ Return a new Index with elements of index not in `other`. @@ -3420,7 +3421,9 @@ def _wrap_difference_result(self, other, result): # We will override for MultiIndex to handle empty results return self._wrap_setop_result(other, result) - def symmetric_difference(self, other, result_name=None, sort=None): + def symmetric_difference( + self, other, result_name: str | None = None, sort: bool | None = None + ): """ Compute the symmetric difference of two Index objects. @@ -6389,7 +6392,7 @@ def _transform_index(self, func, *, level=None) -> Index: items = [func(x) for x in self] return Index(items, name=self.name, tupleize_cols=False) - def isin(self, values, level=None) -> npt.NDArray[np.bool_]: + def isin(self, values, level: str | int | None = None) -> npt.NDArray[np.bool_]: """ Return a boolean array where the index values are in `values`. @@ -6687,7 +6690,12 @@ def get_slice_bound(self, label, side: Literal["left", "right"]) -> int: else: return slc - def slice_locs(self, start=None, end=None, step=None) -> tuple[int, int]: + def slice_locs( + self, + start: Hashable | None = None, + end: Hashable | None = None, + step: int | None = None, + ) -> tuple[int, int]: """ Compute slice locations for input labels. @@ -6781,7 +6789,7 @@ def slice_locs(self, start=None, end=None, step=None) -> tuple[int, int]: return start_slice, end_slice - def delete(self, loc) -> Self: + def delete(self, loc: int | list[int]) -> Self: """ Make new Index with passed location(-s) deleted. @@ -7227,7 +7235,9 @@ def _maybe_disable_logical_methods(self, opname: str_t) -> None: raise TypeError(f"cannot perform {opname} with {type(self).__name__}") @Appender(IndexOpsMixin.argmin.__doc__) - def argmin(self, axis=None, skipna: bool = True, *args, **kwargs) -> int: + def argmin( + self, axis: AxisInt | None = None, skipna: bool = True, *args, **kwargs + ) -> int: nv.validate_argmin(args, kwargs) nv.validate_minmax_axis(axis) @@ -7240,7 +7250,9 @@ def argmin(self, axis=None, skipna: bool = True, *args, **kwargs) -> int: return super().argmin(skipna=skipna) @Appender(IndexOpsMixin.argmax.__doc__) - def argmax(self, axis=None, skipna: bool = True, *args, **kwargs) -> int: + def argmax( + self, axis: AxisInt | None = None, skipna: bool = True, *args, **kwargs + ) -> int: nv.validate_argmax(args, kwargs) nv.validate_minmax_axis(axis) @@ -7251,7 +7263,7 @@ def argmax(self, axis=None, skipna: bool = True, *args, **kwargs) -> int: raise ValueError("Encountered all NA values") return super().argmax(skipna=skipna) - def min(self, axis=None, skipna: bool = True, *args, **kwargs): + def min(self, axis: AxisInt | None = None, skipna: bool = True, *args, **kwargs): """ Return the minimum value of the Index. @@ -7314,7 +7326,7 @@ def min(self, axis=None, skipna: bool = True, *args, **kwargs): return nanops.nanmin(self._values, skipna=skipna) - def max(self, axis=None, skipna: bool = True, *args, **kwargs): + def max(self, axis: AxisInt | None = None, skipna: bool = True, *args, **kwargs): """ Return the maximum value of the Index. From 83a850b415db48e4bcac644ef5e3e814814a886f Mon Sep 17 00:00:00 2001 From: aram Date: Tue, 25 Jun 2024 20:57:31 +0400 Subject: [PATCH 2/3] fixed type hint for argument in astype --- pandas/core/indexes/base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index dfc4ae34bf36a..a04f2a353ff64 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -47,6 +47,7 @@ Axis, AxisInt, DropKeep, + DtypeArg, DtypeObj, F, IgnoreRaise, @@ -1088,7 +1089,7 @@ def view(self, cls=None): result._id = self._id return result - def astype(self, dtype: DtypeObj, copy: bool = True): + def astype(self, dtype: DtypeArg, copy: bool = True): """ Create an Index with values cast to dtypes. From 826336330cee56f3d5eb1b316ed13151635c54dc Mon Sep 17 00:00:00 2001 From: aram Date: Tue, 25 Jun 2024 20:58:46 +0400 Subject: [PATCH 3/3] fixed argument type for `result_name` for symmetric_difference --- pandas/core/indexes/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index a04f2a353ff64..36812afdbec5f 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -3423,7 +3423,7 @@ def _wrap_difference_result(self, other, result): return self._wrap_setop_result(other, result) def symmetric_difference( - self, other, result_name: str | None = None, sort: bool | None = None + self, other, result_name: Hashable | None = None, sort: bool | None = None ): """ Compute the symmetric difference of two Index objects.