diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 4f9ebd44..0480ac05 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -15,6 +15,7 @@ from typing import ( ClassVar, Literal, TypeAlias, + TypeVar, final, overload, ) @@ -43,6 +44,7 @@ from typing_extensions import ( from pandas._libs.interval import _OrderableT from pandas._typing import ( S1, + S2, AnyAll, AxesData, DropKeep, @@ -64,6 +66,8 @@ from pandas._typing import ( type_t, ) +_T_INDEX = TypeVar("_T_INDEX", bound=Index) # ty: ignore[unresolved-reference] + class InvalidIndexError(Exception): ... class Index(IndexOpsMixin[S1]): @@ -401,7 +405,14 @@ class Index(IndexOpsMixin[S1]): ) -> Self: ... @overload def __getitem__(self, idx: int | tuple[np_ndarray_anyint, ...]) -> S1: ... - def append(self, other): ... + @overload + def append(self, other: Index[S1] | Sequence[Index[S1]]) -> Self: ... + @overload + def append(self, other: Index[S2] | Sequence[Index[S2]]) -> Index[S1 | S2]: ... + @overload + def append(self, other: Sequence[_T_INDEX]) -> Self | _T_INDEX: ... + @overload + def append(self, other: Index | Sequence) -> 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..fd852647 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,66 @@ 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 that gives pd.Index[Any]""" + first = pd.Index([1]) + second = pd.Index(["a"]) + third = pd.Index([1, "a"]) + check(assert_type(first.append(second), "pd.Index[int | str]"), pd.Index) + check(assert_type(first.append([second]), "pd.Index[int | str]"), pd.Index) + + check(assert_type(first.append(third), "pd.Index[int | str]"), pd.Index) # type: ignore[assert-type] + check(assert_type(first.append([third]), "pd.Index[int | str]"), pd.Index) # type: ignore[assert-type] + check( + assert_type( # type: ignore[assert-type] + first.append([second, third]), # pyright: ignore[reportAssertTypeFailure] + "pd.Index[int | str]", + ), + pd.Index, + ) + + check(assert_type(third.append([]), "pd.Index[int | str]"), pd.Index) # type: ignore[assert-type] + check( + assert_type(third.append(cast("list[Index[Any]]", [])), "pd.Index[int | str]"), # type: ignore[assert-type] + pd.Index, + ) + check(assert_type(third.append([first]), "pd.Index[int | str]"), pd.Index) # type: ignore[assert-type] + check( + assert_type( # type: ignore[assert-type] + third.append([first, second]), # pyright: ignore[reportAssertTypeFailure] + "pd.Index[int | str]", + ), + 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_append_list_str() -> None: + """Test pd.Index[list[str]].append""" + first = pd.Index([["str", "rts"]]) + second = pd.Index([["srt", "trs"]]) + check(assert_type(first.append([]), "pd.Index[list[str]]"), pd.Index, list) + check(assert_type(first.append(second), "pd.Index[list[str]]"), pd.Index, list) + check(assert_type(first.append([second]), "pd.Index[list[str]]"), pd.Index, list) + + def test_range_index_range() -> None: """Test that pd.RangeIndex can be initialized from range.""" iri = pd.RangeIndex(range(5))