Skip to content

Commit f2ee304

Browse files
authored
gh:576 - fixed .loc (#581)
* fixed .loc * updated test name and func argument
1 parent 44756ab commit f2ee304

File tree

2 files changed

+15
-2
lines changed

2 files changed

+15
-2
lines changed

pandas-stubs/core/frame.pyi

+8-1
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,12 @@ class _LocIndexerFrame(_LocIndexer):
161161
| Callable[[DataFrame], IndexType | MaskType | list[HashableT]]
162162
| list[HashableT]
163163
| tuple[
164-
IndexType | MaskType | list[HashableT] | Hashable,
164+
IndexType
165+
| MaskType
166+
| list[HashableT]
167+
| slice
168+
| _IndexSliceTuple
169+
| Callable,
165170
list[HashableT] | slice | Series[bool] | Callable,
166171
],
167172
) -> DataFrame: ...
@@ -188,6 +193,8 @@ class _LocIndexerFrame(_LocIndexer):
188193
| None,
189194
) -> Series: ...
190195
@overload
196+
def __getitem__(self, idx: tuple[Scalar, slice]) -> Series | DataFrame: ...
197+
@overload
191198
def __setitem__(
192199
self,
193200
idx: MaskType | StrLike | _IndexSliceTuple | list[ScalarT],

tests/test_frame.py

+7-1
Original file line numberDiff line numberDiff line change
@@ -2172,7 +2172,7 @@ def test_loc_slice() -> None:
21722172
{"x": [1, 2, 3, 4]},
21732173
index=pd.MultiIndex.from_product([[1, 2], ["a", "b"]], names=["num", "let"]),
21742174
)
2175-
check(assert_type(df1.loc[1, :], pd.DataFrame), pd.DataFrame)
2175+
check(assert_type(df1.loc[1, :], Union[pd.Series, pd.DataFrame]), pd.DataFrame)
21762176

21772177

21782178
def test_where() -> None:
@@ -2521,3 +2521,9 @@ def test_align() -> None:
25212521
aligned_df0, aligned_df1 = df0.align(df1)
25222522
check(assert_type(aligned_df0, pd.DataFrame), pd.DataFrame)
25232523
check(assert_type(aligned_df1, pd.DataFrame), pd.DataFrame)
2524+
2525+
2526+
def test_loc_returns_series() -> None:
2527+
df1 = pd.DataFrame({"x": [1, 2, 3, 4]}, index=[10, 20, 30, 40])
2528+
df2 = df1.loc[10, :]
2529+
check(assert_type(df2, Union[pd.Series, pd.DataFrame]), pd.Series)

0 commit comments

Comments
 (0)