Merge remote-tracking branch 'upstream/main' into typevar-default

Author: MarcoGorelli

.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/_config/config.pyi

@@ -48,31 +48,31 @@ class DisplayUnicode(DictWrapper):
     east_asian_width: bool
 
 class Display(DictWrapper):
-    chop_threshold: int | None
-    colheader_justify: str
+    chop_threshold: float | None
+    colheader_justify: Literal["left", "right"]
     date_dayfirst: bool
     date_yearfirst: bool
     encoding: str
     expand_frame_repr: bool
     float_format: Callable[[float], str] | None
     html: DisplayHTML
-    large_repr: str
+    large_repr: Literal["truncate", "info"]
     latex: DisplayLaTeX
     max_categories: int
-    max_columns: int
-    max_colwidth: int
-    max_dir_items: int
+    max_columns: int | None
+    max_colwidth: int | None
+    max_dir_items: int | None
     max_info_columns: int
     max_info_rows: int
-    max_rows: int
-    max_seq_items: int
-    memory_usage: bool
-    min_rows: int
+    max_rows: int | None
+    max_seq_items: int | None
+    memory_usage: bool | Literal["deep"] | None
+    min_rows: int | None
     multi_sparse: bool
     notebook_repr_html: bool
     pprint_nest_depth: int
     precision: int
-    show_dimensions: str
+    show_dimensions: bool | Literal["truncate"]
     unicode: DisplayUnicode
     width: int
 

pandas-stubs/_typing.pyi

@@ -18,6 +18,7 @@ from typing import (
     Protocol,
     SupportsIndex,
     TypedDict,
+    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]
@@ -993,7 +995,7 @@ TimeZones: TypeAlias = str | tzinfo | None | int
 
 # Evaluates to a DataFrame column in DataFrame.assign context.
 IntoColumn: TypeAlias = (
-    AnyArrayLike | Scalar | Callable[[DataFrame], AnyArrayLike | Scalar]
+    AnyArrayLike | Scalar | Callable[[DataFrame], AnyArrayLike | Scalar] | None
 )
 
 DatetimeLike: TypeAlias = datetime.datetime | np.datetime64 | Timestamp

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/frame.pyi

@@ -1,3 +1,7 @@
+from builtins import (
+    bool as _bool,
+    str as _str,
+)
 from collections.abc import (
     Callable,
     Hashable,
@@ -160,9 +164,6 @@ from pandas._typing import (
 from pandas.io.formats.style import Styler
 from pandas.plotting import PlotAccessor
 
-_str = str
-_bool = bool
-
 class _iLocIndexerFrame(_iLocIndexer, Generic[_T]):
     @overload
     def __getitem__(self, idx: tuple[int, int]) -> Scalar: ...
@@ -2150,7 +2151,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     def reindex_like(
         self,
         other: DataFrame,
-        method: _str | FillnaOptions | Literal["nearest"] | None = ...,
+        method: FillnaOptions | Literal["nearest"] | None = ...,
         copy: _bool = ...,
         limit: int | None = ...,
         tolerance: Scalar | AnyArrayLike | Sequence[Scalar] = ...,

pandas-stubs/core/generic.pyi

@@ -1,3 +1,7 @@
+from builtins import (
+    bool as _bool,
+    str as _str,
+)
 from collections.abc import (
     Callable,
     Hashable,
@@ -60,17 +64,14 @@ from pandas._typing import (
 from pandas.io.pytables import HDFStore
 from pandas.io.sql import SQLTable
 
-_bool = bool
-_str = str
-
 class NDFrame(indexing.IndexingMixin):
     __hash__: ClassVar[None]  # type: ignore[assignment] # pyright: ignore[reportIncompatibleMethodOverride]
 
     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

@@ -1,3 +1,4 @@
+from builtins import str as _str
 from collections.abc import (
     Callable,
     Hashable,
@@ -66,8 +67,6 @@ from pandas._typing import (
 
 class InvalidIndexError(Exception): ...
 
-_str = str
-
 class Index(IndexOpsMixin[S1]):
     __hash__: ClassVar[None]  # type: ignore[assignment]
     # overloads with additional dtypes
@@ -272,10 +271,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 +477,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: ...