From c24eeb068a845ca0908486ccbf35646c610c9fd8 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Mon, 2 Jan 2023 16:25:28 -0500 Subject: [PATCH 1/4] split Series[IntervalT] into separate __new__ --- pandas-stubs/core/series.pyi | 36 +++++++++++++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 79961bd3b..c56527bb6 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -70,6 +70,7 @@ from typing_extensions import ( ) import xarray as xr +from pandas._libs.interval import Interval from pandas._libs.missing import NAType from pandas._libs.tslibs import BaseOffset from pandas._typing import ( @@ -94,7 +95,6 @@ from pandas._typing import ( IgnoreRaise, IndexingInt, IntervalClosedType, - IntervalT, JoinHow, JsonSeriesOrient, Level, @@ -216,13 +216,43 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def __new__( cls, - data: IntervalIndex[IntervalT], + data: IntervalIndex[Interval[int]], index: Axes | None = ..., dtype=..., name: Hashable | None = ..., copy: bool = ..., fastpath: bool = ..., - ) -> Series[IntervalT]: ... + ) -> Series[Interval[int]]: ... + @overload + def __new__( + cls, + data: IntervalIndex[Interval[float]], + index: Axes | None = ..., + dtype=..., + name: Hashable | None = ..., + copy: bool = ..., + fastpath: bool = ..., + ) -> Series[Interval[float]]: ... + @overload + def __new__( + cls, + data: IntervalIndex[Interval[Timestamp]], + index: Axes | None = ..., + dtype=..., + name: Hashable | None = ..., + copy: bool = ..., + fastpath: bool = ..., + ) -> Series[Interval[Timestamp]]: ... + @overload + def __new__( + cls, + data: IntervalIndex[Interval[Timedelta]], + index: Axes | None = ..., + dtype=..., + name: Hashable | None = ..., + copy: bool = ..., + fastpath: bool = ..., + ) -> Series[Interval[Timedelta]]: ... @overload def __new__( cls, From f5f5d60a13f4459cff6d2e6efba29edff6f2a2a8 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Mon, 2 Jan 2023 17:17:24 -0500 Subject: [PATCH 2/4] allow multiindex for a colun in DataFrame.loc --- pandas-stubs/core/frame.pyi | 4 +++- pandas-stubs/core/indexing.pyi | 2 +- tests/test_frame.py | 13 +++++++++++++ 3 files changed, 17 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 4eadbf022..b5281d4f5 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -166,7 +166,9 @@ class _LocIndexerFrame(_LocIndexer): @overload def __getitem__( self, - idx: tuple[int | StrLike | tuple[ScalarT, ...], int | StrLike], + idx: tuple[ + int | StrLike | tuple[ScalarT, ...], int | StrLike | tuple[ScalarT, ...] + ], ) -> Scalar: ... @overload def __getitem__( diff --git a/pandas-stubs/core/indexing.pyi b/pandas-stubs/core/indexing.pyi index 35a748fcd..c9fef3aaa 100644 --- a/pandas-stubs/core/indexing.pyi +++ b/pandas-stubs/core/indexing.pyi @@ -15,7 +15,7 @@ from pandas._typing import ( ) _IndexSliceTuple: TypeAlias = tuple[ - Union[Index, MaskType, Scalar, list[ScalarT], slice], ... + Union[Index, MaskType, Scalar, list[ScalarT], slice | tuple[ScalarT, ...]], ... ] _IndexSliceUnion: TypeAlias = Union[slice, _IndexSliceTuple] diff --git a/tests/test_frame.py b/tests/test_frame.py index 02f45bc1f..72c7360a7 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2319,3 +2319,16 @@ def test_getattr_and_dataframe_groupby() -> None: assert_type(df.groupby("col1").col3.agg([min, max]), pd.DataFrame), pd.DataFrame, ) + + +def test_getsetitem_multiindex() -> None: + # GH 466 + rows = pd.Index(["project A", "project B", "project C"]) + years: tuple[str, ...] = ("Year 1", "Year 2", "Year 3") + quarters: tuple[str, ...] = ("Q1", "Q2", "Q3", "Q4") + index_tuples: list[tuple[str, ...]] = list(itertools.product(years, quarters)) + cols = pd.MultiIndex.from_tuples(index_tuples) + budget = pd.DataFrame(index=rows, columns=cols) + multi_index: tuple[str, str] = ("Year 1", "Q1") + budget.loc["project A", multi_index] = 4700 + check(assert_type(budget.loc["project A", multi_index], Scalar), int) From fb512037cf8ce9e5ca7ef956e78c62156d0966ba Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Tue, 3 Jan 2023 09:45:28 -0500 Subject: [PATCH 3/4] revert changes related to Series.__new__ and IntervalT --- pandas-stubs/core/series.pyi | 36 +++--------------------------------- 1 file changed, 3 insertions(+), 33 deletions(-) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index c56527bb6..79961bd3b 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -70,7 +70,6 @@ from typing_extensions import ( ) import xarray as xr -from pandas._libs.interval import Interval from pandas._libs.missing import NAType from pandas._libs.tslibs import BaseOffset from pandas._typing import ( @@ -95,6 +94,7 @@ from pandas._typing import ( IgnoreRaise, IndexingInt, IntervalClosedType, + IntervalT, JoinHow, JsonSeriesOrient, Level, @@ -216,43 +216,13 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def __new__( cls, - data: IntervalIndex[Interval[int]], + data: IntervalIndex[IntervalT], index: Axes | None = ..., dtype=..., name: Hashable | None = ..., copy: bool = ..., fastpath: bool = ..., - ) -> Series[Interval[int]]: ... - @overload - def __new__( - cls, - data: IntervalIndex[Interval[float]], - index: Axes | None = ..., - dtype=..., - name: Hashable | None = ..., - copy: bool = ..., - fastpath: bool = ..., - ) -> Series[Interval[float]]: ... - @overload - def __new__( - cls, - data: IntervalIndex[Interval[Timestamp]], - index: Axes | None = ..., - dtype=..., - name: Hashable | None = ..., - copy: bool = ..., - fastpath: bool = ..., - ) -> Series[Interval[Timestamp]]: ... - @overload - def __new__( - cls, - data: IntervalIndex[Interval[Timedelta]], - index: Axes | None = ..., - dtype=..., - name: Hashable | None = ..., - copy: bool = ..., - fastpath: bool = ..., - ) -> Series[Interval[Timedelta]]: ... + ) -> Series[IntervalT]: ... @overload def __new__( cls, From 3417725f9e174a36595017559c1f3e1c2419e807 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Tue, 3 Jan 2023 09:49:54 -0500 Subject: [PATCH 4/4] change tuple[ScalarT,...] to tuple[Scalar, ...] --- pandas-stubs/core/frame.pyi | 2 +- pandas-stubs/core/indexing.pyi | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index b5281d4f5..c6ad8493b 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -167,7 +167,7 @@ class _LocIndexerFrame(_LocIndexer): def __getitem__( self, idx: tuple[ - int | StrLike | tuple[ScalarT, ...], int | StrLike | tuple[ScalarT, ...] + int | StrLike | tuple[Scalar, ...], int | StrLike | tuple[Scalar, ...] ], ) -> Scalar: ... @overload diff --git a/pandas-stubs/core/indexing.pyi b/pandas-stubs/core/indexing.pyi index c9fef3aaa..a3b517150 100644 --- a/pandas-stubs/core/indexing.pyi +++ b/pandas-stubs/core/indexing.pyi @@ -15,7 +15,7 @@ from pandas._typing import ( ) _IndexSliceTuple: TypeAlias = tuple[ - Union[Index, MaskType, Scalar, list[ScalarT], slice | tuple[ScalarT, ...]], ... + Union[Index, MaskType, Scalar, list[ScalarT], slice | tuple[Scalar, ...]], ... ] _IndexSliceUnion: TypeAlias = Union[slice, _IndexSliceTuple]