feat(array): #624 accumulate (#1400)

* feat(array): #624 accumulate

* use pd.array directly

* fix(comment): remove Series[Int64Dtype]

Author: cmp0xff

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
     hooks:
     -   id: isort
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.13.2
+    rev: v0.13.3
     hooks:
     -   id: ruff-check
         args: [

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

@@ -1,10 +1,21 @@
 from collections.abc import Sequence
+from typing import overload
 
 import numpy as np
 from pandas.core.arrays.base import ExtensionArray
+from pandas.core.arrays.integer import IntegerArray
+
+from pandas._libs.missing import NAType
 
 from pandas.core.dtypes.dtypes import ExtensionDtype
 
+@overload
+def array(
+    data: Sequence[int] | Sequence[int | NAType],
+    dtype: str | np.dtype | ExtensionDtype | None = None,
+    copy: bool = True,
+) -> IntegerArray: ...
+@overload
 def array(
     data: Sequence[object],
     dtype: str | np.dtype | ExtensionDtype | None = None,

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,12 @@
+from pandas.core.arrays.integer import IntegerArray
+from pandas.core.construction import array
+from typing_extensions import assert_type
+
+from pandas._libs.missing import NA
+
+from tests import check
+
+
+def test_construction() -> None:
+    check(assert_type(array([1]), IntegerArray), IntegerArray)
+    check(assert_type(array([1, NA]), IntegerArray), IntegerArray)

tests/arrays/test_arrays.py

@@ -0,0 +1,33 @@
+from pandas.core.arrays.integer import IntegerArray
+from pandas.core.construction import array
+from typing_extensions import assert_type
+
+from pandas._libs.missing import NA
+
+from tests import check
+
+
+def test_cumul_int64dtype() -> None:
+    arr = array([1, NA, 2])
+
+    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_cumul.py

@@ -3,19 +3,27 @@
 )
 
 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.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.tslibs.timedeltas import Timedelta
+from pandas._libs.tslibs.timestamps import Timestamp
+
 from tests import (
     TYPE_CHECKING_INVALID_USAGE,
     check,
@@ -25,58 +33,53 @@
     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(
-            pd.Series([1], dtype="category").array,
-            pd.Categorical,
-        ),
-        pd.Categorical,
-        int,
+        assert_type(Series([1], dtype="category").array, Categorical), Categorical, int
     )
     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]