GH1222 Candidate for ty on pandas-stubs (#1240)

* GH1222 Candidate for ty on pandas-stubs

* GH1222 Fix pre-commit

* Update pandas-stubs/core/strings.pyi

Co-authored-by: Irv Lustig <[email protected]>

* GH1222 PR feedback

* GH1222 PR feedback

---------

Co-authored-by: Irv Lustig <[email protected]>

Author: loicdiridollou

.github/workflows/test.yml

@@ -35,6 +35,9 @@ jobs:
       - name: Run mypy on 'tests' (using the local stubs) and on the local stubs
         run: poetry run poe mypy
 
+      - name: Run ty on 'pandas-stubs' (using the local stubs) and on the local stubs
+        run: poetry run poe ty
+
       - name: Run pyright on 'tests' (using the local stubs) and on the local stubs
         run: poetry run poe pyright
 

pandas-stubs/_typing.pyi

@@ -19,6 +19,7 @@ from typing import (
     SupportsIndex,
     TypedDict,
     TypeVar,
+    Union,
     overload,
 )
 
@@ -463,6 +464,10 @@ VoidDtypeArg: TypeAlias = (
     | Literal["V", "void", "void0"]
 )
 
+# DtypeArg specifies all allowable dtypes in a functions its dtype argument
+DtypeArg: TypeAlias = Dtype | Mapping[Hashable, Dtype]
+DtypeObj: TypeAlias = np.dtype[np.generic] | ExtensionDtype
+
 AstypeArg: TypeAlias = (
     BooleanDtypeArg
     | IntDtypeArg
@@ -479,10 +484,6 @@ AstypeArg: TypeAlias = (
     | DtypeObj
 )
 
-# DtypeArg specifies all allowable dtypes in a functions its dtype argument
-DtypeArg: TypeAlias = Dtype | Mapping[Hashable, Dtype]
-DtypeObj: TypeAlias = np.dtype[np.generic] | ExtensionDtype
-
 # converters
 ConvertersArg: TypeAlias = Mapping[Hashable, Callable[[Dtype], Dtype]]
 
@@ -513,7 +514,8 @@ IndexKeyFunc: TypeAlias = Callable[[Index], Index | AnyArrayLike] | None
 
 # types of `func` kwarg for DataFrame.aggregate and Series.aggregate
 # More specific than what is in pandas
-AggFuncTypeBase: TypeAlias = Callable | str | np.ufunc
+# following Union is here to make it ty compliant https://github.com/astral-sh/ty/issues/591
+AggFuncTypeBase: TypeAlias = Union[Callable, str, np.ufunc]  # noqa: UP007
 AggFuncTypeDictSeries: TypeAlias = Mapping[HashableT, AggFuncTypeBase]
 AggFuncTypeDictFrame: TypeAlias = Mapping[
     HashableT, AggFuncTypeBase | list[AggFuncTypeBase]

pandas-stubs/core/arrays/datetimes.pyi

@@ -1,4 +1,4 @@
-from datetime import tzinfo
+from datetime import tzinfo as _tzinfo
 
 import numpy as np
 from pandas.core.arrays.datetimelike import (
@@ -28,7 +28,7 @@ class DatetimeArray(DatetimeLikeArrayMixin, TimelikeOps, DatelikeOps):
     @tz.setter
     def tz(self, value) -> None: ...
     @property
-    def tzinfo(self) -> tzinfo | None: ...
+    def tzinfo(self) -> _tzinfo | None: ...
     @property
     def is_normalized(self): ...
     def __array__(self, dtype=...) -> np.ndarray: ...

pandas-stubs/core/generic.pyi

@@ -69,8 +69,8 @@ class NDFrame(indexing.IndexingMixin):
     def set_flags(
         self,
         *,
-        copy: bool = ...,
-        allows_duplicate_labels: bool | None = ...,
+        copy: _bool = ...,
+        allows_duplicate_labels: _bool | None = ...,
     ) -> Self: ...
     @property
     def attrs(self) -> dict[Hashable | None, Any]: ...

pandas-stubs/core/groupby/groupby.pyi

@@ -75,8 +75,6 @@ from pandas._typing import (
 
 from pandas.plotting import PlotAccessor
 
-_GroupByT = TypeVar("_GroupByT", bound=GroupBy)
-
 _KeysArgType: TypeAlias = (
     Hashable
     | list[Hashable]
@@ -91,67 +89,6 @@ _ResamplerGroupBy: TypeAlias = (
     | TimedeltaIndexResamplerGroupby[NDFrameT]
 )
 
-# GroupByPlot does not really inherit from PlotAccessor but it delegates
-# to it using __call__ and __getattr__. We lie here to avoid repeating the
-# whole stub of PlotAccessor
-@final
-class GroupByPlot(PlotAccessor, Generic[_GroupByT]):
-    def __init__(self, groupby: _GroupByT) -> None: ...
-    # The following methods are inherited from the fake parent class PlotAccessor
-    # def __call__(self, *args, **kwargs): ...
-    # def __getattr__(self, name: str): ...
-
-class BaseGroupBy(SelectionMixin[NDFrameT], GroupByIndexingMixin):
-    axis: AxisInt
-    grouper: ops.BaseGrouper
-    keys: _KeysArgType | None
-    level: IndexLabel | None
-    group_keys: bool
-    @final
-    def __len__(self) -> int: ...
-    @final
-    def __repr__(self) -> str: ...  # noqa: PYI029 __repr__ here is final
-    @final
-    @property
-    def groups(self) -> dict[Hashable, Index]: ...
-    @final
-    @property
-    def ngroups(self) -> int: ...
-    @final
-    @property
-    def indices(self) -> dict[Hashable, Index | npt.NDArray[np.int_] | list[int]]: ...
-    @overload
-    def pipe(
-        self,
-        func: Callable[Concatenate[Self, P], T],
-        *args: P.args,
-        **kwargs: P.kwargs,
-    ) -> T: ...
-    @overload
-    def pipe(
-        self,
-        func: tuple[Callable[..., T], str],
-        *args: Any,
-        **kwargs: Any,
-    ) -> T: ...
-    @final
-    def get_group(self, name, obj: NDFrameT | None = ...) -> NDFrameT: ...
-    @final
-    def __iter__(self) -> Iterator[tuple[Hashable, NDFrameT]]: ...
-    @overload
-    def __getitem__(self: BaseGroupBy[DataFrame], key: Scalar) -> generic.SeriesGroupBy: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
-    @overload
-    def __getitem__(
-        self: BaseGroupBy[DataFrame], key: Iterable[Hashable]
-    ) -> generic.DataFrameGroupBy: ...
-    @overload
-    def __getitem__(
-        self: BaseGroupBy[Series[S1]],
-        idx: list[str] | Index | Series[S1] | MaskType | tuple[Hashable | slice, ...],
-    ) -> generic.SeriesGroupBy: ...
-    @overload
-    def __getitem__(self: BaseGroupBy[Series[S1]], idx: Scalar) -> S1: ...
-
 class GroupBy(BaseGroupBy[NDFrameT]):
     as_index: bool
     sort: bool
@@ -393,3 +330,66 @@ class GroupBy(BaseGroupBy[NDFrameT]):
         weights: Sequence | Series | None = ...,
         random_state: RandomState | None = ...,
     ) -> NDFrameT: ...
+
+_GroupByT = TypeVar("_GroupByT", bound=GroupBy)
+
+# GroupByPlot does not really inherit from PlotAccessor but it delegates
+# to it using __call__ and __getattr__. We lie here to avoid repeating the
+# whole stub of PlotAccessor
+@final
+class GroupByPlot(PlotAccessor, Generic[_GroupByT]):
+    def __init__(self, groupby: _GroupByT) -> None: ...
+    # The following methods are inherited from the fake parent class PlotAccessor
+    # def __call__(self, *args, **kwargs): ...
+    # def __getattr__(self, name: str): ...
+
+class BaseGroupBy(SelectionMixin[NDFrameT], GroupByIndexingMixin):
+    axis: AxisInt
+    grouper: ops.BaseGrouper
+    keys: _KeysArgType | None
+    level: IndexLabel | None
+    group_keys: bool
+    @final
+    def __len__(self) -> int: ...
+    @final
+    def __repr__(self) -> str: ...  # noqa: PYI029 __repr__ here is final
+    @final
+    @property
+    def groups(self) -> dict[Hashable, Index]: ...
+    @final
+    @property
+    def ngroups(self) -> int: ...
+    @final
+    @property
+    def indices(self) -> dict[Hashable, Index | npt.NDArray[np.int_] | list[int]]: ...
+    @overload
+    def pipe(
+        self,
+        func: Callable[Concatenate[Self, P], T],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> T: ...
+    @overload
+    def pipe(
+        self,
+        func: tuple[Callable[..., T], str],
+        *args: Any,
+        **kwargs: Any,
+    ) -> T: ...
+    @final
+    def get_group(self, name, obj: NDFrameT | None = ...) -> NDFrameT: ...
+    @final
+    def __iter__(self) -> Iterator[tuple[Hashable, NDFrameT]]: ...
+    @overload
+    def __getitem__(self: BaseGroupBy[DataFrame], key: Scalar) -> generic.SeriesGroupBy: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
+    @overload
+    def __getitem__(
+        self: BaseGroupBy[DataFrame], key: Iterable[Hashable]
+    ) -> generic.DataFrameGroupBy: ...
+    @overload
+    def __getitem__(
+        self: BaseGroupBy[Series[S1]],
+        idx: list[str] | Index | Series[S1] | MaskType | tuple[Hashable | slice, ...],
+    ) -> generic.SeriesGroupBy: ...
+    @overload
+    def __getitem__(self: BaseGroupBy[Series[S1]], idx: Scalar) -> S1: ...

pandas-stubs/core/indexes/accessors.pyi

@@ -1,7 +1,7 @@
 import datetime as dt
 from datetime import (
     timedelta,
-    tzinfo,
+    tzinfo as _tzinfo,
 )
 from typing import (
     Generic,
@@ -119,7 +119,7 @@ class _FreqProperty(Generic[_DTFreqReturnType]):
 
 class _TZProperty:
     @property
-    def tz(self) -> tzinfo | None: ...
+    def tz(self) -> _tzinfo | None: ...
 
 class _DatetimeObjectOps(
     _FreqProperty[_DTFreqReturnType], _TZProperty, Generic[_DTFreqReturnType]
@@ -416,7 +416,7 @@ class DatetimeIndexProperties(
     @property
     def is_normalized(self) -> bool: ...
     @property
-    def tzinfo(self) -> tzinfo | None: ...
+    def tzinfo(self) -> _tzinfo | None: ...
     def to_pydatetime(self) -> npt.NDArray[np.object_]: ...
     def std(
         self, axis: int | None = ..., ddof: int = ..., skipna: bool = ...

pandas-stubs/core/indexes/base.pyi

@@ -272,10 +272,10 @@ class Index(IndexOpsMixin[S1]):
         Self,
         MultiIndex,
         np_ndarray_bool,
-        Index[list[str]],
+        Index[list[_str]],
         Index[int],
         Index[bytes],
-        Index[str],
+        Index[_str],
         Index[type[object]],
     ]: ...
     def is_(self, other) -> bool: ...
@@ -478,7 +478,7 @@ class Index(IndexOpsMixin[S1]):
 UnknownIndex: TypeAlias = Index[Any]
 
 def ensure_index_from_sequences(
-    sequences: Sequence[Sequence[Dtype]], names: list[str] = ...
+    sequences: Sequence[Sequence[Dtype]], names: list[_str] = ...
 ) -> Index: ...
 def ensure_index(index_like: Sequence | Index, copy: bool = ...) -> Index: ...
 def maybe_extract_name(name, obj, cls) -> Label: ...

pandas-stubs/core/indexes/datetimes.pyi

@@ -5,7 +5,7 @@ from collections.abc import (
 from datetime import (
     datetime,
     timedelta,
-    tzinfo,
+    tzinfo as _tzinfo,
 )
 from typing import overload
 
@@ -85,7 +85,7 @@ class DatetimeIndex(DatetimeTimedeltaMixin[Timestamp], DatetimeIndexProperties):
     def to_julian_date(self) -> Index[float]: ...
     def isocalendar(self) -> DataFrame: ...
     @property
-    def tzinfo(self) -> tzinfo | None: ...
+    def tzinfo(self) -> _tzinfo | None: ...
     @property
     def dtype(self) -> np.dtype | DatetimeTZDtype: ...
     def shift(self, periods: int = ..., freq=...) -> Self: ...