From f78f807cb649341628e2a1c83754df7b03ee0216 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 19 Feb 2025 11:10:27 +0000 Subject: [PATCH 1/5] type multiindex constructors --- pandas-stubs/core/indexes/multi.pyi | 31 ++++++++++++++++++++--------- tests/test_indexes.py | 29 +++++++++++++++++++++++++-- 2 files changed, 49 insertions(+), 11 deletions(-) diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index 75ecd850b..125fb60e2 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -1,6 +1,7 @@ from collections.abc import ( Callable, Hashable, + Iterable, Sequence, ) from typing import ( @@ -16,6 +17,7 @@ from pandas.core.indexes.base import Index from typing_extensions import Self from pandas._typing import ( + AnyArrayLike, Dtype, DtypeArg, HashableT, @@ -27,31 +29,42 @@ from pandas._typing import ( class MultiIndex(Index[Any]): def __new__( cls, - levels=..., - codes=..., - sortorder=..., + levels: SequenceNotStr[SequenceNotStr[Hashable]] = ..., + codes: SequenceNotStr[SequenceNotStr[int]] = ..., + sortorder: int | None = ..., names: SequenceNotStr[Hashable] = ..., - dtype=..., - copy=..., + copy: bool = ..., name: SequenceNotStr[Hashable] = ..., verify_integrity: bool = ..., _set_identity: bool = ..., ) -> Self: ... @classmethod def from_arrays( - cls, arrays, sortorder=..., names: SequenceNotStr[Hashable] = ... + cls, + arrays: SequenceNotStr[SequenceNotStr[Hashable]] | SequenceNotStr[AnyArrayLike], + sortorder: int | None = ..., + names: SequenceNotStr[Hashable] = ..., ) -> Self: ... @classmethod def from_tuples( - cls, tuples, sortorder=..., names: SequenceNotStr[Hashable] = ... + cls, + tuples: Iterable[tuple[Hashable, ...]], + sortorder: int | None = ..., + names: SequenceNotStr[Hashable] = ..., ) -> Self: ... @classmethod def from_product( - cls, iterables, sortorder=..., names: SequenceNotStr[Hashable] = ... + cls, + iterables: SequenceNotStr[Iterable[Hashable]], + sortorder: int | None = ..., + names: SequenceNotStr[Hashable] = ..., ) -> Self: ... @classmethod def from_frame( - cls, df, sortorder=..., names: SequenceNotStr[Hashable] = ... + cls, + df: pd.DataFrame, + sortorder: int | None = ..., + names: SequenceNotStr[Hashable] = ..., ) -> Self: ... @property def shape(self): ... diff --git a/tests/test_indexes.py b/tests/test_indexes.py index 06d6d0d14..dedff599e 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -64,7 +64,7 @@ def test_multiindex_get_level_values() -> None: check(assert_type(i1, pd.Index), pd.Index) -def test_multiindex_constructor() -> None: +def test_multiindex_constructors() -> None: check( assert_type( pd.MultiIndex([[1], [4]], codes=[[0], [0]], name=["a", "b"]), pd.MultiIndex @@ -73,10 +73,35 @@ def test_multiindex_constructor() -> None: ) check( assert_type( - pd.MultiIndex([[1], [4]], codes=[[0], [0]], names=["a", "b"]), pd.MultiIndex + pd.MultiIndex( + [[1], [4]], + codes=[[0], [0]], + names=["a", "b"], + sortorder=0, + copy=True, + verify_integrity=True, + ), + pd.MultiIndex, + ), + pd.MultiIndex, + ) + check( + assert_type(pd.MultiIndex.from_arrays([[1], [4]]), pd.MultiIndex), pd.MultiIndex + ) + check( + assert_type( + pd.MultiIndex.from_arrays([np.arange(3), np.arange(3)]), pd.MultiIndex ), pd.MultiIndex, ) + check( + assert_type(pd.MultiIndex.from_tuples([(1, 3), (2, 4)]), pd.MultiIndex), + pd.MultiIndex, + ) + check( + assert_type(pd.MultiIndex.from_frame(pd.DataFrame({"a": [1]})), pd.MultiIndex), + pd.MultiIndex, + ) def test_index_tolist() -> None: From cabdc3a33d0372ea73c36a1c4fa7310cf0b10632 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 19 Feb 2025 11:27:38 +0000 Subject: [PATCH 2/5] simplify --- pandas-stubs/core/indexes/multi.pyi | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index 125fb60e2..493d126ab 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -29,7 +29,7 @@ from pandas._typing import ( class MultiIndex(Index[Any]): def __new__( cls, - levels: SequenceNotStr[SequenceNotStr[Hashable]] = ..., + levels: SequenceNotStr[SequenceNotStr[Hashable] | AnyArrayLike] = ..., codes: SequenceNotStr[SequenceNotStr[int]] = ..., sortorder: int | None = ..., names: SequenceNotStr[Hashable] = ..., @@ -41,7 +41,7 @@ class MultiIndex(Index[Any]): @classmethod def from_arrays( cls, - arrays: SequenceNotStr[SequenceNotStr[Hashable]] | SequenceNotStr[AnyArrayLike], + arrays: SequenceNotStr[SequenceNotStr[Hashable] | AnyArrayLike], sortorder: int | None = ..., names: SequenceNotStr[Hashable] = ..., ) -> Self: ... From b2d43123b02823f294e3a19927ee588834ae903e Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 19 Feb 2025 15:41:58 +0000 Subject: [PATCH 3/5] use ListLike --- pandas-stubs/core/indexes/multi.pyi | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index 493d126ab..69f1c6288 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -21,6 +21,7 @@ from pandas._typing import ( Dtype, DtypeArg, HashableT, + ListLike, MaskType, np_ndarray_anyint, np_ndarray_bool, @@ -29,8 +30,8 @@ from pandas._typing import ( class MultiIndex(Index[Any]): def __new__( cls, - levels: SequenceNotStr[SequenceNotStr[Hashable] | AnyArrayLike] = ..., - codes: SequenceNotStr[SequenceNotStr[int]] = ..., + levels: SequenceNotStr[ListLike] = ..., + codes: SequenceNotStr[ListLike] = ..., sortorder: int | None = ..., names: SequenceNotStr[Hashable] = ..., copy: bool = ..., From a7c82dedd86b98186c6c9fdd41ea7282ba6277b9 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 19 Feb 2025 15:42:52 +0000 Subject: [PATCH 4/5] use ListLike --- pandas-stubs/core/indexes/multi.pyi | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index 69f1c6288..8e7de09e3 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -17,7 +17,6 @@ from pandas.core.indexes.base import Index from typing_extensions import Self from pandas._typing import ( - AnyArrayLike, Dtype, DtypeArg, HashableT, @@ -42,7 +41,7 @@ class MultiIndex(Index[Any]): @classmethod def from_arrays( cls, - arrays: SequenceNotStr[SequenceNotStr[Hashable] | AnyArrayLike], + arrays: SequenceNotStr[ListLike], sortorder: int | None = ..., names: SequenceNotStr[Hashable] = ..., ) -> Self: ... From 6a88c1a9d05c6220d65eb69572d5a352ec6cc4ad Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 19 Feb 2025 16:06:45 +0000 Subject: [PATCH 5/5] remove outer sequencenotstr --- pandas-stubs/core/indexes/multi.pyi | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi index 8e7de09e3..906df9bbe 100644 --- a/pandas-stubs/core/indexes/multi.pyi +++ b/pandas-stubs/core/indexes/multi.pyi @@ -1,7 +1,6 @@ from collections.abc import ( Callable, Hashable, - Iterable, Sequence, ) from typing import ( @@ -17,10 +16,10 @@ from pandas.core.indexes.base import Index from typing_extensions import Self from pandas._typing import ( + Axes, Dtype, DtypeArg, HashableT, - ListLike, MaskType, np_ndarray_anyint, np_ndarray_bool, @@ -29,33 +28,32 @@ from pandas._typing import ( class MultiIndex(Index[Any]): def __new__( cls, - levels: SequenceNotStr[ListLike] = ..., - codes: SequenceNotStr[ListLike] = ..., + levels: Sequence[SequenceNotStr[Hashable]] = ..., + codes: Sequence[Sequence[int]] = ..., sortorder: int | None = ..., names: SequenceNotStr[Hashable] = ..., copy: bool = ..., name: SequenceNotStr[Hashable] = ..., verify_integrity: bool = ..., - _set_identity: bool = ..., ) -> Self: ... @classmethod def from_arrays( cls, - arrays: SequenceNotStr[ListLike], + arrays: Sequence[Axes], sortorder: int | None = ..., names: SequenceNotStr[Hashable] = ..., ) -> Self: ... @classmethod def from_tuples( cls, - tuples: Iterable[tuple[Hashable, ...]], + tuples: Sequence[tuple[Hashable, ...]], sortorder: int | None = ..., names: SequenceNotStr[Hashable] = ..., ) -> Self: ... @classmethod def from_product( cls, - iterables: SequenceNotStr[Iterable[Hashable]], + iterables: Sequence[SequenceNotStr[Hashable]], sortorder: int | None = ..., names: SequenceNotStr[Hashable] = ..., ) -> Self: ...