Skip to content

ENH: Improve typing for Interval #391

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 16 commits into from
Nov 24, 2022
Merged
64 changes: 57 additions & 7 deletions pandas-stubs/_libs/interval.pyi
Original file line number Diff line number Diff line change
@@ -1,18 +1,23 @@
from typing import (
Any,
Generic,
Literal,
TypeVar,
overload,
)

import numpy as np
from pandas import (
IntervalIndex,
Series,
Timedelta,
Timestamp,
)

from pandas._typing import (
IntervalClosedType,
IntervalT,
np_ndarray_bool,
npt,
)

Expand Down Expand Up @@ -123,28 +128,73 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
@overload
def __mul__(self: Interval[float], y: float) -> Interval[float]: ...
@overload
def __mul__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
@overload
def __rmul__(
self: Interval[int], y: _OrderableScalarT
) -> Interval[_OrderableScalarT]: ...
@overload
def __rmul__(self: Interval[float], y: float) -> Interval[float]: ...
@overload
def __truediv__(
self: Interval[int], y: _OrderableScalarT
) -> Interval[_OrderableScalarT]: ...
def __rmul__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
@overload
def __truediv__(self: Interval[int], y: _OrderableScalarT) -> Interval[float]: ...
@overload
def __truediv__(self: Interval[float], y: float) -> Interval[float]: ...
@overload
def __truediv__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
@overload
def __floordiv__(
self: Interval[int], y: _OrderableScalarT
) -> Interval[_OrderableScalarT]: ...
@overload
def __floordiv__(self: Interval[float], y: float) -> Interval[float]: ...
@overload
def __floordiv__(self: Interval[Timedelta], y: float) -> Interval[Timedelta]: ...
@overload
def overlaps(self: Interval[_OrderableT], other: Interval[_OrderableT]) -> bool: ...

def intervals_to_interval_bounds(
intervals: np.ndarray, validate_closed: bool = ...
) -> tuple[np.ndarray, np.ndarray, str]: ...
@overload
def overlaps(self: Interval[int], other: Interval[float]) -> bool: ...
@overload
def overlaps(self: Interval[float], other: Interval[int]) -> bool: ...
@overload
def __gt__(self, other: Interval[_OrderableT]) -> bool: ...
@overload
def __gt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
@overload
def __gt__(self, other: Series[_OrderableT]) -> Series[bool]: ...
@overload
def __lt__(self, other: Interval[_OrderableT]) -> bool: ...
@overload
def __lt__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
@overload
def __lt__(self, other: Series[_OrderableT]) -> Series[bool]: ...
@overload
def __ge__(self, other: Interval[_OrderableT]) -> bool: ...
@overload
def __ge__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
@overload
def __ge__(self, other: Series[_OrderableT]) -> Series[bool]: ...
@overload
def __le__(self, other: Interval[_OrderableT]) -> bool: ...
@overload
def __le__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ...
@overload
def __eq__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[misc]
@overload
def __eq__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc]
@overload
def __eq__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[misc]
@overload
def __eq__(self, other: object) -> Literal[False]: ...
@overload
def __ne__(self, other: Interval[_OrderableT]) -> bool: ... # type: ignore[misc]
@overload
def __ne__(self: IntervalT, other: IntervalIndex[IntervalT]) -> np_ndarray_bool: ... # type: ignore[misc]
@overload
def __ne__(self, other: Series[_OrderableT]) -> Series[bool]: ... # type: ignore[misc]
@overload
def __ne__(self, other: object) -> Literal[True]: ...

class IntervalTree(IntervalMixin):
def __init__(
Expand Down
13 changes: 12 additions & 1 deletion pandas-stubs/_typing.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ from pandas.core.indexes.base import Index
from pandas.core.series import Series
from typing_extensions import TypeAlias

from pandas._libs.interval import Interval
from pandas._libs.tslibs import (
Period,
Timedelta,
Expand Down Expand Up @@ -196,6 +197,10 @@ S1 = TypeVar(
Timedelta,
np.datetime64,
Period,
Interval[int],
Interval[float],
Interval[Timestamp],
Interval[Timedelta],
)
T1 = TypeVar(
"T1", str, int, np.int64, np.uint64, np.float64, float, np.dtype[np.generic]
Expand All @@ -220,7 +225,13 @@ NDFrameT = TypeVar("NDFrameT", bound=NDFrame)
IndexT = TypeVar("IndexT", bound=Index)

# Interval closed type

IntervalT = TypeVar(
"IntervalT",
Interval[int],
Interval[float],
Interval[Timestamp],
Interval[Timedelta],
)
IntervalClosedType: TypeAlias = Literal["left", "right", "both", "neither"]

IgnoreRaiseCoerce: TypeAlias = Literal["ignore", "raise", "coerce"]
Expand Down
8 changes: 6 additions & 2 deletions pandas-stubs/core/algorithms.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,20 @@ from pandas import (
)
from pandas.api.extensions import ExtensionArray

from pandas._typing import AnyArrayLike
from pandas._typing import (
AnyArrayLike,
IntervalT,
)

# These are type: ignored because the Index types overlap due to inheritance but indices
# with extension types return the same type while standard type return ndarray

@overload
def unique(values: PeriodIndex) -> PeriodIndex: ... # type: ignore[misc]
@overload
def unique(values: CategoricalIndex) -> CategoricalIndex: ... # type: ignore[misc]
@overload
def unique(values: IntervalIndex) -> IntervalIndex: ... # type: ignore[misc]
def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ... # type: ignore[misc]
@overload
def unique(values: Index) -> np.ndarray: ...
@overload
Expand Down
Loading