Add numpy array shapes for return types

Author: hamdanal

pandas-stubs/_libs/interval.pyi

@@ -21,7 +21,7 @@ from pandas.core.series import (
 from pandas._typing import (
     IntervalClosedType,
     IntervalT,
-    np_ndarray_bool,
+    np_1darray,
     npt,
 )
 
@@ -170,7 +170,9 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     @overload
     def __gt__(self, other: Interval[_OrderableT]) -> bool: ...
     @overload
-    def __gt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __gt__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __gt__(
         self,
@@ -179,7 +181,9 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     @overload
     def __lt__(self, other: Interval[_OrderableT]) -> bool: ...
     @overload
-    def __lt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __lt__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __lt__(
         self,
@@ -188,7 +192,9 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     @overload
     def __ge__(self, other: Interval[_OrderableT]) -> bool: ...
     @overload
-    def __ge__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __ge__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __ge__(
         self,
@@ -197,19 +203,25 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
     @overload
     def __le__(self, other: Interval[_OrderableT]) -> bool: ...
     @overload
-    def __le__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __le__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __eq__(self, other: Interval[_OrderableT]) -> bool: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     @overload
-    def __eq__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __eq__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __eq__(self, other: Series[_OrderableT]) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
     def __eq__(self, other: object) -> Literal[False]: ...
     @overload
     def __ne__(self, other: Interval[_OrderableT]) -> bool: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     @overload
-    def __ne__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
+    def __ne__(
+        self: IntervalT, other: IntervalIndex[IntervalT]
+    ) -> np_1darray[np.bool]: ...
     @overload
     def __ne__(self, other: Series[_OrderableT]) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload

pandas-stubs/_libs/tslibs/timestamps.pyi

@@ -40,9 +40,11 @@ from pandas._libs.tslibs import (
     Timedelta,
 )
 from pandas._typing import (
+    ShapeT,
     TimestampNonexistent,
     TimeUnit,
-    np_ndarray_bool,
+    np_1darray,
+    np_ndarray,
     npt,
 )
 
@@ -180,40 +182,48 @@ class Timestamp(datetime, SupportsIndex):
     @overload  # type: ignore[override]
     def __le__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
+    def __le__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ...
+    @overload
     def __le__(
-        self, other: DatetimeIndex | npt.NDArray[np.datetime64]
-    ) -> np_ndarray_bool: ...
+        self, other: np_ndarray[ShapeT, np.datetime64]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
     @overload
     def __le__(self, other: TimestampSeries) -> Series[bool]: ...
     @overload  # type: ignore[override]
     def __lt__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
+    def __lt__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ...
+    @overload
     def __lt__(
-        self, other: DatetimeIndex | npt.NDArray[np.datetime64]
-    ) -> np_ndarray_bool: ...
+        self, other: np_ndarray[ShapeT, np.datetime64]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
     @overload
     def __lt__(self, other: TimestampSeries) -> Series[bool]: ...
     @overload  # type: ignore[override]
     def __ge__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
+    def __ge__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ...
+    @overload
     def __ge__(
-        self, other: DatetimeIndex | npt.NDArray[np.datetime64]
-    ) -> np_ndarray_bool: ...
+        self, other: np_ndarray[ShapeT, np.datetime64]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
     @overload
     def __ge__(self, other: TimestampSeries) -> Series[bool]: ...
     @overload  # type: ignore[override]
     def __gt__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[misc]
     @overload
+    def __gt__(self, other: DatetimeIndex) -> np_1darray[np.bool]: ...
+    @overload
     def __gt__(
-        self, other: DatetimeIndex | npt.NDArray[np.datetime64]
-    ) -> np_ndarray_bool: ...
+        self, other: np_ndarray[ShapeT, np.datetime64]
+    ) -> np_ndarray[ShapeT, np.bool]: ...
     @overload
     def __gt__(self, other: TimestampSeries) -> Series[bool]: ...
     # error: Signature of "__add__" incompatible with supertype "date"/"datetime"
     @overload  # type: ignore[override]
     def __add__(
-        self, other: npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.datetime64]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.datetime64]: ...
     @overload
     def __add__(self, other: timedelta | np.timedelta64 | Tick) -> Self: ...
     @overload
@@ -226,8 +236,8 @@ class Timestamp(datetime, SupportsIndex):
     def __radd__(self, other: TimedeltaIndex) -> DatetimeIndex: ...
     @overload
     def __radd__(
-        self, other: npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.datetime64]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.datetime64]: ...
     # TODO: test dt64
     @overload  # type: ignore[override]
     def __sub__(self, other: Timestamp | datetime | np.datetime64) -> Timedelta: ...
@@ -241,22 +251,26 @@ class Timestamp(datetime, SupportsIndex):
     def __sub__(self, other: TimestampSeries) -> TimedeltaSeries: ...
     @overload
     def __sub__(
-        self, other: npt.NDArray[np.timedelta64]
-    ) -> npt.NDArray[np.datetime64]: ...
+        self, other: np_ndarray[ShapeT, np.timedelta64]
+    ) -> np_ndarray[ShapeT, np.datetime64]: ...
     @overload
     def __eq__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     @overload
     def __eq__(self, other: TimestampSeries) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
-    def __eq__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ...  # type: ignore[overload-overlap]
+    def __eq__(self, other: Index) -> np_1darray[np.bool]: ...  # type: ignore[overload-overlap]
+    @overload  # TODO: using shape-aware arrays similar to other methods doesn't work in mypy
+    def __eq__(self, other: npt.NDArray[np.datetime64]) -> npt.NDArray[np.bool]: ...  # type: ignore[overload-overlap]
     @overload
     def __eq__(self, other: object) -> Literal[False]: ...
     @overload
     def __ne__(self, other: Timestamp | datetime | np.datetime64) -> bool: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     @overload
     def __ne__(self, other: TimestampSeries) -> Series[bool]: ...  # type: ignore[overload-overlap]
     @overload
-    def __ne__(self, other: npt.NDArray[np.datetime64] | Index) -> np_ndarray_bool: ...  # type: ignore[overload-overlap]
+    def __ne__(self, other: Index) -> np_1darray[np.bool]: ...  # type: ignore[overload-overlap]
+    @overload  # TODO: using shape-aware arrays similar to other methods doesn't work in mypy
+    def __ne__(self, other: npt.NDArray[np.datetime64]) -> npt.NDArray[np.bool]: ...  # type: ignore[overload-overlap]
     @overload
     def __ne__(self, other: object) -> Literal[True]: ...
     def __hash__(self) -> int: ...

pandas-stubs/_typing.pyi

@@ -819,6 +819,15 @@ np_ndarray_complex: TypeAlias = npt.NDArray[np.complexfloating]
 np_ndarray_bool: TypeAlias = npt.NDArray[np.bool_]
 np_ndarray_str: TypeAlias = npt.NDArray[np.str_]
 
+# Define shape and generic type variables with defaults similar to numpy
+GenericT = TypeVar("GenericT", bound=np.generic, default=Any)
+ShapeT = TypeVar("ShapeT", bound=tuple[int, ...], default=tuple[Any, ...])
+# Numpy ndarray with more ergonomic typevar
+np_ndarray: TypeAlias = np.ndarray[ShapeT, np.dtype[GenericT]]
+# Numpy arrays with known shape (Do not use as argument types, only as return types)
+np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[GenericT]]
+np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]]
+
 IndexType: TypeAlias = slice | np_ndarray_anyint | Index | list[int] | Series[int]
 MaskType: TypeAlias = Series[bool] | np_ndarray_bool | list[bool]
 

pandas-stubs/core/algorithms.pyi

@@ -18,6 +18,7 @@ from pandas._typing import (
     AnyArrayLike,
     IntervalT,
     TakeIndexer,
+    np_1darray,
 )
 
 # These are type: ignored because the Index types overlap due to inheritance but indices
@@ -54,14 +55,14 @@ def factorize(
     sort: bool = ...,
     use_na_sentinel: bool = ...,
     size_hint: int | None = ...,
-) -> tuple[np.ndarray, Index]: ...
+) -> tuple[np_1darray, Index]: ...
 @overload
 def factorize(
     values: Categorical,
     sort: bool = ...,
     use_na_sentinel: bool = ...,
     size_hint: int | None = ...,
-) -> tuple[np.ndarray, Categorical]: ...
+) -> tuple[np_1darray, Categorical]: ...
 def value_counts(
     values: AnyArrayLike | list | tuple,
     sort: bool = True,

pandas-stubs/core/arrays/base.pyi

@@ -12,6 +12,7 @@ from pandas._typing import (
     ScalarIndexer,
     SequenceIndexer,
     TakeIndexer,
+    np_1darray,
     npt,
 )
 
@@ -31,7 +32,7 @@ class ExtensionArray:
         dtype: npt.DTypeLike | None = ...,
         copy: bool = False,
         na_value: Scalar = ...,
-    ) -> np.ndarray: ...
+    ) -> np_1darray: ...
     @property
     def dtype(self) -> ExtensionDtype: ...
     @property
@@ -44,13 +45,13 @@ class ExtensionArray:
     def isna(self) -> ArrayLike: ...
     def argsort(
         self, *, ascending: bool = ..., kind: str = ..., **kwargs
-    ) -> np.ndarray: ...
+    ) -> np_1darray: ...
     def fillna(self, value=..., method=None, limit=None): ...
     def dropna(self): ...
     def shift(self, periods: int = 1, fill_value: object = ...) -> Self: ...
     def unique(self): ...
     def searchsorted(self, value, side: str = ..., sorter=...): ...
-    def factorize(self, use_na_sentinel: bool = True) -> tuple[np.ndarray, Self]: ...
+    def factorize(self, use_na_sentinel: bool = True) -> tuple[np_1darray, Self]: ...
     def repeat(self, repeats, axis=...): ...
     def take(
         self,
@@ -60,7 +61,7 @@ class ExtensionArray:
         fill_value=...,
     ) -> Self: ...
     def copy(self) -> Self: ...
-    def view(self, dtype=...) -> Self | np.ndarray: ...
+    def view(self, dtype=...) -> Self | np_1darray: ...
     def ravel(self, order="C") -> Self: ...
     def tolist(self) -> list: ...
     def _reduce(

pandas-stubs/core/arrays/categorical.pyi

@@ -25,8 +25,7 @@ from pandas._typing import (
     ScalarIndexer,
     SequenceIndexer,
     TakeIndexer,
-    np_ndarray_bool,
-    np_ndarray_int,
+    np_1darray,
 )
 
 from pandas.core.dtypes.dtypes import CategoricalDtype as CategoricalDtype
@@ -63,7 +62,7 @@ class Categorical(ExtensionArray):
         fastpath: bool = ...,
     ) -> Categorical: ...
     @property
-    def codes(self) -> np_ndarray_int: ...
+    def codes(self) -> np_1darray[np.signedinteger]: ...
     def set_ordered(self, value) -> Categorical: ...
     def as_ordered(self) -> Categorical: ...
     def as_unordered(self) -> Categorical: ...
@@ -90,18 +89,18 @@ class Categorical(ExtensionArray):
     @property
     def shape(self): ...
     def shift(self, periods=1, fill_value=...): ...
-    def __array__(self, dtype=...) -> np.ndarray: ...
+    def __array__(self, dtype=...) -> np_1darray: ...
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): ...
     @property
     def T(self): ...
     @property
     def nbytes(self) -> int: ...
     def memory_usage(self, deep: bool = ...): ...
     def searchsorted(self, value, side: str = ..., sorter=...): ...
-    def isna(self) -> np_ndarray_bool: ...
-    def isnull(self) -> np_ndarray_bool: ...
-    def notna(self) -> np_ndarray_bool: ...
-    def notnull(self) -> np_ndarray_bool: ...
+    def isna(self) -> np_1darray[np.bool]: ...
+    def isnull(self) -> np_1darray[np.bool]: ...
+    def notna(self) -> np_1darray[np.bool]: ...
+    def notnull(self) -> np_1darray[np.bool]: ...
     def dropna(self): ...
     def value_counts(self, dropna: bool = True): ...
     def check_for_ordered(self, op) -> None: ...

pandas-stubs/core/arrays/interval.pyi

@@ -21,7 +21,7 @@ from pandas._typing import (
     ScalarIndexer,
     SequenceIndexer,
     TakeIndexer,
-    np_ndarray_bool,
+    np_1darray,
 )
 
 IntervalOrNA: TypeAlias = Interval | float
@@ -99,7 +99,7 @@ class IntervalArray(IntervalMixin, ExtensionArray):
     def mid(self) -> Index: ...
     @property
     def is_non_overlapping_monotonic(self) -> bool: ...
-    def __array__(self, dtype=...) -> np.ndarray: ...
+    def __array__(self, dtype=...) -> np_1darray: ...
     def __arrow_array__(self, type=...): ...
     def to_tuples(self, na_tuple: bool = True): ...
     def repeat(self, repeats, axis: Axis | None = ...): ...
@@ -108,5 +108,5 @@ class IntervalArray(IntervalMixin, ExtensionArray):
     @overload
     def contains(
         self, other: Scalar | ExtensionArray | Index | np.ndarray
-    ) -> np_ndarray_bool: ...
+    ) -> np_1darray[np.bool]: ...
     def overlaps(self, other: Interval) -> bool: ...

pandas-stubs/core/base.pyi

@@ -26,6 +26,7 @@ from pandas._typing import (
     DropKeep,
     NDFrameT,
     Scalar,
+    np_1darray,
     npt,
 )
 from pandas.util._decorators import cache_readonly
@@ -63,7 +64,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1]):
         copy: bool = False,
         na_value: Scalar = ...,
         **kwargs,
-    ) -> np.ndarray: ...
+    ) -> np_1darray: ...
     @property
     def empty(self) -> bool: ...
     def max(self, axis=..., skipna: bool = ..., **kwargs): ...
@@ -114,7 +115,7 @@ class IndexOpsMixin(OpsMixin, Generic[S1]):
     def is_monotonic_increasing(self) -> bool: ...
     def factorize(
         self, sort: bool = False, use_na_sentinel: bool = True
-    ) -> tuple[np.ndarray, np.ndarray | Index | Categorical]: ...
+    ) -> tuple[np_1darray, np_1darray | Index | Categorical]: ...
     def searchsorted(
         self, value, side: Literal["left", "right"] = ..., sorter=...
     ) -> int | list[int]: ...