Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
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
2 changes: 1 addition & 1 deletion pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -224,7 +224,7 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]):
| slice
| _IndexSliceTuple
| Callable,
MaskType | list[HashableT] | slice | Callable,
MaskType | list[HashableT] | IndexType | Callable,
]
),
) -> _T: ...
Expand Down
16 changes: 15 additions & 1 deletion tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -281,6 +281,7 @@ def test_types_loc_at() -> None:
df.loc[0, "col1"]



def test_types_boolean_indexing() -> None:
df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
check(assert_type(df[df > 1], pd.DataFrame), pd.DataFrame)
Expand Down Expand Up @@ -3737,13 +3738,26 @@ def test_xs_key() -> None:


def test_loc_slice() -> None:
# GH 277
"""Test DataFrame.loc with a slice, Index, Series."""
# GH277
df1 = pd.DataFrame(
{"x": [1, 2, 3, 4]},
index=pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=["num", "let"]),
)
check(assert_type(df1.loc[1, :], Union[pd.Series, pd.DataFrame]), pd.DataFrame)

# GH1299
ind = pd.Index(["a", "b"])
mask = pd.Series([True, False])
mask_col = pd.Series([True, False], index=pd.Index(["a", "b"]))
df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})

# loc with index for columns
check(assert_type(df.loc[mask, ind], pd.DataFrame), pd.DataFrame)
# loc with index for columns
check(assert_type(df.loc[mask, mask_col], pd.DataFrame), pd.DataFrame)



def test_where() -> None:
df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
Expand Down
Loading