Skip to content

GH1222 Candidate for ty on pandas-stubs #1240

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jun 6, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
12 changes: 7 additions & 5 deletions pandas-stubs/_typing.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ from typing import (
SupportsIndex,
TypedDict,
TypeVar,
Union,
overload,
)

Expand Down Expand Up @@ -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
Expand All @@ -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]]

Expand Down Expand Up @@ -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]
Expand Down
4 changes: 2 additions & 2 deletions pandas-stubs/core/arrays/datetimes.pyi
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from datetime import tzinfo
from datetime import tzinfo as _tzinfo

import numpy as np
from pandas.core.arrays.datetimelike import (
Expand Down Expand Up @@ -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: ...
Expand Down
4 changes: 2 additions & 2 deletions pandas-stubs/core/generic.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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]: ...
Expand Down
126 changes: 63 additions & 63 deletions pandas-stubs/core/groupby/groupby.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,6 @@ from pandas._typing import (

from pandas.plotting import PlotAccessor

_GroupByT = TypeVar("_GroupByT", bound=GroupBy)

_KeysArgType: TypeAlias = (
Hashable
| list[Hashable]
Expand All @@ -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
Expand Down Expand Up @@ -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: ...
6 changes: 3 additions & 3 deletions pandas-stubs/core/indexes/accessors.pyi
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import datetime as dt
from datetime import (
timedelta,
tzinfo,
tzinfo as _tzinfo,
)
from typing import (
Generic,
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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 = ...
Expand Down
6 changes: 3 additions & 3 deletions pandas-stubs/core/indexes/base.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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: ...
Expand Down Expand Up @@ -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: ...
4 changes: 2 additions & 2 deletions pandas-stubs/core/indexes/datetimes.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ from collections.abc import (
from datetime import (
datetime,
timedelta,
tzinfo,
tzinfo as _tzinfo,
)
from typing import overload

Expand Down Expand Up @@ -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: ...
Expand Down
Loading