From 21a7794ae86767bf0051f9aea6dcea8c124c06d1 Mon Sep 17 00:00:00 2001 From: Loic Diridollou Date: Tue, 12 Nov 2024 21:45:55 -0500 Subject: [PATCH 1/2] GH804 Index names typing --- pandas-stubs/core/indexes/base.pyi | 5 +++-- tests/test_indexes.py | 11 +++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 406d1ad48..53df00da3 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) -> SequenceNotStr: ... @names.setter - def names(self, names: list[_str]): ... + def names(self, names: SequenceNotStr): ... 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..12f9cafa9 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -1160,3 +1160,14 @@ def test_value_counts() -> None: pd.Series, float, ) + + +def test_index_naming() -> None: + """Test that we can set the names of an index as a hashable.""" + df = pd.DataFrame({"a": ["a", "b", "c"], "i": [10, 11, 12]}) + + df.index.names = ["idx"] + + df.index.names = ("idx2",) + df.index.names = [None] + df.index.names = (None,) From bd1dfbd687a697edd95c95b2db3121c3bddfc5cb Mon Sep 17 00:00:00 2001 From: Loic Diridollou Date: Wed, 13 Nov 2024 19:03:25 -0500 Subject: [PATCH 2/2] GH804 Potential Index.names typehint --- pandas-stubs/core/indexes/base.pyi | 4 ++-- tests/test_indexes.py | 28 +++++++++++++++++++++------- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/pandas-stubs/core/indexes/base.pyi b/pandas-stubs/core/indexes/base.pyi index 53df00da3..2a775b4bd 100644 --- a/pandas-stubs/core/indexes/base.pyi +++ b/pandas-stubs/core/indexes/base.pyi @@ -289,9 +289,9 @@ class Index(IndexOpsMixin[S1]): @name.setter def name(self, value) -> None: ... @property - def names(self) -> SequenceNotStr: ... + def names(self) -> list[_str | None]: ... @names.setter - def names(self, names: SequenceNotStr): ... + 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 12f9cafa9..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: @@ -1163,11 +1169,19 @@ def test_value_counts() -> None: def test_index_naming() -> None: - """Test that we can set the names of an index as a hashable.""" + """ + 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)