Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 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
9 changes: 8 additions & 1 deletion pandas-stubs/core/arrays/base.pyi
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from collections.abc import Iterator
from typing import (
Any,
Literal,
overload,
)

Expand Down Expand Up @@ -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
Expand Down
11 changes: 11 additions & 0 deletions pandas-stubs/core/construction.pyi
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
8 changes: 8 additions & 0 deletions pandas-stubs/core/indexes/accessors.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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: ...
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this should exist because Int64Dtype is in S1, and I don't want to introduce that now, because it will open up a whole set of new issues.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@overload
def __get__(
self, instance: IndexOpsMixin[CategoricalDtype], owner: type[IndexOpsMixin]
) -> Categorical: ...
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
Empty file added tests/arrays/__init__.py
Empty file.
12 changes: 12 additions & 0 deletions tests/arrays/test_arrays.py
Original file line number Diff line number Diff line change
@@ -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)
33 changes: 33 additions & 0 deletions tests/arrays/test_cumul.py
Original file line number Diff line number Diff line change
@@ -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,
)
57 changes: 37 additions & 20 deletions tests/series/test_properties.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should not have a test like this for 2 reasons:

  1. It encourages the use of Series[Int64Dtype] and I just don't think we would introduce that. Also, then Int64Dtype should really be in S1
  2. We can have a test with dtype="Int64", but then the results will just be an ExtensionArray for now, and if someone wants the real type, they can cast the result.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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]