diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 7b2b0081..f14c393f 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -847,6 +847,26 @@ S1 = TypeVar("S1", bound=SeriesDType, default=Any) S2 = TypeVar("S2", bound=SeriesDType) S3 = TypeVar("S3", bound=SeriesDType) +# Constraint, instead of bound +C2 = TypeVar( + "C2", + str, + bytes, + datetime.date, + datetime.time, + bool, + int, + float, + complex, + Dtype, + datetime.datetime, # includes pd.Timestamp + datetime.timedelta, # includes pd.Timedelta + Period, + Interval, + CategoricalDtype, + BaseOffset, +) + IndexingInt: TypeAlias = ( int | np.int_ | np.integer | np.unsignedinteger | np.signedinteger | np.int8 ) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 4f9ebd44..cea53a64 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -42,6 +42,7 @@ from typing_extensions import ( from pandas._libs.interval import _OrderableT from pandas._typing import ( + C2, S1, AnyAll, AxesData, @@ -401,7 +402,12 @@ class Index(IndexOpsMixin[S1]): ) -> Self: ... @overload def __getitem__(self, idx: int | tuple[np_ndarray_anyint, ...]) -> S1: ... - def append(self, other): ... + @overload + def append( + self: Index[C2], other: Index[C2] | Sequence[Index[C2]] + ) -> Index[C2]: ... + @overload + def append(self, other: Index | Sequence[Index]) -> Index: ... def putmask(self, mask, value): ... def equals(self, other) -> bool: ... @final diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index d313852d..5278d081 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -135,7 +135,7 @@ class MultiIndex(Index): def take( self, indices, axis: int = ..., allow_fill: bool = ..., fill_value=..., **kwargs ): ... - def append(self, other): ... + def append(self, other): ... # pyrefly: ignore def argsort(self, *args, **kwargs): ... def repeat(self, repeats, axis=...): ... @final diff --git a/pyproject.toml b/pyproject.toml index 71400163..d916bfa9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ mypy = "1.17.0" pandas = "2.3.0" pyarrow = ">=10.0.1" pytest = ">=7.1.2" -pyright = ">=1.1.400" +pyright = ">=1.1.403" ty = "^0.0.1a8" pyrefly = "^0.21.0" poethepoet = ">=0.16.5" diff --git a/tests/test_indexes.py b/tests/test_indexes.py index b3f566cd..57042c29 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -3,7 +3,9 @@ import datetime as dt from typing import ( TYPE_CHECKING, + Any, Union, + cast, ) import numpy as np @@ -1028,6 +1030,42 @@ def test_getitem() -> None: check(assert_type(i0[[0, 2]], "pd.Index[str]"), pd.Index, str) +def test_append_mix() -> None: + """Test pd.Index.append with mixed types""" + first = pd.Index([1]) + second = pd.Index(["a"]) + third = pd.Index([1, "a"]) + check(assert_type(first.append(second), pd.Index), pd.Index) + check(assert_type(first.append([second]), pd.Index), pd.Index) + + check(assert_type(first.append(third), pd.Index), pd.Index) + check(assert_type(first.append([third]), pd.Index), pd.Index) + check(assert_type(first.append([second, third]), pd.Index), pd.Index) + + check(assert_type(third.append([]), pd.Index), pd.Index) + check(assert_type(third.append(cast("list[Index[Any]]", [])), pd.Index), pd.Index) + check(assert_type(third.append([first]), pd.Index), pd.Index) + check(assert_type(third.append([first, second]), pd.Index), pd.Index) + + +def test_append_int() -> None: + """Test pd.Index[int].append""" + first = pd.Index([1]) + second = pd.Index([2]) + check(assert_type(first.append([]), "pd.Index[int]"), pd.Index, np.int64) + check(assert_type(first.append(second), "pd.Index[int]"), pd.Index, np.int64) + check(assert_type(first.append([second]), "pd.Index[int]"), pd.Index, np.int64) + + +def test_append_str() -> None: + """Test pd.Index[str].append""" + first = pd.Index(["str"]) + second = pd.Index(["rts"]) + check(assert_type(first.append([]), "pd.Index[str]"), pd.Index, str) + check(assert_type(first.append(second), "pd.Index[str]"), pd.Index, str) + check(assert_type(first.append([second]), "pd.Index[str]"), pd.Index, str) + + def test_range_index_range() -> None: """Test that pd.RangeIndex can be initialized from range.""" iri = pd.RangeIndex(range(5))