diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 406d1ad48..2a775b4bd 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -49,6 +49,7 @@ from pandas._typing import ( Label, Level, NaPosition, + SequenceNotStr, TimedeltaDtypeArg, TimestampDtypeArg, np_ndarray_anyint, @@ -288,9 +289,9 @@ class Index(IndexOpsMixin[S1]): @name.setter def name(self, value) -> None: ... @property - def names(self) -> list[_str]: ... + def names(self) -> list[_str | None]: ... @names.setter - def names(self, names: list[_str]): ... + def names(self, names: SequenceNotStr[_str | None]): ... def set_names(self, names, *, level=..., inplace: bool = ...): ... @overload def rename(self, name, inplace: Literal[False] = False) -> Self: ... diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 75c9fc118..b43c21333 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -81,11 +81,17 @@ def test_column_contains() -> None: # https://github.com/microsoft/python-type-stubs/issues/199 df = pd.DataFrame({"A": [1, 2], "B": ["c", "d"], "E": [3, 4]}) - collist = [column for column in df.columns] - - collist2 = [column for column in df.columns[df.columns.str.contains("A|B")]] - - length = len(df.columns[df.columns.str.contains("A|B")]) + # check accessing the columns as Scalar + check(assert_type([column for column in df.columns], list[str]), list) + # check slicing the columns with a Series[bool] + check( + assert_type( + [column for column in df.columns[df.columns.str.contains("A|B")]], list[str] + ), + list, + ) + # check that generic methods are defined on a slice of an index + check(assert_type(len(df.columns[df.columns.str.contains("A|B")]), int), int) def test_column_sequence() -> None: @@ -1160,3 +1166,22 @@ def test_value_counts() -> None: pd.Series, float, ) + + +def test_index_naming() -> None: + """ + Test index names type both for the getter and the setter. + + The names of an index should be settable with a sequence (not str) and names + property is a list[str | None] (FrozenList). + """ + df = pd.DataFrame({"a": ["a", "b", "c"], "i": [10, 11, 12]}) + + df.index.names = ["idx"] + check(assert_type(df.index.names, list[str | None]), list) + df.index.names = ("idx2",) + check(assert_type(df.index.names, list[str | None]), list) + df.index.names = [None] + check(assert_type(df.index.names, list[str | None]), list) + df.index.names = (None,) + check(assert_type(df.index.names, list[str | None]), list)