From 86c94ebdca7e9e21d11e13bb75cc5d95380dc88a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Diridollou?= Date: Thu, 9 Oct 2025 19:35:34 -0400 Subject: [PATCH 1/4] GH1419 Allow Series and Index for other in Index.where(..., other) --- pandas-stubs/core/indexes/base.pyi | 4 ++-- tests/indexes/test_indexes.py | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index b41353ef4..1c4a13caf 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -61,7 +61,7 @@ from pandas._typing import ( T_COMPLEX, T_INT, AnyAll, - ArrayLike, + AnyArrayLike, AxesData, CategoryDtypeArg, DropKeep, @@ -434,7 +434,7 @@ class Index(IndexOpsMixin[S1]): @property def values(self) -> np_1darray: ... def memory_usage(self, deep: bool = False): ... - def where(self, cond, other: Scalar | ArrayLike | None = None): ... + def where(self, cond, other: Scalar | AnyArrayLike | None = None) -> Self: ... def __contains__(self, key) -> bool: ... @final def __setitem__(self, key, value) -> None: ... diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 892cea004..fee3f2928 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -19,6 +19,7 @@ from pandas.core.arrays.timedeltas import TimedeltaArray from pandas.core.indexes.base import Index from pandas.core.indexes.category import CategoricalIndex +from pandas.core.indexes.datetimes import DatetimeIndex from typing_extensions import ( Never, assert_type, @@ -1608,3 +1609,18 @@ def test_to_series() -> None: np.complexfloating, ) check(assert_type(Index(["1"]).to_series(), "pd.Series[str]"), pd.Series, str) + + +def test_index_where() -> None: + """Test Index.where with multiple types of other GH1419.""" + datetime_index = pd.DatetimeIndex( + pd.date_range(start="2025-01-01", freq="h", periods=48) + ) + mask = np.ones(48, dtype=bool) + val_idx = datetime_index.where(mask, datetime_index - pd.Timedelta(days=1)) + check(assert_type(val_idx, DatetimeIndex), DatetimeIndex) + + val_sr = datetime_index.where( + mask, (datetime_index - pd.Timedelta(days=1)).to_series() + ) + check(assert_type(val_sr, DatetimeIndex), DatetimeIndex) From a6653dc18240b90a1578959df75659451df4f8df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Diridollou?= Date: Fri, 10 Oct 2025 15:55:01 -0400 Subject: [PATCH 2/4] GH1419 PR Feedback --- tests/indexes/test_indexes.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index fee3f2928..d08ff17e2 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1613,6 +1613,17 @@ def test_to_series() -> None: def test_index_where() -> None: """Test Index.where with multiple types of other GH1419.""" + idx = pd.Index(range(48)) + mask = np.ones(48, dtype=bool) + val_idx = idx.where(mask, idx) + check(assert_type(val_idx, "pd.Index[int]"), pd.Index, np.integer) + + val_sr = idx.where(mask, (idx).to_series()) + check(assert_type(val_sr, "pd.Index[int]"), pd.Index, np.integer) + + +def test_datetimeindex_where() -> None: + """Test DatetimeIndex.where with multiple types of other GH1419.""" datetime_index = pd.DatetimeIndex( pd.date_range(start="2025-01-01", freq="h", periods=48) ) From 06591660b3cf85c586a14635d903f8bdb0633a2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Diridollou?= Date: Fri, 10 Oct 2025 15:57:10 -0400 Subject: [PATCH 3/4] GH1419 PR Feedback --- tests/indexes/test_indexes.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index d08ff17e2..9e12ee4ca 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1624,9 +1624,7 @@ def test_index_where() -> None: def test_datetimeindex_where() -> None: """Test DatetimeIndex.where with multiple types of other GH1419.""" - datetime_index = pd.DatetimeIndex( - pd.date_range(start="2025-01-01", freq="h", periods=48) - ) + datetime_index = pd.date_range(start="2025-01-01", freq="h", periods=48) mask = np.ones(48, dtype=bool) val_idx = datetime_index.where(mask, datetime_index - pd.Timedelta(days=1)) check(assert_type(val_idx, DatetimeIndex), DatetimeIndex) From 96c8a145286ac429b2f657f47adbbc4cd1c0cd57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Diridollou?= Date: Fri, 10 Oct 2025 16:00:31 -0400 Subject: [PATCH 4/4] GH1419 PR Feedback --- tests/indexes/test_indexes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/indexes/test_indexes.py b/tests/indexes/test_indexes.py index 9e12ee4ca..d079c9a8e 100644 --- a/tests/indexes/test_indexes.py +++ b/tests/indexes/test_indexes.py @@ -1616,10 +1616,10 @@ def test_index_where() -> None: idx = pd.Index(range(48)) mask = np.ones(48, dtype=bool) val_idx = idx.where(mask, idx) - check(assert_type(val_idx, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(val_idx, "pd.Index[int]"), pd.Index, int) val_sr = idx.where(mask, (idx).to_series()) - check(assert_type(val_sr, "pd.Index[int]"), pd.Index, np.integer) + check(assert_type(val_sr, "pd.Index[int]"), pd.Index, int) def test_datetimeindex_where() -> None: