Skip to content

Commit 12d8d84

Browse files
cmp0xffcmp0xff
committed
resolve unique
1 parent a6faa39 commit 12d8d84

File tree

15 files changed

+85
-63
lines changed

15 files changed

+85
-63
lines changed

‎pandas-stubs/_typing.pyi‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -938,6 +938,8 @@ C2 = TypeVar(
938938
BaseOffset,
939939
)
940940

941+
A1_co = TypeVar("A1_co", default=Any, covariant=True)
942+
941943
IndexingInt: TypeAlias = (
942944
int | np.int_ | np.integer | np.unsignedinteger | np.signedinteger | np.int8
943945
)

‎pandas-stubs/core/algorithms.pyi‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,18 +5,19 @@ from typing import (
55
)
66

77
import numpy as np
8-
from pandas import (
9-
Categorical,
10-
CategoricalIndex,
11-
Index,
12-
IntervalIndex,
13-
PeriodIndex,
14-
Series,
15-
)
8+
from numpy import typing as npt
169
from pandas.api.extensions import ExtensionArray
10+
from pandas.core.arrays.categorical import Categorical
11+
from pandas.core.indexes.base import Index
12+
from pandas.core.indexes.category import CategoricalIndex
13+
from pandas.core.indexes.interval import IntervalIndex
14+
from pandas.core.indexes.period import PeriodIndex
15+
from pandas.core.series import Series
1716

1817
from pandas._typing import (
18+
S1,
1919
AnyArrayLike,
20+
GenericT_co,
2021
IntervalT,
2122
TakeIndexer,
2223
np_1darray,
@@ -26,21 +27,21 @@ from pandas._typing import (
2627
# with extension types return the same type while standard type return ndarray
2728

2829
@overload
29-
def unique( # pyright: ignore[reportOverlappingOverload]
30-
values: PeriodIndex,
31-
) -> PeriodIndex: ...
30+
def unique(values: PeriodIndex) -> PeriodIndex: ...
3231
@overload
33-
def unique(values: CategoricalIndex) -> CategoricalIndex: ... # type: ignore[overload-overlap]
32+
def unique(
33+
values: CategoricalIndex[S1, GenericT_co],
34+
) -> CategoricalIndex[S1, GenericT_co]: ...
3435
@overload
3536
def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ...
3637
@overload
37-
def unique(values: Index) -> np.ndarray: ...
38+
def unique(values: Index[S1, np_1darray, GenericT_co]) -> np_1darray[GenericT_co]: ...
3839
@overload
3940
def unique(values: Categorical) -> Categorical: ...
4041
@overload
4142
def unique(values: Series) -> np.ndarray | ExtensionArray: ...
4243
@overload
43-
def unique(values: np.ndarray) -> np.ndarray: ...
44+
def unique(values: npt.NDArray[GenericT_co]) -> np_1darray[GenericT_co]: ...
4445
@overload
4546
def unique(values: ExtensionArray) -> ExtensionArray: ...
4647
@overload

‎pandas-stubs/core/base.pyi‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ from pandas._libs.tslibs.timedeltas import Timedelta
2929
from pandas._typing import (
3030
S1,
3131
S2,
32+
A1_co,
3233
AxisIndex,
3334
DropKeep,
3435
DTypeLike,
@@ -59,7 +60,7 @@ class SelectionMixin(Generic[NDFrameT]):
5960
def __getitem__(self, key): ...
6061
def aggregate(self, func, *args: Any, **kwargs: Any): ...
6162

62-
class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
63+
class IndexOpsMixin(OpsMixin, Generic[S1, A1_co, GenericT_co]):
6364
__array_priority__: int = ...
6465
@property
6566
def T(self) -> Self: ...

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

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ from pathlib import Path
1414
from typing import (
1515
Any,
1616
ClassVar,
17-
Generic,
1817
Literal,
1918
final,
2019
overload,
@@ -67,6 +66,7 @@ from pandas._typing import (
6766
S2,
6867
S2_NSDT,
6968
T_COMPLEX,
69+
A1_co,
7070
AnyAll,
7171
AnyArrayLike,
7272
AnyArrayLikeInt,
@@ -109,7 +109,7 @@ from pandas._typing import (
109109

110110
class InvalidIndexError(Exception): ...
111111

112-
class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
112+
class Index(IndexOpsMixin[S1, A1_co, GenericT_co], ElementOpsMixin[S1]):
113113
__hash__: ClassVar[None] # type: ignore[assignment] # pyright: ignore[reportIncompatibleMethodOverride]
114114
# overloads with additional dtypes
115115
@overload
@@ -475,15 +475,19 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
475475
@overload
476476
def where(
477477
self,
478-
cond: Sequence[bool] | np_ndarray_bool | BooleanArray | IndexOpsMixin[bool],
479-
other: S1 | Series[S1] | Self,
478+
cond: (
479+
Sequence[bool] | np_ndarray_bool | BooleanArray | IndexOpsMixin[bool, A1_co]
480+
),
481+
other: S1 | Series[S1, A1_co] | Self,
480482
) -> Self: ...
481483
@overload
482484
def where(
483485
self,
484-
cond: Sequence[bool] | np_ndarray_bool | BooleanArray | IndexOpsMixin[bool],
486+
cond: (
487+
Sequence[bool] | np_ndarray_bool | BooleanArray | IndexOpsMixin[bool, A1_co]
488+
),
485489
other: Scalar | AnyArrayLike | None = None,
486-
) -> Index: ...
490+
) -> Index[Any, A1_co]: ...
487491
def __contains__(self, key: Hashable) -> bool: ...
488492
@final
489493
def __setitem__(self, key, value) -> None: ...
@@ -541,7 +545,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
541545
@overload
542546
def insert(self, loc: int, item: S1) -> Self: ...
543547
@overload
544-
def insert(self, loc: int, item: object) -> Index: ...
548+
def insert(self, loc: int, item: Any) -> Index[Any, A1_co]: ...
545549
def drop(self, labels, errors: IgnoreRaise = "raise") -> Self: ...
546550
@property
547551
def shape(self) -> tuple[int, ...]: ...
@@ -731,7 +735,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
731735
),
732736
) -> Index[complex]: ...
733737
@overload
734-
def __rsub__(self: Index[Never], other: DatetimeIndex) -> Never: ... # type: ignore[misc]
738+
def __rsub__(self: Index[Never], other: DatetimeIndex) -> Never: ...
735739
@overload
736740
def __rsub__(
737741
self: Index[Never], other: complex | ArrayLike | SequenceNotStr[S1] | Index
@@ -1051,7 +1055,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
10511055
def infer_objects(self, copy: bool = True) -> Self: ...
10521056

10531057
@type_check_only
1054-
class _IndexSubclassBase(Index[S1], Generic[S1, GenericT_co]):
1058+
class _IndexSubclassBase(Index[S1, A1_co, GenericT_co]):
10551059
@overload
10561060
def to_numpy(
10571061
self,

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

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ from collections.abc import (
33
Iterable,
44
)
55

6-
import numpy as np
76
from pandas.core import accessor
87
from pandas.core.arrays.categorical import Categorical
98
from pandas.core.indexes.base import Index
@@ -13,12 +12,16 @@ from typing_extensions import Self
1312
from pandas._typing import (
1413
S1,
1514
Dtype,
15+
GenericT_co,
1616
ListLike,
17+
np_1darray,
1718
)
1819

19-
class CategoricalIndex(ExtensionIndex[S1], accessor.PandasDelegate):
20-
codes: np.ndarray = ...
21-
categories: Index = ...
20+
class CategoricalIndex(
21+
ExtensionIndex[S1, Categorical, GenericT_co], accessor.PandasDelegate
22+
):
23+
codes: np_1darray[GenericT_co] = ...
24+
categories: Index[S1] = ...
2225
@property
2326
def array(self) -> Categorical: ... # type: ignore[override] # pyrefly: ignore[bad-override]
2427
def __new__(

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,14 @@ from typing_extensions import (
1010
from pandas._libs.tslibs import BaseOffset
1111
from pandas._typing import (
1212
S1,
13+
A1_co,
1314
AxisIndex,
1415
GenericT_co,
1516
TimeUnit,
1617
np_ndarray_complex,
1718
)
1819

19-
class DatetimeIndexOpsMixin(ExtensionIndex[S1, GenericT_co]):
20+
class DatetimeIndexOpsMixin(ExtensionIndex[S1, A1_co, GenericT_co]):
2021
@property
2122
def freq(self) -> BaseOffset | None: ...
2223
@property
@@ -58,7 +59,7 @@ class DatetimeIndexOpsMixin(ExtensionIndex[S1, GenericT_co]):
5859
self, other: np_ndarray_complex
5960
) -> Never: ...
6061

61-
class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1, GenericT_co]):
62+
class DatetimeTimedeltaMixin(DatetimeIndexOpsMixin[S1, A1_co, GenericT_co]):
6263
@property
6364
def unit(self) -> TimeUnit: ...
6465
def as_unit(self, unit: TimeUnit) -> Self: ...

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ from pandas.core.series import Series
2626
from typing_extensions import Self
2727

2828
from pandas._typing import (
29+
A1_co,
2930
AxesData,
3031
DateAndDatetimeLike,
3132
Dtype,
@@ -43,7 +44,7 @@ from pandas.core.dtypes.dtypes import DatetimeTZDtype
4344
from pandas.tseries.offsets import BaseOffset
4445

4546
class DatetimeIndex(
46-
DatetimeTimedeltaMixin[Timestamp, np.datetime64], DatetimeIndexProperties
47+
DatetimeTimedeltaMixin[Timestamp, A1_co, np.datetime64], DatetimeIndexProperties
4748
):
4849
def __new__(
4950
cls,

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@ from pandas.core.indexes.base import _IndexSubclassBase
22

33
from pandas._typing import (
44
S1,
5+
A1_co,
56
GenericT_co,
67
)
78

8-
class ExtensionIndex(_IndexSubclassBase[S1, GenericT_co]): ...
9+
class ExtensionIndex(_IndexSubclassBase[S1, A1_co, GenericT_co]): ...

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ from typing import (
1111

1212
import numpy as np
1313
import pandas as pd
14-
from pandas import Index
14+
from pandas.core.indexes.base import Index
1515
from pandas.core.indexes.extension import ExtensionIndex
1616

1717
from pandas._libs.interval import (
@@ -60,7 +60,9 @@ _EdgesTimedelta: TypeAlias = (
6060
_TimestampLike: TypeAlias = pd.Timestamp | np.datetime64 | dt.datetime
6161
_TimedeltaLike: TypeAlias = pd.Timedelta | np.timedelta64 | dt.timedelta
6262

63-
class IntervalIndex(ExtensionIndex[IntervalT, np.object_], IntervalMixin):
63+
class IntervalIndex(
64+
ExtensionIndex[IntervalT, IntervalIndex, np.object_], IntervalMixin
65+
):
6466
closed: IntervalClosedType
6567

6668
def __new__(

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,9 @@ from pandas._typing import (
2525
np_1darray,
2626
)
2727

28-
class PeriodIndex(DatetimeIndexOpsMixin[pd.Period, np.object_], PeriodIndexFieldOps):
28+
class PeriodIndex(
29+
DatetimeIndexOpsMixin[pd.Period, np_1darray, np.object_], PeriodIndexFieldOps
30+
):
2931
def __new__(
3032
cls,
3133
data: AxesData[Any] | None = None,

0 commit comments

Comments
 (0)