From ed68fdada5779aef1af3c4e3573238ccd7178346 Mon Sep 17 00:00:00 2001 From: Brock Date: Mon, 5 Apr 2021 20:24:02 -0700 Subject: [PATCH 1/2] TYP: implement typing.Positional --- pandas/_typing.py | 8 ++++++++ pandas/core/arrays/_mixins.py | 4 +++- pandas/core/arrays/base.py | 5 ++--- pandas/core/arrays/datetimelike.py | 3 ++- pandas/core/arrays/masked.py | 5 ++--- pandas/core/arrays/string_arrow.py | 3 ++- pandas/core/internals/blocks.py | 4 +--- 7 files changed, 20 insertions(+), 12 deletions(-) diff --git a/pandas/_typing.py b/pandas/_typing.py index f90ef33434773..ca3a529842614 100644 --- a/pandas/_typing.py +++ b/pandas/_typing.py @@ -187,3 +187,11 @@ # internals Manager = Union["ArrayManager", "BlockManager"] SingleManager = Union["SingleArrayManager", "SingleBlockManager"] + +# indexing +# Positional -> valid 1D positional indexer, e.g. can pass to ndarray.__getitem__ +# TODO: add Ellipsis, see +# https://github.com/python/typing/issues/684#issuecomment-548203158 +# https://bugs.python.org/issue41810 +Positional = Union[int, np.integer, slice, Sequence[int], np.ndarray] +Positional2D = Union[Positional, Tuple[Positional, Positional]] diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 08061eb1ec28c..24cde76673788 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -16,6 +16,7 @@ from pandas._libs import lib from pandas._typing import ( F, + Positional2D, Shape, ) from pandas.compat.numpy import function as nv @@ -277,7 +278,8 @@ def _validate_setitem_value(self, value): return value def __getitem__( - self: NDArrayBackedExtensionArrayT, key: Union[int, slice, np.ndarray] + self: NDArrayBackedExtensionArrayT, + key: Positional2D, ) -> Union[NDArrayBackedExtensionArrayT, Any]: if lib.is_integer(key): # fast-path diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index c45528d657404..6ba101ce720b7 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -29,6 +29,7 @@ from pandas._typing import ( ArrayLike, Dtype, + Positional, Shape, ) from pandas.compat import set_function_name @@ -293,9 +294,7 @@ def _from_factorized(cls, values, original): # Must be a Sequence # ------------------------------------------------------------------------ - def __getitem__( - self, item: Union[int, slice, np.ndarray] - ) -> Union[ExtensionArray, Any]: + def __getitem__(self, item: Positional) -> Union[ExtensionArray, Any]: """ Select a subset of self. diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 25939bcdc7c6a..238bf10fa8daf 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -50,6 +50,7 @@ Dtype, DtypeObj, NpDtype, + Positional2D, ) from pandas.compat.numpy import function as nv from pandas.errors import ( @@ -341,7 +342,7 @@ def __array__(self, dtype: Optional[NpDtype] = None) -> np.ndarray: return self._ndarray def __getitem__( - self, key: Union[int, slice, np.ndarray] + self, key: Positional2D ) -> Union[DatetimeLikeArrayMixin, DTScalarOrNaT]: """ This getitem defers to the underlying array, which by-definition can diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index 31d58d9d89d49..16c0c513c8e6d 100644 --- a/pandas/core/arrays/masked.py +++ b/pandas/core/arrays/masked.py @@ -21,6 +21,7 @@ ArrayLike, Dtype, NpDtype, + Positional, Scalar, ) from pandas.errors import AbstractMethodError @@ -137,9 +138,7 @@ def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False): def dtype(self) -> BaseMaskedDtype: raise AbstractMethodError(self) - def __getitem__( - self, item: Union[int, slice, np.ndarray] - ) -> Union[BaseMaskedArray, Any]: + def __getitem__(self, item: Positional) -> Union[BaseMaskedArray, Any]: if is_integer(item): if self._mask[item]: return self.dtype.na_value diff --git a/pandas/core/arrays/string_arrow.py b/pandas/core/arrays/string_arrow.py index e1262d691128f..e47da8bce27b3 100644 --- a/pandas/core/arrays/string_arrow.py +++ b/pandas/core/arrays/string_arrow.py @@ -21,6 +21,7 @@ from pandas._typing import ( Dtype, NpDtype, + Positional, ) from pandas.util._decorators import doc from pandas.util._validators import validate_fillna_kwargs @@ -315,7 +316,7 @@ def _concat_same_type(cls, to_concat) -> ArrowStringArray: ) ) - def __getitem__(self, item: Any) -> Any: + def __getitem__(self, item: Positional) -> Any: """Select a subset of self. Parameters diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index 9a2b3be4b66e2..275c5218c3cb7 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -1394,9 +1394,7 @@ def iget(self, col): elif isinstance(col, slice): if col != slice(None): raise NotImplementedError(col) - # error: Invalid index type "List[Any]" for "ExtensionArray"; expected - # type "Union[int, slice, ndarray]" - return self.values[[loc]] # type: ignore[index] + return self.values[[loc]] return self.values[loc] else: if col != 0: From 0ef129343046103681fd5bf3f02da53cdd1b3af8 Mon Sep 17 00:00:00 2001 From: Brock Date: Wed, 7 Apr 2021 13:15:18 -0700 Subject: [PATCH 2/2] Positional->PositionalIndexere --- pandas/_typing.py | 9 ++++++--- pandas/core/arrays/_mixins.py | 4 ++-- pandas/core/arrays/base.py | 4 ++-- pandas/core/arrays/datetimelike.py | 6 ++++-- pandas/core/arrays/masked.py | 4 ++-- pandas/core/arrays/string_arrow.py | 4 ++-- 6 files changed, 18 insertions(+), 13 deletions(-) diff --git a/pandas/_typing.py b/pandas/_typing.py index 6ba17138e8fa6..d6561176deb71 100644 --- a/pandas/_typing.py +++ b/pandas/_typing.py @@ -189,9 +189,12 @@ SingleManager = Union["SingleArrayManager", "SingleBlockManager"] # indexing -# Positional -> valid 1D positional indexer, e.g. can pass to ndarray.__getitem__ +# PositionalIndexer -> valid 1D positional indexer, e.g. can pass +# to ndarray.__getitem__ # TODO: add Ellipsis, see # https://github.com/python/typing/issues/684#issuecomment-548203158 # https://bugs.python.org/issue41810 -Positional = Union[int, np.integer, slice, Sequence[int], np.ndarray] -Positional2D = Union[Positional, Tuple[Positional, Positional]] +PositionalIndexer = Union[int, np.integer, slice, Sequence[int], np.ndarray] +PositionalIndexer2D = Union[ + PositionalIndexer, Tuple[PositionalIndexer, PositionalIndexer] +] diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 3eea7cf2a7133..d4e5ca00b06dd 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -13,7 +13,7 @@ from pandas._libs import lib from pandas._typing import ( F, - Positional2D, + PositionalIndexer2D, Shape, ) from pandas.compat.numpy import function as nv @@ -276,7 +276,7 @@ def _validate_setitem_value(self, value): def __getitem__( self: NDArrayBackedExtensionArrayT, - key: Positional2D, + key: PositionalIndexer2D, ) -> NDArrayBackedExtensionArrayT | Any: if lib.is_integer(key): # fast-path diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index 89ca8ef415cc6..933b829e0b29f 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -24,7 +24,7 @@ from pandas._typing import ( ArrayLike, Dtype, - Positional, + PositionalIndexer, Shape, ) from pandas.compat import set_function_name @@ -289,7 +289,7 @@ def _from_factorized(cls, values, original): # Must be a Sequence # ------------------------------------------------------------------------ - def __getitem__(self, item: Positional) -> ExtensionArray | Any: + def __getitem__(self, item: PositionalIndexer) -> ExtensionArray | Any: """ Select a subset of self. diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index dbf32262e22c6..b039df12ee3df 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -47,7 +47,7 @@ Dtype, DtypeObj, NpDtype, - Positional2D, + PositionalIndexer2D, ) from pandas.compat.numpy import function as nv from pandas.errors import ( @@ -338,7 +338,9 @@ def __array__(self, dtype: NpDtype | None = None) -> np.ndarray: return np.array(list(self), dtype=object) return self._ndarray - def __getitem__(self, key: Positional2D) -> DatetimeLikeArrayMixin | DTScalarOrNaT: + def __getitem__( + self, key: PositionalIndexer2D + ) -> DatetimeLikeArrayMixin | DTScalarOrNaT: """ This getitem defers to the underlying array, which by-definition can only handle list-likes, slices, and integer scalars diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index 98c24267bcaab..93de1cd91d625 100644 --- a/pandas/core/arrays/masked.py +++ b/pandas/core/arrays/masked.py @@ -17,7 +17,7 @@ ArrayLike, Dtype, NpDtype, - Positional, + PositionalIndexer, Scalar, type_t, ) @@ -135,7 +135,7 @@ def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False): def dtype(self) -> BaseMaskedDtype: raise AbstractMethodError(self) - def __getitem__(self, item: Positional) -> BaseMaskedArray | Any: + def __getitem__(self, item: PositionalIndexer) -> BaseMaskedArray | Any: if is_integer(item): if self._mask[item]: return self.dtype.na_value diff --git a/pandas/core/arrays/string_arrow.py b/pandas/core/arrays/string_arrow.py index 3c2179f2ee7c3..faca868873efa 100644 --- a/pandas/core/arrays/string_arrow.py +++ b/pandas/core/arrays/string_arrow.py @@ -17,7 +17,7 @@ from pandas._typing import ( Dtype, NpDtype, - Positional, + PositionalIndexer, type_t, ) from pandas.util._decorators import doc @@ -311,7 +311,7 @@ def _concat_same_type(cls, to_concat) -> ArrowStringArray: ) ) - def __getitem__(self, item: Positional) -> Any: + def __getitem__(self, item: PositionalIndexer) -> Any: """Select a subset of self. Parameters