Skip to content

Commit cd5be53

Browse files
cmp0xffcmp0xff
committed
feat(index): arithmetic addition
1 parent d20612e commit cd5be53

File tree

19 files changed

+824
-150
lines changed

19 files changed

+824
-150
lines changed

‎pandas-stubs/_typing.pyi‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -845,6 +845,8 @@ MaskType: TypeAlias = Series[bool] | np_ndarray_bool | list[bool]
845845

846846
# Scratch types for generics
847847

848+
T_INT = TypeVar("T_INT", bound=int)
849+
T_COMPLEX = TypeVar("T_COMPLEX", bound=complex)
848850
SeriesDType: TypeAlias = (
849851
str
850852
| bytes

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

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,12 +39,15 @@ from pandas.core.strings.accessor import StringMethods
3939
from typing_extensions import (
4040
Never,
4141
Self,
42+
TypeAlias,
4243
)
4344

4445
from pandas._libs.interval import _OrderableT
4546
from pandas._typing import (
4647
C2,
4748
S1,
49+
T_COMPLEX,
50+
T_INT,
4851
AnyAll,
4952
ArrayLike,
5053
AxesData,
@@ -70,18 +73,32 @@ from pandas._typing import (
7073
TimestampDtypeArg,
7174
np_1darray,
7275
np_ndarray_anyint,
76+
np_ndarray_bool,
7377
np_ndarray_complex,
7478
np_ndarray_float,
79+
np_ndarray_str,
7580
type_t,
7681
)
7782

7883
class InvalidIndexError(Exception): ...
7984

85+
_ListLike: TypeAlias = ArrayLike | dict[_str, np.ndarray] | SequenceNotStr[S1]
86+
8087
class Index(IndexOpsMixin[S1]):
8188
__hash__: ClassVar[None] # type: ignore[assignment]
8289
# overloads with additional dtypes
8390
@overload
8491
def __new__( # pyright: ignore[reportOverlappingOverload]
92+
cls,
93+
data: Sequence[bool | np.bool_] | IndexOpsMixin[bool] | np_ndarray_bool,
94+
*,
95+
dtype: Literal["bool"] | type_t[bool | np.bool_] = ...,
96+
copy: bool = ...,
97+
name: Hashable = ...,
98+
tupleize_cols: bool = ...,
99+
) -> Index[bool]: ...
100+
@overload
101+
def __new__(
85102
cls,
86103
data: Sequence[int | np.integer] | IndexOpsMixin[int] | np_ndarray_anyint,
87104
*,
@@ -460,6 +477,155 @@ class Index(IndexOpsMixin[S1]):
460477
def __gt__(self, other: Self | S1) -> np_1darray[np.bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
461478
# overwrite inherited methods from OpsMixin
462479
@overload
480+
def __add__(self: Index[Never], other: _str) -> Never: ...
481+
@overload
482+
def __add__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ...
483+
@overload
484+
def __add__(self, other: Index[Never]) -> Index: ...
485+
@overload
486+
def __add__(
487+
self: Index[bool],
488+
other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX],
489+
) -> Index[T_COMPLEX]: ...
490+
@overload
491+
def __add__(self: Index[bool], other: np_ndarray_bool) -> Index[bool]: ...
492+
@overload
493+
def __add__(self: Index[bool], other: np_ndarray_anyint) -> Index[int]: ...
494+
@overload
495+
def __add__(self: Index[bool], other: np_ndarray_float) -> Index[float]: ...
496+
@overload
497+
def __add__(self: Index[bool], other: np_ndarray_complex) -> Index[complex]: ...
498+
@overload
499+
def __add__(
500+
self: Index[int],
501+
other: (
502+
bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Index[bool]
503+
),
504+
) -> Index[int]: ...
505+
@overload
506+
def __add__(
507+
self: Index[int],
508+
other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX],
509+
) -> Index[T_COMPLEX]: ...
510+
@overload
511+
def __add__(self: Index[int], other: np_ndarray_float) -> Index[float]: ...
512+
@overload
513+
def __add__(self: Index[int], other: np_ndarray_complex) -> Index[complex]: ...
514+
@overload
515+
def __add__(
516+
self: Index[float],
517+
other: (
518+
int
519+
| Sequence[int]
520+
| np_ndarray_bool
521+
| np_ndarray_anyint
522+
| np_ndarray_float
523+
| Index[T_INT]
524+
),
525+
) -> Index[float]: ...
526+
@overload
527+
def __add__(
528+
self: Index[float],
529+
other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX],
530+
) -> Index[T_COMPLEX]: ...
531+
@overload
532+
def __add__(self: Index[float], other: np_ndarray_complex) -> Index[complex]: ...
533+
@overload
534+
def __add__(
535+
self: Index[complex],
536+
other: (
537+
T_COMPLEX
538+
| Sequence[T_COMPLEX]
539+
| np_ndarray_bool
540+
| np_ndarray_anyint
541+
| np_ndarray_float
542+
| np_ndarray_complex
543+
| Index[T_COMPLEX]
544+
),
545+
) -> Index[complex]: ...
546+
@overload
547+
def __add__(
548+
self: Index[_str],
549+
other: (
550+
np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex
551+
),
552+
) -> Never: ...
553+
@overload
554+
def __add__(
555+
self: Index[_str], other: _str | Sequence[_str] | np_ndarray_str | Index[_str]
556+
) -> Index[_str]: ...
557+
@overload # type: ignore[override]
558+
def __radd__(self: Index[Never], other: _str) -> Never: ...
559+
@overload
560+
def __radd__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ...
561+
@overload
562+
def __radd__(
563+
self: Index[bool],
564+
other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX],
565+
) -> Index[T_COMPLEX]: ...
566+
@overload
567+
def __radd__(self: Index[bool], other: np_ndarray_bool) -> Index[bool]: ...
568+
@overload
569+
def __radd__(self: Index[bool], other: np_ndarray_anyint) -> Index[int]: ...
570+
@overload
571+
def __radd__(self: Index[bool], other: np_ndarray_float) -> Index[float]: ...
572+
@overload
573+
def __radd__(
574+
self: Index[int],
575+
other: (
576+
bool | Sequence[bool] | np_ndarray_bool | np_ndarray_anyint | Index[bool]
577+
),
578+
) -> Index[int]: ...
579+
@overload
580+
def __radd__(
581+
self: Index[int], other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX]
582+
) -> Index[T_COMPLEX]: ...
583+
@overload
584+
def __radd__(self: Index[int], other: np_ndarray_float) -> Index[float]: ...
585+
@overload
586+
def __radd__(
587+
self: Index[float],
588+
other: (
589+
int
590+
| Sequence[int]
591+
| np_ndarray_bool
592+
| np_ndarray_anyint
593+
| np_ndarray_float
594+
| Index[T_INT]
595+
),
596+
) -> Index[float]: ...
597+
@overload
598+
def __radd__(
599+
self: Index[float], other: T_COMPLEX | Sequence[T_COMPLEX] | Index[T_COMPLEX]
600+
) -> Index[T_COMPLEX]: ...
601+
@overload
602+
def __radd__(
603+
self: Index[complex],
604+
other: (
605+
T_COMPLEX
606+
| Sequence[T_COMPLEX]
607+
| np_ndarray_bool
608+
| np_ndarray_anyint
609+
| np_ndarray_float
610+
| Index[T_COMPLEX]
611+
),
612+
) -> Index[complex]: ...
613+
@overload
614+
def __radd__(
615+
self: Index[T_COMPLEX], other: np_ndarray_complex
616+
) -> Index[complex]: ...
617+
@overload
618+
def __radd__(
619+
self: Index[_str],
620+
other: (
621+
np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex
622+
),
623+
) -> Never: ...
624+
@overload
625+
def __radd__(
626+
self: Index[_str], other: _str | Sequence[_str] | np_ndarray_str | Index[_str]
627+
) -> Index[_str]: ...
628+
@overload
463629
def __mul__(
464630
self: Index[int] | Index[float], other: timedelta
465631
) -> TimedeltaIndex: ...

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,10 @@ class DatetimeIndex(
6161
def __reduce__(self): ...
6262
# various ignores needed for mypy, as we do want to restrict what can be used in
6363
# arithmetic for these types
64-
@overload
64+
@overload # type: ignore[override]
6565
def __add__(self, other: TimedeltaSeries) -> TimestampSeries: ...
6666
@overload
67-
def __add__(
67+
def __add__( # pyright: ignore[reportIncompatibleMethodOverride]
6868
self, other: timedelta | Timedelta | TimedeltaIndex | BaseOffset
6969
) -> DatetimeIndex: ...
7070
@overload

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,12 +49,14 @@ class TimedeltaIndex(
4949
) -> Self: ...
5050
# various ignores needed for mypy, as we do want to restrict what can be used in
5151
# arithmetic for these types
52-
@overload
52+
@overload # type: ignore[override]
5353
def __add__(self, other: Period) -> PeriodIndex: ...
5454
@overload
5555
def __add__(self, other: DatetimeIndex) -> DatetimeIndex: ...
5656
@overload
57-
def __add__(self, other: dt.timedelta | Timedelta | Self) -> Self: ...
57+
def __add__( # pyright: ignore[reportIncompatibleMethodOverride]
58+
self, other: dt.timedelta | Timedelta | Self
59+
) -> Self: ...
5860
def __radd__(self, other: dt.datetime | Timestamp | DatetimeIndex) -> DatetimeIndex: ... # type: ignore[override]
5961
def __sub__(self, other: dt.timedelta | Timedelta | Self) -> Self: ...
6062
def __mul__(self, other: num) -> Self: ...

0 commit comments

Comments
 (0)