From c1cac48c724d850b0f9055d712f9cf900d4e9fe9 Mon Sep 17 00:00:00 2001 From: Loic Diridollou Date: Tue, 8 Oct 2024 08:45:05 -0400 Subject: [PATCH] GH935 Create overload for pd.Index.rename --- pandas-stubs/core/indexes/base.pyi | 5 ++++- tests/test_indexes.py | 9 +++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 297e94057..ea95aa82a 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -292,7 +292,10 @@ class Index(IndexOpsMixin[S1]): @names.setter def names(self, names: list[_str]): ... def set_names(self, names, *, level=..., inplace: bool = ...): ... - def rename(self, name, inplace: bool = ...) -> Self: ... + @overload + def rename(self, name, inplace: Literal[False] = False) -> Self: ... + @overload + def rename(self, name, inplace: Literal[True]) -> None: ... @property def nlevels(self) -> int: ... def sortlevel(self, level=..., ascending: bool = ..., sort_remaining=...): ... diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 46ec843af..75c9fc118 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -121,11 +121,20 @@ def test_str_match() -> None: def test_index_rename() -> None: + """Test that index rename returns an element of type Index.""" ind = pd.Index([1, 2, 3], name="foo") ind2 = ind.rename("goo") check(assert_type(ind2, "pd.Index[int]"), pd.Index, np.integer) +def test_index_rename_inplace() -> None: + """Test that index rename in-place does not return anything (None).""" + ind = pd.Index([1, 2, 3], name="foo") + ind2 = ind.rename("goo", inplace=True) + check(assert_type(ind2, None), type(None)) + assert ind2 is None + + def test_index_dropna(): idx = pd.Index([1, 2])