pyright happiness

Author: cmp0xff

pandas-stubs/core/algorithms.pyi

@@ -10,8 +10,10 @@ 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.indexes.timedeltas import TimedeltaIndex
 from pandas.core.series import Series
 
 from pandas._typing import (
@@ -37,7 +39,11 @@ def unique(
 @overload
 def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ...
 @overload
-def unique(values: Index[S1, np_1darray, GenericT_co]) -> np_1darray[GenericT_co]: ...
+def unique(values: DatetimeIndex) -> DatetimeIndex | np_1darray[np.datetime64]: ...
+@overload
+def unique(values: TimedeltaIndex) -> TimedeltaIndex | np_1darray[np.timedelta64]: ...
+@overload
+def unique(values: Index[S1, Any, GenericT_co]) -> np_1darray[GenericT_co]: ...
 @overload
 def unique(values: Categorical) -> Categorical: ...
 @overload

pandas-stubs/core/indexes/accessors.pyi

@@ -5,6 +5,7 @@ from datetime import (
     tzinfo as _tzinfo,
 )
 from typing import (
+    Any,
     Generic,
     Literal,
     TypeVar,
@@ -486,5 +487,5 @@ class ArrayDescriptor:
     ) -> TimedeltaArray: ...
     @overload
     def __get__(
-        self, instance: IndexOpsMixin[int, Never], owner: type[IndexOpsMixin]
+        self, instance: IndexOpsMixin[Any, Never], owner: type[IndexOpsMixin]
     ) -> NumpyExtensionArray: ...

tests/test_pandas.py

@@ -609,7 +609,7 @@ def test_unique() -> None:
                     ]
                 )
             ),
-            np_1darray[np.datetime64],
+            pd.DatetimeIndex | np_1darray[np.datetime64],
         ),
         pd.DatetimeIndex,
     )
@@ -663,9 +663,9 @@ def test_unique() -> None:
     check(
         assert_type(
             pd.unique(pd.timedelta_range(start="1 day", periods=4)),
-            np_1darray[np.timedelta64],
+            pd.TimedeltaIndex | np_1darray[np.timedelta64],
         ),
-        np_1darray if PD_LTE_23 else pd.Index,
+        np_1darray if PD_LTE_23 else pd.TimedeltaIndex,
         np.timedelta64 if PD_LTE_23 else pd.Timedelta,
     )
 
@@ -987,8 +987,8 @@ def test_index_unqiue() -> None:
 
     check(assert_type(pd.unique(ci), "pd.CategoricalIndex[str]"), pd.CategoricalIndex)
     check(
-        assert_type(pd.unique(dti), np_1darray[np.datetime64]),
-        np.ndarray if PD_LTE_23 else pd.Index,
+        assert_type(pd.unique(dti), pd.DatetimeIndex | np_1darray[np.datetime64]),
+        np.ndarray if PD_LTE_23 else pd.DatetimeIndex,
         np.datetime64 if PD_LTE_23 else pd.Timestamp,
     )
     check(assert_type(pd.unique(i), np_1darray), np.ndarray if PD_LTE_23 else pd.Index)
@@ -998,8 +998,9 @@ def test_index_unqiue() -> None:
         np.ndarray if PD_LTE_23 else pd.Index,
     )
     check(
-        assert_type(pd.unique(tdi), np_1darray[np.timedelta64]),
-        np.ndarray if PD_LTE_23 else pd.Index,
+        assert_type(pd.unique(tdi), pd.TimedeltaIndex | np_1darray[np.timedelta64]),
+        np_1darray if PD_LTE_23 else pd.TimedeltaIndex,
+        np.timedelta64 if PD_LTE_23 else pd.Timedelta,
     )
     check(assert_type(pd.unique(mi), np_1darray), np.ndarray if PD_LTE_23 else pd.Index)
     check(