Skip to content

Commit 8429b2d

Browse files
Dr-IrvDr-Irv
authored
fix issues 302, 314, 320, 321 (#327)
1 parent 10b87cd commit 8429b2d

File tree

9 files changed

+65
-62
lines changed

9 files changed

+65
-62
lines changed

‎pandas-stubs/_libs/tslibs/offsets.pyi

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,8 @@ class BusinessHour(BusinessMixin):
136136
offset: timedelta = ...,
137137
): ...
138138

139-
class WeekOfMonthMixin(SingleConstructorOffset): ...
139+
class WeekOfMonthMixin(SingleConstructorOffset):
140+
def __init__(self, n: int = ..., weekday: Literal[0, 1, 2, 3, 4, 5, 6] = ...): ...
140141

141142
class YearOffset(SingleConstructorOffset):
142143
def __init__(

‎pandas-stubs/core/frame.pyi

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ from pandas.core.groupby.grouper import Grouper
3030
from pandas.core.indexes.base import Index
3131
from pandas.core.indexing import (
3232
_iLocIndexer,
33+
_IndexSliceTuple,
3334
_LocIndexer,
3435
)
3536
from pandas.core.resample import Resampler
@@ -150,28 +151,25 @@ class _LocIndexerFrame(_LocIndexer):
150151
@overload
151152
def __getitem__(
152153
self,
153-
idx: tuple[int | StrLike | tuple[HashableT, ...], int | StrLike],
154+
idx: tuple[int | StrLike | tuple[ScalarT, ...], int | StrLike],
154155
) -> Scalar: ...
155156
@overload
156157
def __getitem__(
157158
self,
158159
idx: ScalarT
159-
| tuple[IndexType | MaskType | tuple[slice, ...], ScalarT | None]
160+
| tuple[IndexType | MaskType | _IndexSliceTuple, ScalarT | None]
160161
| None,
161162
) -> Series: ...
162163
@overload
163164
def __setitem__(
164165
self,
165-
idx: MaskType
166-
| StrLike
167-
| tuple[MaskType | Index | Sequence[ScalarT] | Scalar | slice, ...]
168-
| list[ScalarT],
166+
idx: MaskType | StrLike | _IndexSliceTuple | list[ScalarT],
169167
value: S1 | ArrayLike | Series | DataFrame,
170168
) -> None: ...
171169
@overload
172170
def __setitem__(
173171
self,
174-
idx: tuple[tuple[HashableT | Scalar | slice, ...], HashableT],
172+
idx: tuple[_IndexSliceTuple, HashableT],
175173
value: S1 | ArrayLike | Series[S1] | list,
176174
) -> None: ...
177175

@@ -550,10 +548,10 @@ class DataFrame(NDFrame, OpsMixin):
550548
@overload
551549
def drop(
552550
self,
553-
labels: Hashable | Sequence[Hashable] = ...,
551+
labels: Hashable | Sequence[Hashable] | Index = ...,
554552
*,
555553
axis: Axis = ...,
556-
index: Hashable | Sequence[Hashable] = ...,
554+
index: Hashable | Sequence[Hashable] | Index = ...,
557555
columns: Hashable | Sequence[Hashable] | Index = ...,
558556
level: Level | None = ...,
559557
inplace: Literal[True],
@@ -562,10 +560,10 @@ class DataFrame(NDFrame, OpsMixin):
562560
@overload
563561
def drop(
564562
self,
565-
labels: Hashable | Sequence[Hashable] = ...,
563+
labels: Hashable | Sequence[Hashable] | Index = ...,
566564
*,
567565
axis: Axis = ...,
568-
index: Hashable | Sequence[Hashable] = ...,
566+
index: Hashable | Sequence[Hashable] | Index = ...,
569567
columns: Hashable | Sequence[Hashable] | Index = ...,
570568
level: Level | None = ...,
571569
inplace: Literal[False] = ...,
@@ -574,10 +572,10 @@ class DataFrame(NDFrame, OpsMixin):
574572
@overload
575573
def drop(
576574
self,
577-
labels: Hashable | Sequence[Hashable] = ...,
575+
labels: Hashable | Sequence[Hashable] | Index = ...,
578576
*,
579577
axis: Axis = ...,
580-
index: Hashable | Sequence[Hashable] = ...,
578+
index: Hashable | Sequence[Hashable] | Index = ...,
581579
columns: Hashable | Sequence[Hashable] | Index = ...,
582580
level: Level | None = ...,
583581
inplace: bool = ...,

‎pandas-stubs/core/indexes/accessors.pyi

Lines changed: 2 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -183,16 +183,6 @@ class _DatetimeRoundingMethods(Generic[_DTRoundingMethodReturnType]):
183183
| Timedelta = ...,
184184
) -> _DTRoundingMethodReturnType: ...
185185

186-
class _TZConversionMethods:
187-
def tz_localize(
188-
self,
189-
tz: str | None,
190-
ambiguous: Literal["raise", "infer", "NaT"] | np_ndarray_bool = ...,
191-
nonexistent: Literal["shift_forward", "shift_backward", "NaT", "raise"]
192-
| Timedelta = ...,
193-
) -> DatetimeArray: ...
194-
def tz_convert(self, tz: str | None) -> TimestampSeries: ...
195-
196186
_DTNormalizeReturnType = TypeVar(
197187
"_DTNormalizeReturnType", TimestampSeries, DatetimeIndex
198188
)
@@ -217,8 +207,8 @@ class _DatetimeLikeNoTZMethods(
217207
ambiguous: Literal["raise", "infer", "NaT"] | np_ndarray_bool = ...,
218208
nonexistent: Literal["shift_forward", "shift_backward", "NaT", "raise"]
219209
| Timedelta = ...,
220-
) -> DatetimeArray: ...
221-
def tz_convert(self, tz: str | None) -> TimestampSeries: ...
210+
) -> _DTNormalizeReturnType: ...
211+
def tz_convert(self, tz: str | None) -> _DTNormalizeReturnType: ...
222212
def normalize(self) -> _DTNormalizeReturnType: ...
223213
def strftime(self, date_format: str) -> _DTStrKindReturnType: ...
224214
def month_name(self, locale: str | None = ...) -> _DTStrKindReturnType: ...
@@ -264,7 +254,6 @@ class DatetimeProperties(
264254
_DTStrKindReturnType,
265255
_DTToPeriodReturnType,
266256
],
267-
_TZConversionMethods,
268257
Generic[
269258
_DTFieldOpsReturnType,
270259
_DTBoolOpsReturnType,

‎pandas-stubs/core/indexes/datetimes.pyi

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@ from pandas.core.dtypes.dtypes import DatetimeTZDtype
3232
from pandas.tseries.offsets import BaseOffset
3333

3434
class DatetimeIndex(DatetimeTimedeltaMixin, DatetimeIndexProperties):
35-
tz: tzinfo | None
3635
def __init__(
3736
self,
3837
data: ArrayLike | AnyArrayLike | list | tuple,

‎pandas-stubs/core/indexing.pyi

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,14 @@ from pandas._typing import (
1313
ScalarT,
1414
)
1515

16-
_IndexSliceTuple = Union[
17-
slice, tuple[Union[Index, MaskType, Scalar, list[ScalarT], slice], ...]
18-
]
16+
_IndexSliceTuple = tuple[Union[Index, MaskType, Scalar, list[ScalarT], slice], ...]
1917

20-
_IndexSliceTupleT = TypeVar("_IndexSliceTupleT", bound=_IndexSliceTuple)
18+
_IndexSliceUnion = Union[slice, _IndexSliceTuple]
19+
20+
_IndexSliceUnionT = TypeVar("_IndexSliceUnionT", bound=_IndexSliceUnion)
2121

2222
class _IndexSlice:
23-
def __getitem__(self, arg: _IndexSliceTupleT) -> _IndexSliceTupleT: ...
23+
def __getitem__(self, arg: _IndexSliceUnionT) -> _IndexSliceUnionT: ...
2424

2525
IndexSlice: _IndexSlice
2626

‎pandas-stubs/core/series.pyi

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -727,35 +727,35 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
727727
@overload
728728
def drop(
729729
self,
730-
labels: Hashable | list[HashableT] = ...,
730+
labels: Hashable | list[HashableT] | Index = ...,
731731
*,
732732
axis: Axis = ...,
733-
index: Hashable | list[HashableT] = ...,
734-
columns: Hashable | list[HashableT] = ...,
733+
index: Hashable | list[HashableT] | Index = ...,
734+
columns: Hashable | list[HashableT] | Index = ...,
735735
level: Level | None = ...,
736736
inplace: Literal[True],
737737
errors: IgnoreRaise = ...,
738738
) -> None: ...
739739
@overload
740740
def drop(
741741
self,
742-
labels: Hashable | list[HashableT] = ...,
742+
labels: Hashable | list[HashableT] | Index = ...,
743743
*,
744744
axis: Axis = ...,
745-
index: Hashable | list[HashableT] = ...,
746-
columns: Hashable | list[HashableT] = ...,
745+
index: Hashable | list[HashableT] | Index = ...,
746+
columns: Hashable | list[HashableT] | Index = ...,
747747
level: Level | None = ...,
748748
inplace: Literal[False] = ...,
749749
errors: IgnoreRaise = ...,
750750
) -> Series: ...
751751
@overload
752752
def drop(
753753
self,
754-
labels: Hashable | list[HashableT] = ...,
754+
labels: Hashable | list[HashableT] | Index = ...,
755755
*,
756756
axis: Axis = ...,
757-
index: Hashable | list[HashableT] = ...,
758-
columns: Hashable | list[HashableT] = ...,
757+
index: Hashable | list[HashableT] | Index = ...,
758+
columns: Hashable | list[HashableT] | Index = ...,
759759
level: Level | None = ...,
760760
inplace: bool = ...,
761761
errors: IgnoreRaise = ...,

‎tests/test_frame.py

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -244,16 +244,20 @@ def test_types_setting() -> None:
244244

245245
def test_types_drop() -> None:
246246
df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
247-
res: pd.DataFrame = df.drop("col1", axis=1)
248-
res2: pd.DataFrame = df.drop(columns=["col1"])
249-
res3: pd.DataFrame = df.drop([0])
250-
res4: pd.DataFrame = df.drop(index=[0])
251-
res5: pd.DataFrame = df.drop(columns=["col1"])
252-
res6: pd.DataFrame = df.drop(index=1)
253-
res7: pd.DataFrame = df.drop(labels=0)
254-
res8: None = df.drop([0, 0], inplace=True)
247+
check(assert_type(df.drop("col1", axis=1), pd.DataFrame), pd.DataFrame)
248+
check(assert_type(df.drop(columns=["col1"]), pd.DataFrame), pd.DataFrame)
249+
check(assert_type(df.drop([0]), pd.DataFrame), pd.DataFrame)
250+
check(assert_type(df.drop(index=[0]), pd.DataFrame), pd.DataFrame)
251+
check(assert_type(df.drop(columns=["col1"]), pd.DataFrame), pd.DataFrame)
252+
check(assert_type(df.drop(index=1), pd.DataFrame), pd.DataFrame)
253+
check(assert_type(df.drop(labels=0), pd.DataFrame), pd.DataFrame)
254+
assert assert_type(df.drop([0, 0], inplace=True), None) is None
255255
to_drop: list[str] = ["col1"]
256-
res9: pd.DataFrame = df.drop(columns=to_drop)
256+
check(assert_type(df.drop(columns=to_drop), pd.DataFrame), pd.DataFrame)
257+
# GH 302
258+
check(assert_type(df.drop(pd.Index([1])), pd.DataFrame), pd.DataFrame)
259+
check(assert_type(df.drop(index=pd.Index([1])), pd.DataFrame), pd.DataFrame)
260+
check(assert_type(df.drop(columns=pd.Index(["col1"])), pd.DataFrame), pd.DataFrame)
257261

258262

259263
def test_types_dropna() -> None:
@@ -1262,6 +1266,8 @@ def test_indexslice_setitem():
12621266
s = pd.Series([-1, -2])
12631267
df.loc[pd.IndexSlice[2, :]] = s.values
12641268
df.loc[pd.IndexSlice[2, :], "z"] = [200, 300]
1269+
# GH 314
1270+
df.loc[pd.IndexSlice[pd.Index([2, 3]), :], "z"] = 99
12651271

12661272

12671273
def test_indexslice_getitem():
@@ -1284,6 +1290,11 @@ def test_indexslice_getitem():
12841290
assert_type(df.loc[pd.IndexSlice[[2, 4], [20, 40]], :], pd.DataFrame),
12851291
pd.DataFrame,
12861292
)
1293+
# GH 314
1294+
check(
1295+
assert_type(df.loc[pd.IndexSlice[pd.Index([2, 4]), :], "z"], pd.Series),
1296+
pd.Series,
1297+
)
12871298

12881299

12891300
def test_compute_values():

‎tests/test_series.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,10 @@ def test_types_drop() -> None:
183183
check(assert_type(s.drop(0, axis=0), pd.Series), pd.Series)
184184
assert assert_type(s.drop([0, 1], inplace=True, errors="raise"), None) is None
185185
assert assert_type(s.drop([0, 1], inplace=True, errors="ignore"), None) is None
186+
# GH 302
187+
s = pd.Series([0, 1, 2])
188+
check(assert_type(s.drop(pd.Index([0, 1])), pd.Series), pd.Series)
189+
check(assert_type(s.drop(index=pd.Index([0, 1])), pd.Series), pd.Series)
186190

187191

188192
def test_types_drop_multilevel() -> None:

‎tests/test_timefuncs.py

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -336,8 +336,8 @@ def test_series_dt_accessors() -> None:
336336
check(assert_type(s0.dt.isocalendar(), pd.DataFrame), pd.DataFrame)
337337
check(assert_type(s0.dt.to_period("D"), "PeriodSeries"), pd.Series, pd.Period)
338338
check(assert_type(s0.dt.to_pydatetime(), np.ndarray), np.ndarray, dt.datetime)
339-
local_dtarray = s0.dt.tz_localize("UTC")
340-
slocal = pd.Series(local_dtarray)
339+
slocal = s0.dt.tz_localize("UTC")
340+
check(assert_type(slocal, "TimestampSeries"), pd.Series, pd.Timestamp)
341341
check(
342342
assert_type(slocal.dt.tz_convert("EST"), "TimestampSeries"),
343343
pd.Series,
@@ -414,7 +414,7 @@ def test_datetimeindex_accessors() -> None:
414414
check(assert_type(i0.is_leap_year, npt.NDArray[np.bool_]), np.ndarray, np.bool_)
415415
check(assert_type(i0.daysinmonth, IntegerIndex), IntegerIndex, int)
416416
check(assert_type(i0.days_in_month, IntegerIndex), IntegerIndex, int)
417-
assert assert_type(i0.tz, Optional[dt.tzinfo]) is None
417+
assert assert_type(i0.tz, Optional[Union[dt.tzinfo, BaseTzInfo]]) is None
418418
check(assert_type(i0.freq, Optional[BaseOffset]), BaseOffset)
419419
check(assert_type(i0.isocalendar(), pd.DataFrame), pd.DataFrame)
420420
check(assert_type(i0.to_period("D"), pd.PeriodIndex), pd.PeriodIndex, pd.Period)
@@ -423,14 +423,10 @@ def test_datetimeindex_accessors() -> None:
423423
np.ndarray,
424424
dt.datetime,
425425
)
426-
local_dtarray = i0.tz_localize("UTC")
427-
slocal = pd.Series(local_dtarray)
428-
check(
429-
assert_type(slocal.dt.tz_convert("EST"), "TimestampSeries"),
430-
pd.Series,
431-
pd.Timestamp,
432-
)
433-
check(assert_type(slocal.dt.tz, Optional[Union[dt.tzinfo, BaseTzInfo]]), BaseTzInfo)
426+
slocal = i0.tz_localize("UTC")
427+
check(assert_type(slocal, pd.DatetimeIndex), pd.DatetimeIndex)
428+
check(assert_type(slocal.tz_convert("EST"), pd.DatetimeIndex), pd.DatetimeIndex)
429+
check(assert_type(slocal.tz, Optional[Union[dt.tzinfo, BaseTzInfo]]), BaseTzInfo)
434430
check(assert_type(i0.normalize(), pd.DatetimeIndex), pd.DatetimeIndex, pd.Timestamp)
435431
check(assert_type(i0.strftime("%Y"), pd.Index), pd.Index, str)
436432
check(assert_type(i0.round("D"), pd.DatetimeIndex), pd.DatetimeIndex, pd.Timestamp)
@@ -517,6 +513,11 @@ def test_some_offsets() -> None:
517513
),
518514
pd.DatetimeIndex,
519515
)
516+
# GH 320
517+
tswm1 = pd.Timestamp("9/23/2022") + pd.offsets.WeekOfMonth(2, 3)
518+
check(assert_type(tswm1, pd.Timestamp), pd.Timestamp)
519+
tswm2 = pd.Timestamp("9/23/2022") + pd.offsets.LastWeekOfMonth(2, 3)
520+
check(assert_type(tswm2, pd.Timestamp), pd.Timestamp)
520521

521522

522523
def test_types_to_numpy() -> None:

0 commit comments

Comments
 (0)