Skip to content

Commit 6bf34de

Browse files
committed
Merge branch 'main' of github.com:pandas-dev/pandas-stubs into feature/floating-array
2 parents ee8615a + 50bed6e commit 6bf34de

File tree

5 files changed

+97
-44
lines changed

5 files changed

+97
-44
lines changed

pandas-stubs/core/frame.pyi

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -225,27 +225,6 @@ class _iLocIndexerFrame(_iLocIndexer, Generic[_T]):
225225
) -> None: ...
226226

227227
class _LocIndexerFrame(_LocIndexer, Generic[_T]):
228-
@overload
229-
def __getitem__(self, idx: Scalar) -> Series | _T: ...
230-
@overload
231-
def __getitem__( # type: ignore[overload-overlap]
232-
self,
233-
idx: (
234-
IndexType
235-
| MaskType
236-
| Callable[[DataFrame], IndexType | MaskType | Sequence[Hashable]]
237-
| list[HashableT]
238-
| tuple[
239-
IndexType
240-
| MaskType
241-
| list[HashableT]
242-
| slice
243-
| _IndexSliceTuple
244-
| Callable,
245-
MaskType | list[HashableT] | IndexType | Callable,
246-
]
247-
),
248-
) -> _T: ...
249228
@overload
250229
def __getitem__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
251230
self,
@@ -277,7 +256,28 @@ class _LocIndexerFrame(_LocIndexer, Generic[_T]):
277256
),
278257
) -> Series: ...
279258
@overload
259+
def __getitem__(self, idx: Scalar) -> Series | _T: ...
260+
@overload
280261
def __getitem__(self, idx: tuple[Scalar, slice]) -> Series | _T: ...
262+
@overload
263+
def __getitem__(
264+
self,
265+
idx: (
266+
IndexType
267+
| MaskType
268+
| Callable[[DataFrame], IndexType | MaskType | Sequence[Hashable]]
269+
| list[HashableT]
270+
| tuple[
271+
IndexType
272+
| MaskType
273+
| list[HashableT]
274+
| slice
275+
| _IndexSliceTuple
276+
| Callable,
277+
MaskType | Iterable[HashableT] | IndexType | Callable,
278+
]
279+
),
280+
) -> _T: ...
281281

282282
# Keep in sync with `DataFrame.__setitem__`
283283
@overload

pandas-stubs/core/indexes/base.pyi

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ from collections.abc import (
33
Callable,
44
Hashable,
55
Iterable,
6-
MutableMapping,
76
Sequence,
87
)
98
from datetime import (
@@ -84,6 +83,7 @@ from pandas._typing import (
8483
AnyArrayLikeInt,
8584
ArrayLike,
8685
AxesData,
86+
Axis,
8787
BuiltinFloatDtypeArg,
8888
CategoryDtypeArg,
8989
DropKeep,
@@ -111,6 +111,7 @@ from pandas._typing import (
111111
SequenceNotStr,
112112
SliceType,
113113
SupportsDType,
114+
TakeIndexer,
114115
TimedeltaDtypeArg,
115116
TimestampDtypeArg,
116117
np_1darray,
@@ -362,16 +363,15 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
362363
Index,
363364
]: ...
364365
@final
365-
def is_(self, other) -> bool: ...
366+
def is_(self, other: object) -> bool: ...
366367
def __len__(self) -> int: ...
367368
def __array__(
368369
self, dtype: _str | np.dtype = ..., copy: bool | None = ...
369370
) -> np_1darray: ...
370-
def __array_wrap__(self, result, context=...): ...
371371
@property
372372
def dtype(self) -> DtypeObj: ...
373373
@final
374-
def ravel(self, order: _str = ...): ...
374+
def ravel(self, order: _str = "C") -> Self: ...
375375
def view(self, cls=...): ...
376376
@overload
377377
def astype(
@@ -383,26 +383,23 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
383383
def astype(self, dtype: DtypeArg, copy: bool = True) -> Index: ...
384384
def take(
385385
self,
386-
indices,
387-
axis: int = 0,
386+
indices: TakeIndexer,
387+
axis: Axis = 0,
388388
allow_fill: bool = True,
389389
fill_value: Scalar | None = None,
390390
**kwargs: Any,
391-
): ...
392-
def repeat(self, repeats, axis=...): ...
391+
) -> Self: ...
392+
def repeat(
393+
self, repeats: int | AnyArrayLikeInt | Sequence[int], axis: None = None
394+
) -> Self: ...
393395
def copy(self, name: Hashable = ..., deep: bool = False) -> Self: ...
394-
@final
395-
def __copy__(self, **kwargs: Any): ...
396-
@final
397-
def __deepcopy__(self, memo: MutableMapping[int, Any] | None = None) -> Self: ...
398396
def format(
399397
self, name: bool = ..., formatter: Callable | None = ..., na_rep: _str = ...
400398
) -> list[_str]: ...
401-
def to_flat_index(self): ...
402399
def to_series(
403400
self, index: Index | None = None, name: Hashable | None = None
404401
) -> Series[S1]: ...
405-
def to_frame(self, index: bool = True, name=...) -> DataFrame: ...
402+
def to_frame(self, index: bool = True, name: Hashable = ...) -> DataFrame: ...
406403
@property
407404
def name(self) -> Hashable | None: ...
408405
@name.setter
@@ -411,11 +408,17 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
411408
def names(self) -> list[Hashable | None]: ...
412409
@names.setter
413410
def names(self, names: SequenceNotStr[Hashable | None]) -> None: ...
414-
def set_names(self, names, *, level=..., inplace: bool = ...): ...
411+
def set_names(
412+
self,
413+
names: Hashable | Sequence[Hashable],
414+
*,
415+
level: Level | Sequence[Level] | None = None,
416+
inplace: bool = False,
417+
) -> Self: ...
415418
@overload
416-
def rename(self, name, *, inplace: Literal[False] = False) -> Self: ...
419+
def rename(self, name: Hashable, *, inplace: Literal[False] = False) -> Self: ...
417420
@overload
418-
def rename(self, name, *, inplace: Literal[True]) -> None: ...
421+
def rename(self, name: Hashable, *, inplace: Literal[True]) -> None: ...
419422
@property
420423
def nlevels(self) -> int: ...
421424
def get_level_values(self, level: int | _str) -> Index: ...

pandas-stubs/core/indexes/multi.pyi

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,11 @@ from collections.abc import (
22
Callable,
33
Hashable,
44
Iterable,
5+
Mapping,
56
Sequence,
67
)
78
from typing import (
9+
Any,
810
final,
911
overload,
1012
)
@@ -109,13 +111,13 @@ class MultiIndex(Index):
109111
def dropna(self, how: AnyAll = "any") -> Self: ...
110112
def get_level_values(self, level: str | int) -> Index: ...
111113
def unique(self, level=...): ...
112-
def to_frame(
114+
def to_frame( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
113115
self,
114116
index: bool = True,
115117
name: list[HashableT] = ...,
116118
allow_duplicates: bool = False,
117119
) -> pd.DataFrame: ...
118-
def to_flat_index(self): ...
120+
def to_flat_index(self) -> Index: ...
119121
def remove_unused_levels(self): ...
120122
@property
121123
def nlevels(self) -> int: ...
@@ -163,3 +165,10 @@ class MultiIndex(Index):
163165
def insert(self, loc, item): ...
164166
def delete(self, loc): ...
165167
def isin(self, values, level=...) -> np_1darray_bool: ...
168+
def set_names(
169+
self,
170+
names: Hashable | Sequence[Hashable] | Mapping[Any, Hashable],
171+
*,
172+
level: Level | Sequence[Level] | None = None,
173+
inplace: bool = False,
174+
) -> Self: ...

tests/indexes/test_indexes.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1523,3 +1523,26 @@ def test_datetimeindex_where() -> None:
15231523

15241524
val_range = pd.RangeIndex(2).where(pd.Series([True, False]), 3)
15251525
check(assert_type(val_range, pd.Index), pd.RangeIndex)
1526+
1527+
1528+
def test_index_set_names() -> None:
1529+
"""Test Index.where with multiple types of other GH1419."""
1530+
idx = pd.Index([1, 2])
1531+
check(
1532+
assert_type(idx.set_names("chinchilla"), "pd.Index[int]"), pd.Index, np.integer
1533+
)
1534+
check(assert_type(idx.set_names((0,)), "pd.Index[int]"), pd.Index, np.integer)
1535+
check(
1536+
assert_type(idx.set_names(["chinchilla"]), "pd.Index[int]"),
1537+
pd.Index,
1538+
np.integer,
1539+
)
1540+
1541+
mi = pd.MultiIndex.from_arrays([[1, 2, 3], [4, 5, 6]], names=["elk", "owl"])
1542+
check(assert_type(mi.set_names(["beluga", "pig"]), pd.MultiIndex), pd.MultiIndex)
1543+
check(
1544+
assert_type(mi.set_names({"elk": "beluga", "owl": "pig"}), pd.MultiIndex),
1545+
pd.MultiIndex,
1546+
)
1547+
mi = cast("pd.MultiIndex", pd.Index([(1,)]))
1548+
check(assert_type(mi.set_names(1), pd.MultiIndex), pd.MultiIndex, tuple)

tests/test_frame.py

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22

33
from collections import (
44
OrderedDict,
5+
UserList,
56
defaultdict,
7+
deque,
68
)
79
from collections.abc import (
810
Callable,
@@ -3718,12 +3720,28 @@ def test_loc_int_set() -> None:
37183720
df.loc[np.uint64(1)] = [2, 3]
37193721

37203722

3721-
def test_loclist() -> None:
3722-
# GH 189
3723+
@pytest.mark.parametrize("col", [1, None])
3724+
@pytest.mark.parametrize("typ", [list, tuple, deque, UserList, iter])
3725+
def test_loc_iterable(col: Hashable, typ: type) -> None:
3726+
# GH 189, GH 1410
37233727
df = pd.DataFrame({1: [1, 2], None: 5}, columns=pd.Index([1, None], dtype=object))
3728+
check(df.loc[:, typ([col])], pd.DataFrame)
37243729

3725-
check(assert_type(df.loc[:, [None]], pd.DataFrame), pd.DataFrame)
3726-
check(assert_type(df.loc[:, [1]], pd.DataFrame), pd.DataFrame)
3730+
if TYPE_CHECKING:
3731+
assert_type(df.loc[:, [None]], pd.DataFrame)
3732+
assert_type(df.loc[:, [1]], pd.DataFrame)
3733+
3734+
assert_type(df.loc[:, (None,)], pd.DataFrame)
3735+
assert_type(df.loc[:, (1,)], pd.DataFrame)
3736+
3737+
assert_type(df.loc[:, deque([None])], pd.DataFrame)
3738+
assert_type(df.loc[:, deque([1])], pd.DataFrame)
3739+
3740+
assert_type(df.loc[:, UserList([None])], pd.DataFrame)
3741+
assert_type(df.loc[:, UserList([1])], pd.DataFrame)
3742+
3743+
assert_type(df.loc[:, (None for _ in [0])], pd.DataFrame)
3744+
assert_type(df.loc[:, (1 for _ in [0])], pd.DataFrame)
37273745

37283746

37293747
def test_dict_items() -> None:

0 commit comments

Comments
 (0)