Merge branch 'main' of github.com:loicdiridollou/pandas-stubs into gh203_groupby

Author: loicdiridollou

pandas-stubs/_libs/tslibs/dtypes.pyi

@@ -1,4 +1,5 @@
 from enum import Enum
+from typing import cast
 
 from .offsets import BaseOffset
 
@@ -29,16 +30,16 @@ class FreqGroup:
     def get_freq_group(code: int) -> int: ...
 
 class Resolution(Enum):
-    RESO_NS: int
-    RESO_US: int
-    RESO_MS: int
-    RESO_SEC: int
-    RESO_MIN: int
-    RESO_HR: int
-    RESO_DAY: int
-    RESO_MTH: int
-    RESO_QTR: int
-    RESO_YR: int
+    RESO_NS = cast(int, ...)
+    RESO_US = cast(int, ...)
+    RESO_MS = cast(int, ...)
+    RESO_SEC = cast(int, ...)
+    RESO_MIN = cast(int, ...)
+    RESO_HR = cast(int, ...)
+    RESO_DAY = cast(int, ...)
+    RESO_MTH = cast(int, ...)
+    RESO_QTR = cast(int, ...)
+    RESO_YR = cast(int, ...)
 
     def __lt__(self, other) -> bool: ...
     def __ge__(self, other) -> bool: ...

pandas-stubs/_libs/tslibs/offsets.pyi

@@ -14,6 +14,7 @@ from typing import (
 
 from dateutil.relativedelta import weekday as WeekdayClass
 import numpy as np
+from pandas import Timestamp
 from pandas.core.indexes.datetimes import DatetimeIndex
 from typing_extensions import Self
 
@@ -22,7 +23,7 @@ from pandas._typing import npt
 
 from pandas.tseries.holiday import AbstractHolidayCalendar
 
-_DatetimeT = TypeVar("_DatetimeT", bound=date)
+_DatetimeT = TypeVar("_DatetimeT", bound=datetime)
 _TimedeltaT = TypeVar("_TimedeltaT", bound=timedelta)
 
 prefix_mapping: dict[str, type]
@@ -42,26 +43,32 @@ class BaseOffset:
     @overload
     def __add__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
     @overload
-    def __add__(self, other: BaseOffset) -> Self: ...
+    def __add__(self, other: _DatetimeT) -> _DatetimeT: ...  # type: ignore[overload-overlap]  # pyright: ignore[reportOverlappingOverload]
+    @overload
+    def __add__(self, other: date) -> Timestamp: ...
     @overload
-    def __add__(self, other: _DatetimeT) -> _DatetimeT: ...
+    def __add__(self, other: BaseOffset) -> Self: ...
     @overload
     def __add__(self, other: _TimedeltaT) -> _TimedeltaT: ...
     @overload
     def __radd__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
     @overload
-    def __radd__(self, other: BaseOffset) -> Self: ...
+    def __radd__(self, other: _DatetimeT) -> _DatetimeT: ...  # type: ignore[overload-overlap]  # pyright: ignore[reportOverlappingOverload]
     @overload
-    def __radd__(self, other: _DatetimeT) -> _DatetimeT: ...
+    def __radd__(self, other: date) -> Timestamp: ...
+    @overload
+    def __radd__(self, other: BaseOffset) -> Self: ...
     @overload
     def __radd__(self, other: _TimedeltaT) -> _TimedeltaT: ...
     def __sub__(self, other: BaseOffset) -> Self: ...
     @overload
     def __rsub__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
     @overload
-    def __rsub__(self, other: BaseOffset) -> Self: ...
+    def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ...  # type: ignore[overload-overlap]  # pyright: ignore[reportOverlappingOverload]
     @overload
-    def __rsub__(self, other: _DatetimeT) -> _DatetimeT: ...
+    def __rsub__(self, other: date) -> Timestamp: ...
+    @overload
+    def __rsub__(self, other: BaseOffset) -> Self: ...
     @overload
     def __rsub__(self, other: _TimedeltaT) -> _TimedeltaT: ...
     def __call__(self, other): ...

pandas-stubs/_version.pyi

@@ -4,4 +4,4 @@ version_json: str = ...
 
 def get_versions(): ...
 
-_stub_version: Literal["2.2.2.240909"]
+_stub_version: Literal["2.2.3.241009"]

pandas-stubs/core/frame.pyi

@@ -605,7 +605,9 @@ class DataFrame(NDFrame, OpsMixin):
     @overload
     def __getitem__(self, key: Scalar | tuple[Hashable, ...]) -> Series: ...  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
     @overload
-    def __getitem__(self, key: Iterable[Hashable] | slice) -> DataFrame: ...
+    def __getitem__(  # pyright: ignore[reportOverlappingOverload]
+        self, key: Iterable[Hashable] | slice
+    ) -> DataFrame: ...
     @overload
     def __getitem__(self, key: Hashable) -> Series: ...
     def isetitem(

pandas-stubs/core/groupby/generic.pyi

@@ -238,7 +238,7 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
         self, key: Scalar | Hashable | tuple[Hashable, ...]
     ) -> SeriesGroupBy[Any, ByT]: ...
     @overload
-    def __getitem__(  # pyright: ignore[reportIncompatibleMethodOverride]
+    def __getitem__(  # pyright: ignore[reportIncompatibleMethodOverride, reportOverlappingOverload]
         self, key: Iterable[Hashable] | slice
     ) -> DataFrameGroupBy[ByT, bool]: ...
     def nunique(self, dropna: bool = ...) -> DataFrame: ...

pandas-stubs/core/groupby/groupby.pyi

@@ -141,7 +141,7 @@ class BaseGroupBy(SelectionMixin[NDFrameT], GroupByIndexingMixin):
     @overload
     def __getitem__(self: BaseGroupBy[DataFrame], key: Scalar | Hashable | tuple[Hashable, ...]) -> generic.SeriesGroupBy: ...  # type: ignore[overload-overlap]
     @overload
-    def __getitem__(
+    def __getitem__(  # pyright: ignore[reportOverlappingOverload]
         self: BaseGroupBy[DataFrame], key: Iterable[Hashable] | slice
     ) -> generic.DataFrameGroupBy: ...
     @overload

pandas-stubs/core/interchange/dataframe_protocol.pyi

@@ -11,33 +11,34 @@ import enum
 from typing import (
     Any,
     TypedDict,
+    cast,
 )
 
 class DlpackDeviceType(enum.IntEnum):
-    CPU: int
-    CUDA: int
-    CPU_PINNED: int
-    OPENCL: int
-    VULKAN: int
-    METAL: int
-    VPI: int
-    ROCM: int
+    CPU = cast(int, ...)
+    CUDA = cast(int, ...)
+    CPU_PINNED = cast(int, ...)
+    OPENCL = cast(int, ...)
+    VULKAN = cast(int, ...)
+    METAL = cast(int, ...)
+    VPI = cast(int, ...)
+    ROCM = cast(int, ...)
 
 class DtypeKind(enum.IntEnum):
-    INT: int
-    UINT: int
-    FLOAT: int
-    BOOL: int
-    STRING: int
-    DATETIME: int
-    CATEGORICAL: int
+    INT = cast(int, ...)
+    UINT = cast(int, ...)
+    FLOAT = cast(int, ...)
+    BOOL = cast(int, ...)
+    STRING = cast(int, ...)
+    DATETIME = cast(int, ...)
+    CATEGORICAL = cast(int, ...)
 
 class ColumnNullType(enum.IntEnum):
-    NON_NULLABLE: int
-    USE_NAN: int
-    USE_SENTINEL: int
-    USE_BITMASK: int
-    USE_BYTEMASK: int
+    NON_NULLABLE = cast(int, ...)
+    USE_NAN = cast(int, ...)
+    USE_SENTINEL = cast(int, ...)
+    USE_BITMASK = cast(int, ...)
+    USE_BYTEMASK = cast(int, ...)
 
 class ColumnBuffers(TypedDict):
     data: tuple[Buffer, Any]

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pandas-stubs"
-version = "2.2.2.240909"
+version = "2.2.3.241009"
 description = "Type annotations for pandas"
 authors = ["The Pandas Development Team <[email protected]>"]
 license = "BSD-3-Clause"
@@ -37,7 +37,7 @@ mypy = "1.10.1"
 pandas = "2.2.3"
 pyarrow = ">=10.0.1"
 pytest = ">=7.1.2"
-pyright = ">= 1.1.383"
+pyright = ">= 1.1.387"
 poethepoet = ">=0.16.5"
 loguru = ">=0.6.0"
 typing-extensions = ">=4.4.0"

scripts/__init__.py

@@ -1,9 +1,10 @@
 import sys
+from typing import Any
 
 from loguru import logger
 
 # Config the format of log message
-config = {
+config: dict[str, Any] = {
     "handlers": [
         {
             "sink": sys.stderr,

tests/test_timefuncs.py

@@ -664,8 +664,13 @@ def test_some_offsets() -> None:
         ),
         pd.DatetimeIndex,
     )
-    # GH 224
-    check(assert_type(dt.date.today() - Day(), dt.date), dt.date)
+    # GH 755
+    check(assert_type(dt.date.today() - Day(), pd.Timestamp), pd.Timestamp)
+    check(assert_type(dt.date.today() + Day(), pd.Timestamp), pd.Timestamp)
+    check(assert_type(Day() + dt.date.today(), pd.Timestamp), pd.Timestamp)
+    check(assert_type(dt.datetime.now() - Day(), dt.datetime), dt.datetime)
+    check(assert_type(dt.datetime.now() + Day(), dt.datetime), dt.datetime)
+    check(assert_type(Day() + dt.datetime.now(), dt.datetime), dt.datetime)
     # GH 235
     check(
         assert_type(