Skip to content

allow multiindex for a column in DataFrame.loc #494

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jan 3, 2023
Merged

allow multiindex for a column in DataFrame.loc #494

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c24eeb0
split Series[IntervalT] into separate __new__
Dr-Irv Jan 2, 2023
f5f5d60
allow multiindex for a colun in DataFrame.loc
Dr-Irv Jan 2, 2023
fb51203
revert changes related to Series.__new__ and IntervalT
Dr-Irv Jan 3, 2023
3417725
change tuple[ScalarT,...] to tuple[Scalar, ...]
Dr-Irv Jan 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion pandas-stubs/core/frame.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, ...]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would Scalar also work here instead of ScalarT?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would Scalar also work here instead of ScalarT?

Yes, seems so. Fixed in commit 3417725

],
) -> Scalar: ...
@overload
def __getitem__(
Expand Down
2 changes: 1 addition & 1 deletion pandas-stubs/core/indexing.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, ...]], ...
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here (I think we need it for list but not sure about tuple).

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here (I think we need it for list but not sure about tuple).

Yes, fixed in commit 3417725

]

_IndexSliceUnion: TypeAlias = Union[slice, _IndexSliceTuple]
Expand Down
36 changes: 33 additions & 3 deletions pandas-stubs/core/series.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -94,7 +95,6 @@ from pandas._typing import (
IgnoreRaise,
IndexingInt,
IntervalClosedType,
IntervalT,
JoinHow,
JsonSeriesOrient,
Level,
Expand Down Expand Up @@ -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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would be nice to add the TODO again

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same issue - included the changes from the other PR, so I've reverted that here in commit fb51203

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,
Expand Down
13 changes: 13 additions & 0 deletions tests/test_frame.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)