Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 6 additions & 25 deletions pandas-stubs/_libs/interval.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ from typing import (
Literal,
TypeVar,
overload,
type_check_only,
)

import numpy as np
Expand All @@ -18,7 +19,6 @@ from pandas._typing import (
IntervalClosedType,
IntervalT,
np_1darray,
npt,
)

VALID_CLOSED: frozenset[str]
Expand All @@ -27,6 +27,7 @@ _OrderableScalarT = TypeVar("_OrderableScalarT", bound=int | float)
_OrderableTimesT = TypeVar("_OrderableTimesT", bound=Timestamp | Timedelta)
_OrderableT = TypeVar("_OrderableT", bound=int | float | Timestamp | Timedelta)

@type_check_only
class _LengthDescriptor:
@overload
def __get__(
Expand All @@ -36,18 +37,15 @@ class _LengthDescriptor:
def __get__(
self, instance: Interval[_OrderableTimesT], owner: Any
) -> Timedelta: ...
@overload
def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ...

@type_check_only
class _MidDescriptor:
@overload
def __get__(self, instance: Interval[_OrderableScalarT], owner: Any) -> float: ...
@overload
def __get__(
self, instance: Interval[_OrderableTimesT], owner: Any
) -> _OrderableTimesT: ...
@overload
def __get__(self, instance: IntervalTree, owner: Any) -> np.ndarray: ...

class IntervalMixin:
@property
Expand All @@ -68,8 +66,8 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
def right(self: Interval[_OrderableT]) -> _OrderableT: ...
@property
def closed(self) -> IntervalClosedType: ...
mid: _MidDescriptor
length: _LengthDescriptor
mid = _MidDescriptor()
length = _LengthDescriptor()
def __init__(
self,
left: _OrderableT,
Expand Down Expand Up @@ -223,21 +221,4 @@ class Interval(IntervalMixin, Generic[_OrderableT]):
@overload
def __ne__(self, other: object) -> Literal[True]: ...

class IntervalTree(IntervalMixin):
def __init__(
self,
left: np.ndarray,
right: np.ndarray,
closed: IntervalClosedType = ...,
leaf_size: int = ...,
) -> None: ...
def get_indexer(self, target) -> npt.NDArray[np.intp]: ...
def get_indexer_non_unique(
self, target
) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]: ...
_na_count: int
@property
def is_overlapping(self) -> bool: ...
@property
def is_monotonic_increasing(self) -> bool: ...
def clear_mapping(self) -> None: ...
class IntervalTree(IntervalMixin): ...
22 changes: 17 additions & 5 deletions pandas-stubs/_libs/tslibs/offsets.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,14 @@ from typing import (

from dateutil.relativedelta import weekday as WeekdayClass
import numpy as np
from numpy import typing as npt
from pandas import Timestamp
from typing_extensions import Self

from pandas._typing import npt
from pandas._typing import (
ShapeT,
np_ndarray,
)

from pandas.tseries.holiday import AbstractHolidayCalendar

Expand All @@ -37,7 +41,9 @@ class BaseOffset:
@property
def base(self) -> BaseOffset: ...
@overload
def __add__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
def __add__(
self, other: np_ndarray[ShapeT, np.object_]
) -> np_ndarray[ShapeT, np.object_]: ...
@overload
def __add__(self, other: _DatetimeT) -> _DatetimeT: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
@overload
Expand All @@ -47,7 +53,9 @@ class BaseOffset:
@overload
def __add__(self, other: _TimedeltaT) -> _TimedeltaT: ...
@overload
def __radd__(self, other: npt.NDArray[np.object_]) -> npt.NDArray[np.object_]: ...
def __radd__(
self, other: np_ndarray[ShapeT, np.object_]
) -> np_ndarray[ShapeT, np.object_]: ...
@overload
def __radd__(self, other: _DatetimeT) -> _DatetimeT: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
@overload
Expand All @@ -68,11 +76,15 @@ class BaseOffset:
@overload
def __rsub__(self, other: _TimedeltaT) -> _TimedeltaT: ...
@overload
def __mul__(self, other: np.ndarray) -> np.ndarray: ...
def __mul__(
self, other: np_ndarray[ShapeT, np.object_]
) -> np_ndarray[ShapeT, np.object_]: ...
@overload
def __mul__(self, other: int) -> Self: ...
@overload
def __rmul__(self, other: np.ndarray) -> np.ndarray: ...
def __rmul__(
self, other: np_ndarray[ShapeT, np.object_]
) -> np_ndarray[ShapeT, np.object_]: ...
@overload
def __rmul__(self, other: int) -> Self: ...
def __neg__(self) -> Self: ...
Expand Down
6 changes: 3 additions & 3 deletions pandas-stubs/_testing/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ from typing import (
import warnings

from matplotlib.artist import Artist
import numpy as np
from numpy import typing as npt
from pandas import (
Categorical,
DataFrame,
Expand Down Expand Up @@ -61,7 +61,7 @@ def assert_attr_equal(
attr: str, left: object, right: object, obj: str = "Attributes"
) -> None: ...
def assert_is_valid_plot_return_object(
objs: Series | np.ndarray | Artist | tuple | dict,
objs: Series | npt.NDArray[Any] | Artist | tuple | dict,
) -> None: ...
def assert_is_sorted(seq: AnyArrayLike) -> None: ...
def assert_categorical_equal(
Expand Down Expand Up @@ -96,7 +96,7 @@ def assert_extension_array_equal(
left: ExtensionArray,
right: ExtensionArray,
check_dtype: bool | Literal["equiv"] = True,
index_values: Index | np.ndarray | None = None,
index_values: Index | npt.NDArray[Any] | None = None,
check_exact: bool = False,
rtol: float = 1e-5,
atol: float = 1e-8,
Expand Down
30 changes: 15 additions & 15 deletions pandas-stubs/_typing.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ HashableT5 = TypeVar("HashableT5", bound=Hashable)

# array-like

ArrayLike: TypeAlias = ExtensionArray | np.ndarray
ArrayLike: TypeAlias = ExtensionArray | npt.NDArray[Any]
AnyArrayLike: TypeAlias = ArrayLike | Index | Series
AnyArrayLikeInt: TypeAlias = (
IntegerArray | Index[int] | Series[int] | npt.NDArray[np.integer]
Expand Down Expand Up @@ -166,6 +166,7 @@ ToTimestampHow: TypeAlias = Literal["s", "e", "start", "end"]
NDFrameT = TypeVar("NDFrameT", bound=NDFrame)

IndexT = TypeVar("IndexT", bound=Index)
T_EXTENSION_ARRAY = TypeVar("T_EXTENSION_ARRAY", bound=ExtensionArray)

# From _typing.py, not used here:
# FreqIndexT = TypeVar("FreqIndexT", "DatetimeIndex", "PeriodIndex", "TimedeltaIndex")
Expand Down Expand Up @@ -697,11 +698,11 @@ InterpolateOptions: TypeAlias = Literal[
# Using List[int] here rather than Sequence[int] to disallow tuples.

ScalarIndexer: TypeAlias = int | np.integer
SequenceIndexer: TypeAlias = slice | list[int] | np.ndarray
SequenceIndexer: TypeAlias = slice | list[int] | npt.NDArray[np.integer | np.bool]
PositionalIndexer: TypeAlias = ScalarIndexer | SequenceIndexer
PositionalIndexerTuple: TypeAlias = tuple[PositionalIndexer, PositionalIndexer]
# PositionalIndexer2D = Union[PositionalIndexer, PositionalIndexerTuple] Not used in stubs
TakeIndexer: TypeAlias = Sequence[int] | Sequence[np.integer] | npt.NDArray[np.integer]
TakeIndexer: TypeAlias = Sequence[int | np.integer] | npt.NDArray[np.integer | np.bool]

# Shared by functions such as drop and astype
IgnoreRaise: TypeAlias = Literal["ignore", "raise"]
Expand Down Expand Up @@ -832,18 +833,6 @@ SliceType: TypeAlias = Hashable | None
## All types below this point are only used in pandas-stubs
######

num: TypeAlias = complex

DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic])
KeysArgType: TypeAlias = Any
ListLikeT = TypeVar("ListLikeT", bound=ListLike)
ListLikeExceptSeriesAndStr: TypeAlias = (
MutableSequence[Any] | np.ndarray | tuple[Any, ...] | Index
)
ListLikeU: TypeAlias = Sequence | np.ndarray | Series | Index
ListLikeHashable: TypeAlias = (
MutableSequence[HashableT] | np.ndarray | tuple[HashableT, ...] | range
)
StrLike: TypeAlias = str | np.str_

ScalarT = TypeVar("ScalarT", bound=Scalar)
Expand Down Expand Up @@ -871,6 +860,17 @@ np_ndarray: TypeAlias = np.ndarray[ShapeT, np.dtype[GenericT]]
np_1darray: TypeAlias = np.ndarray[tuple[int], np.dtype[GenericT]]
np_2darray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[GenericT]]

DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic])
KeysArgType: TypeAlias = Any
ListLikeT = TypeVar("ListLikeT", bound=ListLike)
ListLikeExceptSeriesAndStr: TypeAlias = (
MutableSequence[Any] | np_1darray[Any] | tuple[Any, ...] | Index
)
ListLikeU: TypeAlias = Sequence | np_1darray[Any] | Series | Index
ListLikeHashable: TypeAlias = (
MutableSequence[HashableT] | np_1darray[Any] | tuple[HashableT, ...] | range
)

class SupportsDType(Protocol[GenericT_co]):
@property
def dtype(self) -> np.dtype[GenericT_co]: ...
Expand Down
58 changes: 34 additions & 24 deletions pandas-stubs/core/algorithms.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -5,65 +5,75 @@ from typing import (
)

import numpy as np
from pandas import (
Categorical,
CategoricalIndex,
Index,
IntervalIndex,
PeriodIndex,
Series,
)
from numpy import typing as npt
from pandas.api.extensions import ExtensionArray
from pandas.core.arrays.categorical import Categorical
from pandas.core.indexes.base import Index
from pandas.core.indexes.category import CategoricalIndex
from pandas.core.indexes.datetimes import DatetimeIndex
from pandas.core.indexes.interval import IntervalIndex
from pandas.core.indexes.period import PeriodIndex
from pandas.core.series import Series

from pandas._typing import (
T_EXTENSION_ARRAY,
AnyArrayLike,
GenericT,
IntervalT,
TakeIndexer,
np_1darray,
np_ndarray,
)

# 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( # pyright: ignore[reportOverlappingOverload]
values: PeriodIndex,
) -> PeriodIndex: ...
@overload
def unique(values: CategoricalIndex) -> CategoricalIndex: ... # type: ignore[overload-overlap]
def unique(values: CategoricalIndex) -> CategoricalIndex: ...
@overload
def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ...
@overload
def unique(values: Index) -> np.ndarray: ...
def unique(values: PeriodIndex) -> PeriodIndex: ...
@overload
def unique(values: DatetimeIndex) -> np_1darray[np.datetime64] | DatetimeIndex: ...
@overload
def unique(values: Index) -> np_1darray[Any] | Index: ...
@overload
def unique(values: Categorical) -> Categorical: ...

# @overload
# def unique(values: Series[Never]) -> np_1darray[Any] | ExtensionArray: ...
# TODO: DatetimeArray python/mypy#19952
# @overload
# def unique(values: Series[Timestamp]) -> np_1darray[np.datetime64] | ExtensionArray: ...
# @overload
# def unique(values: Series[int]) -> np_1darray[np.integer] | ExtensionArray: ...
@overload
def unique(values: Series) -> np.ndarray | ExtensionArray: ...
def unique(values: Series) -> np_1darray[Any] | ExtensionArray: ...
@overload
def unique(values: np.ndarray) -> np.ndarray: ...
def unique(values: npt.NDArray[GenericT]) -> np_1darray[GenericT]: ...
@overload
def unique(values: ExtensionArray) -> ExtensionArray: ...
def unique(values: T_EXTENSION_ARRAY) -> T_EXTENSION_ARRAY: ...
@overload
def factorize(
values: np.ndarray,
values: npt.NDArray[GenericT],
sort: bool = ...,
use_na_sentinel: bool = ...,
size_hint: int | None = ...,
) -> tuple[np.ndarray, np.ndarray]: ...
) -> tuple[np_1darray[np.int64], np_1darray[GenericT]]: ...
@overload
def factorize(
values: Index | Series,
sort: bool = ...,
use_na_sentinel: bool = ...,
size_hint: int | None = ...,
) -> tuple[np_1darray, Index]: ...
) -> tuple[np_1darray[np.int64], Index]: ...
@overload
def factorize(
values: Categorical,
sort: bool = ...,
use_na_sentinel: bool = ...,
size_hint: int | None = ...,
) -> tuple[np_1darray, Categorical]: ...
) -> tuple[np_1darray[np.int64], Categorical]: ...
def value_counts(
values: AnyArrayLike | list | tuple,
sort: bool = True,
Expand All @@ -73,9 +83,9 @@ def value_counts(
dropna: bool = True,
) -> Series: ...
def take(
arr: np.ndarray | ExtensionArray | Index | Series,
arr: np_ndarray[Any] | ExtensionArray | Index | Series,
indices: TakeIndexer,
axis: Literal[0, 1] = 0,
allow_fill: bool = False,
fill_value: Any = None,
) -> np_1darray | ExtensionArray: ...
) -> np_1darray[Any] | ExtensionArray: ...
6 changes: 4 additions & 2 deletions pandas-stubs/core/arrays/base.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ class ExtensionArray:
def __getitem__(self, item: ScalarIndexer) -> Any: ...
@overload
def __getitem__(self, item: SequenceIndexer) -> Self: ...
def __setitem__(self, key: int | slice | np.ndarray, value: Any) -> None: ...
def __setitem__(self, key: int | slice | npt.NDArray[Any], value: Any) -> None: ...
def __len__(self) -> int: ...
def __iter__(self) -> Iterator[Any]: ...
def __contains__(self, item: object) -> bool | np.bool_: ...
Expand Down Expand Up @@ -83,7 +83,9 @@ class ExtensionArray:
side: Literal["left", "right"] = ...,
sorter: ListLike | None = ...,
) -> np.intp: ...
def factorize(self, use_na_sentinel: bool = True) -> tuple[np_1darray, Self]: ...
def factorize(
self, use_na_sentinel: bool = True
) -> tuple[np_1darray[Any], Self]: ...
def repeat(
self, repeats: int | AnyArrayLikeInt | Sequence[int], axis: None = None
) -> Self: ...
Expand Down
8 changes: 5 additions & 3 deletions pandas-stubs/core/arrays/boolean.pyi
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
from typing import Any

import numpy as np
from pandas.core.arrays.masked import BaseMaskedArray as BaseMaskedArray

from pandas._libs.missing import NAType
from pandas._typing import type_t
from pandas._typing import (
np_ndarray_bool,
type_t,
)

from pandas.core.dtypes.base import ExtensionDtype as ExtensionDtype

Expand All @@ -16,7 +18,7 @@ class BooleanDtype(ExtensionDtype):

class BooleanArray(BaseMaskedArray):
def __init__(
self, values: np.ndarray, mask: np.ndarray, copy: bool = ...
self, values: np_ndarray_bool, mask: np_ndarray_bool, copy: bool = ...
) -> None: ...
@property
def dtype(self): ...
Expand Down
3 changes: 1 addition & 2 deletions pandas-stubs/core/arrays/datetimelike.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ from typing import (
overload,
)

import numpy as np
from pandas.core.arrays.base import (
ExtensionArray,
ExtensionOpsMixin,
Expand Down Expand Up @@ -65,7 +64,7 @@ class DatetimeLikeArrayMixin(ExtensionOpsMixin, ExtensionArray):
def ravel(self, *args: Any, **kwargs: Any): ...
def __iter__(self): ...
@property
def asi8(self) -> np.ndarray: ...
def asi8(self) -> np_1darray[Any]: ...
@property
def nbytes(self): ...
def __array__(
Expand Down
Loading