feat(array): #624 accumulate

Author: cmp0xff

pandas-stubs/core/arrays/base.pyi

@@ -1,6 +1,7 @@
 from collections.abc import Iterator
 from typing import (
     Any,
+    Literal,
     overload,
 )
 
@@ -68,7 +69,13 @@ class ExtensionArray:
     def _reduce(
         self, name: str, *, skipna: bool = ..., keepdims: bool = ..., **kwargs
     ) -> object: ...
-    def _accumulate(self, name: str, *, skipna: bool = ..., **kwargs) -> Self: ...
+    def _accumulate(
+        self,
+        name: Literal["cummin", "cummax", "cumsum", "cumprod"],
+        *,
+        skipna: bool = True,
+        **kwargs,
+    ) -> Self: ...
 
 class ExtensionOpsMixin:
     @classmethod

pandas-stubs/core/indexes/accessors.pyi

@@ -17,6 +17,10 @@ from pandas.core.accessor import PandasDelegate
 from pandas.core.arrays.base import ExtensionArray
 from pandas.core.arrays.categorical import Categorical
 from pandas.core.arrays.datetimes import DatetimeArray
+from pandas.core.arrays.integer import (
+    Int64Dtype,
+    IntegerArray,
+)
 from pandas.core.arrays.interval import IntervalArray
 from pandas.core.arrays.period import PeriodArray
 from pandas.core.arrays.timedeltas import TimedeltaArray
@@ -464,6 +468,10 @@ class ArrayDescriptor:
         self, instance: IndexOpsMixin[Never], owner: type[IndexOpsMixin]
     ) -> ExtensionArray: ...
     @overload
+    def __get__(
+        self, instance: IndexOpsMixin[Int64Dtype], owner: type[IndexOpsMixin]
+    ) -> IntegerArray: ...
+    @overload
     def __get__(
         self, instance: IndexOpsMixin[CategoricalDtype], owner: type[IndexOpsMixin]
     ) -> Categorical: ...

pyproject.toml

@@ -40,7 +40,7 @@ mypy = "1.18.2"
 pandas = "2.3.2"
 pyarrow = ">=10.0.1"
 pytest = ">=8.4.2"
-pyright = ">=1.1.405"
+pyright = ">=1.1.406"
 ty = ">=0.0.1a21"
 pyrefly = ">=0.35.0"
 poethepoet = ">=0.16.5"

tests/arrays/__init__.py

@@ -0,0 +1,39 @@
+from typing import cast
+
+from pandas.core.arrays.integer import (
+    Int64Dtype,
+    IntegerArray,
+)
+from pandas.core.series import Series
+from typing_extensions import assert_type
+
+from pandas._libs.missing import NA
+
+from tests import check
+
+
+def test_cumul_int64dtype():
+    # cast will be removed if pandas-dev/pandas-stubs#1395 is resolved
+    arr = cast("Series[Int64Dtype]", Series([1, NA, 2], dtype="Int64")).array
+
+    check(assert_type(arr._accumulate("cummin"), IntegerArray), IntegerArray)
+    check(assert_type(arr._accumulate("cummax"), IntegerArray), IntegerArray)
+    check(assert_type(arr._accumulate("cumsum"), IntegerArray), IntegerArray)
+    check(assert_type(arr._accumulate("cumprod"), IntegerArray), IntegerArray)
+
+    check(
+        assert_type(arr._accumulate("cummin", skipna=False), IntegerArray),
+        IntegerArray,
+    )
+    check(
+        assert_type(arr._accumulate("cummax", skipna=False), IntegerArray),
+        IntegerArray,
+    )
+    check(
+        assert_type(arr._accumulate("cumsum", skipna=False), IntegerArray),
+        IntegerArray,
+    )
+    check(
+        assert_type(arr._accumulate("cumprod", skipna=False), IntegerArray),
+        IntegerArray,
+    )

tests/arrays/test_arrays.py

@@ -1,21 +1,35 @@
 from typing import (
     TYPE_CHECKING,
+    cast,
 )
 
 import numpy as np
-import pandas as pd
 from pandas.core.arrays import DatetimeArray
 from pandas.core.arrays.base import ExtensionArray
+from pandas.core.arrays.categorical import Categorical
+from pandas.core.arrays.integer import (
+    Int64Dtype,
+    IntegerArray,
+)
 from pandas.core.arrays.interval import IntervalArray
 from pandas.core.arrays.timedeltas import TimedeltaArray
+from pandas.core.frame import DataFrame
 from pandas.core.indexes.accessors import (
     DatetimeProperties,
     PeriodProperties,
     Properties,
     TimedeltaProperties,
 )
+from pandas.core.indexes.interval import interval_range
+from pandas.core.indexes.period import period_range
+from pandas.core.series import Series
 from typing_extensions import assert_type
 
+from pandas._libs.interval import Interval
+from pandas._libs.missing import NA
+from pandas._libs.tslibs.timedeltas import Timedelta
+from pandas._libs.tslibs.timestamps import Timestamp
+
 from tests import (
     TYPE_CHECKING_INVALID_USAGE,
     check,
@@ -25,58 +39,61 @@
     from pandas.core.indexes.accessors import TimestampProperties  # noqa: F401
 
 
-def test_dt_property() -> None:
+def test_property_dt() -> None:
     """Test the Series.dt property"""
     check(
-        assert_type(pd.Series([pd.Timestamp(2025, 9, 28)]).dt, "TimestampProperties"),
+        assert_type(Series([Timestamp(2025, 9, 28)]).dt, "TimestampProperties"),
         DatetimeProperties,
     )
     check(
-        assert_type(pd.Series([pd.Timedelta(1, "s")]).dt, TimedeltaProperties),
+        assert_type(Series([Timedelta(1, "s")]).dt, TimedeltaProperties),
         TimedeltaProperties,
     )
     check(
         assert_type(
-            pd.period_range(start="2022-06-01", periods=10).to_series().dt,
+            period_range(start="2022-06-01", periods=10).to_series().dt,
             PeriodProperties,
         ),
         PeriodProperties,
     )
 
     if TYPE_CHECKING_INVALID_USAGE:
-        s = pd.DataFrame({"a": [1]})["a"]
+        s = DataFrame({"a": [1]})["a"]
         # python/mypy#19952: mypy believes Properties and its subclasses have a
         # conflict and gives Any for s.dt
         assert_type(s.dt, Properties)  # type: ignore[assert-type]
-        _1 = pd.Series([1]).dt  # type: ignore[arg-type] # pyright: ignore[reportAttributeAccessIssue]
+        _1 = Series([1]).dt  # type: ignore[arg-type] # pyright: ignore[reportAttributeAccessIssue]
 
 
-def test_array_property() -> None:
+def test_property_array() -> None:
     """Test that Series.array returns ExtensionArray and its subclasses"""
+    check(
+        assert_type(Series([1], dtype="category").array, Categorical), Categorical, int
+    )
+    # cast will be removed if pandas-dev/pandas-stubs#1395 is resolved
     check(
         assert_type(
-            pd.Series([1], dtype="category").array,
-            pd.Categorical,
+            cast("Series[Int64Dtype]", Series([1, NA], dtype="Int64")).array,
+            IntegerArray,
         ),
-        pd.Categorical,
-        int,
+        IntegerArray,
     )
     check(
-        assert_type(pd.Series(pd.interval_range(0, 1)).array, IntervalArray),
+        assert_type(Series(interval_range(0, 1)).array, IntervalArray),
         IntervalArray,
-        pd.Interval,
+        Interval,
     )
     check(
-        assert_type(pd.Series([pd.Timestamp(2025, 9, 28)]).array, DatetimeArray),
+        assert_type(Series([Timestamp(2025, 9, 28)]).array, DatetimeArray),
         DatetimeArray,
-        pd.Timestamp,
+        Timestamp,
     )
     check(
-        assert_type(pd.Series([pd.Timedelta(1, "s")]).array, TimedeltaArray),
+        assert_type(Series([Timedelta(1, "s")]).array, TimedeltaArray),
         TimedeltaArray,
-        pd.Timedelta,
+        Timedelta,
     )
-    check(assert_type(pd.Series([1]).array, ExtensionArray), ExtensionArray, np.integer)
+    check(assert_type(Series([1]).array, ExtensionArray), ExtensionArray, np.integer)
     # python/mypy#19952: mypy believes ExtensionArray and its subclasses have a
     # conflict and gives Any for s.array
-    check(assert_type(pd.Series([1, "s"]).array, ExtensionArray), ExtensionArray)  # type: ignore[assert-type]
+    check(assert_type(Series([1, "s"]).array, ExtensionArray), ExtensionArray)  # type: ignore[assert-type]