Skip to content

Commit 8d71ade

Browse files
Dr-IrvDr-Irv
authored
fix up pandas stubs to make tests pass (#149)
Fix up pandas stubs to make tests pass: * fix issue with pd.Interval.length * fix divmod * move interval specific types out of _typing.pyi * on interval limit scalar types Thanks to Dr Irv for this PR!
1 parent 518bade commit 8d71ade

File tree

12 files changed

+1131
-603
lines changed

12 files changed

+1131
-603
lines changed

‎pandas/_libs/interval.pyi

Lines changed: 49 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,43 +1,71 @@
11
from __future__ import annotations
22
import sys
3-
from typing import Any, Union
3+
from typing import Generic, overload, Union, Protocol, TypeVar
4+
from _typing import Timedelta, Timestamp
5+
import datetime
6+
47
if sys.version_info >= (3, 8):
58
from typing import Literal
69
else:
710
from typing_extensions import Literal
811

12+
OrderableScalar = TypeVar("OrderableScalar", int, float)
13+
OrderableTimes = TypeVar("OrderableTimes", datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta)
14+
Orderable = TypeVar("Orderable", int, float, datetime.date, datetime.datetime, datetime.timedelta, Timestamp, Timedelta)
15+
16+
class IntervalMixinProtocol(Protocol): ...
17+
918
class IntervalMixin:
1019
@property
11-
def closed_left(self) -> bool: ...
20+
def closed_left(self: IntervalMixinProtocol) -> bool: ...
1221
@property
13-
def closed_right(self) -> bool: ...
22+
def closed_right(self: IntervalMixinProtocol) -> bool: ...
1423
@property
15-
def open_left(self) -> bool: ...
24+
def open_left(self: IntervalMixinProtocol) -> bool: ...
1625
@property
17-
def open_right(self) -> bool: ...
26+
def open_right(self: IntervalMixinProtocol) -> bool: ...
1827
@property
19-
def mid(self) -> float: ...
28+
def mid(self: IntervalMixinProtocol) -> float: ...
2029
@property
21-
def length(self) -> float: ...
22-
@property
23-
def is_empty(self) -> bool: ...
24-
30+
def is_empty(self: IntervalMixinProtocol) -> bool: ...
2531

26-
class Interval(IntervalMixin):
32+
class Interval(IntervalMixin, Generic[Orderable]):
33+
@overload
34+
def left(self: Interval[OrderableScalar]) -> OrderableScalar: ...
35+
@overload
36+
def left(self: Interval[OrderableTimes]) -> OrderableTimes: ...
2737
@property
28-
def left(self) -> object: ...
38+
@overload
39+
def left(self: Interval[Orderable]) -> Orderable: ...
40+
@overload
41+
def right(self: Interval[OrderableScalar]) -> OrderableScalar: ...
42+
@overload
43+
def right(self: Interval[OrderableTimes]) -> OrderableTimes: ...
2944
@property
30-
def right(self) -> object: ...
45+
@overload
46+
def right(self: Interval[Orderable]) -> Orderable: ...
3147
@property
3248
def closed(self) -> str: ...
33-
def __init__(self, left: object, right: object, closed: Union[str, Literal['left', 'right', 'both', 'neither']] = ...) -> None: ...
49+
def __init__(
50+
self,
51+
left: Orderable,
52+
right: Orderable,
53+
closed: Union[str, Literal["left", "right", "both", "neither"]] = ...,
54+
) -> None: ...
55+
@overload
56+
def length(self: Interval[OrderableScalar]) -> float: ...
57+
@overload
58+
def length(self: Interval[OrderableTimes]) -> Timedelta: ...
59+
@property
60+
@overload
61+
def length(self: Interval[Orderable]) -> Orderable: ...
3462
def __hash__(self) -> int: ...
35-
def __contains__(self, key: Any) -> bool: ...
63+
def __contains__(self, key: Orderable) -> bool: ...
3664
def __repr__(self) -> str: ...
3765
def __str__(self) -> str: ...
38-
def __add__(self, y: float) -> Interval: ...
39-
def __sub__(self, y: float) -> Interval: ...
40-
def __mul__(self, y: float) -> Interval: ...
41-
def __truediv__(self, y: float) -> Interval: ...
42-
def __floordiv__(self, y: float) -> Interval: ...
43-
def overlaps(self, other: Interval) -> bool: ...
66+
def __add__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ...
67+
def __sub__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ...
68+
def __mul__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ...
69+
def __truediv__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ...
70+
def __floordiv__(self: Interval[Orderable], y: float) -> Interval[Orderable]: ...
71+
def overlaps(self: Interval[Orderable], other: Interval[Orderable]) -> bool: ...

‎pandas/_typing.pyi

Lines changed: 29 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,40 +5,57 @@ import numpy as np
55
import sys
66
from os import PathLike
77
from pathlib import Path
8-
from typing import Any, AnyStr, Callable, Collection, Dict, Hashable, IO, List, Mapping, NewType, Optional, Sequence, Type, TypeVar, Union
8+
from typing import (
9+
Any,
10+
AnyStr,
11+
Callable,
12+
Collection,
13+
Dict,
14+
Hashable,
15+
IO,
16+
List,
17+
Mapping,
18+
NewType,
19+
Optional,
20+
Sequence,
21+
Type,
22+
TypeVar,
23+
Union,
24+
)
925

1026
from pandas.core.generic import NDFrame
1127
from pandas._libs.tslibs import Period, Timedelta, Timestamp
1228
from pandas.core.arrays import ExtensionArray
1329
from pandas.core.series import Series
1430
from pandas.core.indexes.base import Index
31+
1532
if sys.version_info >= (3, 8):
1633
from typing import Literal
1734
else:
1835
from typing_extensions import Literal
1936

20-
AnyArrayLike = TypeVar('AnyArrayLike', ExtensionArray, Index, Series, np.ndarray)
21-
ArrayLike = TypeVar('ArrayLike', ExtensionArray, np.ndarray)
37+
AnyArrayLike = TypeVar("AnyArrayLike", ExtensionArray, Index, Series, np.ndarray)
38+
ArrayLike = TypeVar("ArrayLike", ExtensionArray, np.ndarray)
2239
PythonScalar = Union[str, int, float, bool, complex]
23-
DatetimeLikeScalar = TypeVar('DatetimeLikeScalar', Period, Timestamp, Timedelta)
40+
DatetimeLikeScalar = TypeVar("DatetimeLikeScalar", Period, Timestamp, Timedelta)
2441
PandasScalar = Union[bytes, datetime.date, datetime.datetime, datetime.timedelta]
25-
#Scalar = Union[PythonScalar, PandasScalar]
26-
#Dtype: Any
42+
# Scalar = Union[PythonScalar, PandasScalar]
43+
# Dtype: Any
2744

2845
# filenames and file-like-objects
2946
Buffer = Union[IO[AnyStr], RawIOBase, BufferedIOBase, TextIOBase, TextIOWrapper, mmap]
3047
FileOrBuffer = Union[str, Buffer[AnyStr]]
3148
FilePathOrBuffer = Union["PathLike[str]", FileOrBuffer[AnyStr]]
3249

33-
FrameOrSeries = TypeVar('FrameOrSeries', bound=NDFrame)
50+
FrameOrSeries = TypeVar("FrameOrSeries", bound=NDFrame)
3451
Axis = Union[str, int]
3552
Label = Optional[Hashable]
3653
Level = Union[Label, int]
3754
Ordered = Optional[bool]
3855
JSONSerializable = Union[PythonScalar, List, Dict]
3956
Axes = Collection
4057
Renamer = Union[Mapping[Label, Any], Callable[[Label], Label]]
41-
T = TypeVar('T')
58+
T = TypeVar("T")
4259
FuncType = Callable[..., Any]
4360
F = TypeVar("F", bound=FuncType)
4461

@@ -56,7 +73,7 @@ AxisType = Union[str, int, Literal["columns", "index", 0, 1]]
5673
Dtype = TypeVar("Dtype", bool, int, float, object)
5774
DtypeNp = TypeVar("DtypeNp", bound=np.dtype)
5875
KeysArgType = Any
59-
ListLike = TypeVar("ListLike", Sequence, np.ndarray, 'Series')
76+
ListLike = TypeVar("ListLike", Sequence, np.ndarray, "Series")
6077
StrLike = Union[str, np.str_]
6178
Scalar = Union[str, bytes, datetime.date, datetime.datetime, datetime.timedelta, bool, int, float, complex, Timestamp, Timedelta]
6279
# Refine the next 3 in 3.9 to use the specialized type.
@@ -66,8 +83,8 @@ np_ndarray_str = NewType("np_ndarray_str", Type[np.ndarray])
6683
IndexType = Union[slice, np_ndarray_int64, Index[int], List[int], Series[int]]
6784
MaskType = Union[Series[bool], np_ndarray_bool, Sequence[bool]]
6885
# Scratch types for generics
69-
S1 = TypeVar("S1", str, bytes, datetime.date, datetime.datetime, datetime.timedelta, bool, int, float, complex, Timestamp, Timedelta)
86+
S1 = TypeVar(
87+
"S1", str, bytes, datetime.date, datetime.datetime, datetime.timedelta, bool, int, float, complex, Timestamp, Timedelta
88+
)
7089
T1 = TypeVar("T1", str, int)
7190
T2 = TypeVar("T2", str, int)
72-
73-

‎pandas/core/arraylike.pyi

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
from typing import Any, Protocol, Tuple
2+
from pandas._typing import FrameOrSeries, DataFrame
3+
4+
class OpsMixinProtocol(Protocol): ...
5+
6+
class OpsMixin:
7+
def __eq__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
8+
def __ne__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
9+
def __lt__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
10+
def __le__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
11+
def __gt__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
12+
def __ge__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
13+
# -------------------------------------------------------------
14+
# Logical Methods
15+
def __and__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
16+
def __rand__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
17+
def __or__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
18+
def __ror__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
19+
def __xor__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
20+
def __rxor__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
21+
# -------------------------------------------------------------
22+
# Arithmetic Methods
23+
def __add__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
24+
def __radd__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
25+
def __sub__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
26+
def __rsub__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
27+
def __mul__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
28+
def __rmul__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
29+
def __truediv__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
30+
def __rtruediv__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
31+
def __floordiv__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
32+
def __rfloordiv__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
33+
def __mod__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
34+
def __rmod__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
35+
def __divmod__(self: OpsMixinProtocol, other: DataFrame) -> Tuple[DataFrame, DataFrame]: ...
36+
def __rdivmod__(self: OpsMixinProtocol, other: DataFrame) -> Tuple[DataFrame, DataFrame]: ...
37+
def __pow__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
38+
def __rpow__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...

0 commit comments

Comments
 (0)