Allow sorted() to work on Index (#501)

* trying SupportsRichComparisonT

* create IndexIterScalar type for sorted(Index) to work

* Use List not list, and Tuple not tuple for 3.8 compat

Author: Dr-Irv

pandas-stubs/_typing.pyi

@@ -165,7 +165,7 @@ ListLikeExceptSeriesAndStr = TypeVar(
 )
 ListLikeU: TypeAlias = Union[Sequence, np.ndarray, Series, Index]
 StrLike: TypeAlias = Union[str, np.str_]
-Scalar: TypeAlias = Union[
+IndexIterScalar: TypeAlias = Union[
     str,
     bytes,
     datetime.date,
@@ -174,10 +174,13 @@ Scalar: TypeAlias = Union[
     bool,
     int,
     float,
-    complex,
     Timestamp,
     Timedelta,
 ]
+Scalar: TypeAlias = Union[
+    IndexIterScalar,
+    complex,
+]
 ScalarT = TypeVar("ScalarT", bound=Scalar)
 # Refine the definitions below in 3.9 to use the specialized type.
 np_ndarray_int64: TypeAlias = npt.NDArray[np.int64]

pandas-stubs/core/indexes/base.pyi

@@ -32,6 +32,7 @@ from pandas._typing import (
     DtypeObj,
     FillnaOptions,
     HashableT,
+    IndexIterScalar,
     IndexT,
     Label,
     Level,
@@ -222,7 +223,7 @@ class Index(IndexOpsMixin, PandasObject):
     def shape(self) -> tuple[int, ...]: ...
     # Extra methods from old stubs
     def __eq__(self, other: object) -> np_ndarray_bool: ...  # type: ignore[override]
-    def __iter__(self) -> Iterator[Scalar | tuple[Hashable, ...]]: ...
+    def __iter__(self) -> Iterator[IndexIterScalar | tuple[Hashable, ...]]: ...
     def __ne__(self, other: object) -> np_ndarray_bool: ...  # type: ignore[override]
     def __le__(self, other: Index | Scalar) -> np_ndarray_bool: ...
     def __ge__(self, other: Index | Scalar) -> np_ndarray_bool: ...

tests/test_indexes.py

@@ -2,6 +2,7 @@
 
 import datetime as dt
 from typing import (
+    TYPE_CHECKING,
     Hashable,
     List,
     Tuple,
@@ -16,6 +17,9 @@
 
 from pandas._typing import Scalar
 
+if TYPE_CHECKING:
+    from pandas._typing import IndexIterScalar
+
 from tests import (
     check,
     pytest_warns_bounded,
@@ -92,7 +96,11 @@ def test_column_contains() -> None:
 def test_column_sequence() -> None:
     df = pd.DataFrame([1, 2, 3])
     col_list = list(df.columns)
-    check(assert_type(col_list, List[Union[Scalar, Tuple[Hashable, ...]]]), list, int)
+    check(
+        assert_type(col_list, List[Union["IndexIterScalar", Tuple[Hashable, ...]]]),
+        list,
+        int,
+    )
 
 
 def test_difference_none() -> None:
@@ -657,3 +665,22 @@ def test_interval_index_tuples():
         pd.IntervalIndex,
         pd.Interval,
     )
+
+
+def test_sorted_and_list() -> None:
+    # GH 497
+    i1 = pd.Index([3, 2, 1])
+    check(
+        assert_type(
+            sorted(i1),
+            List[Union["IndexIterScalar", Tuple[Hashable, ...]]],
+        ),
+        list,
+    )
+    check(
+        assert_type(
+            list(i1),
+            List[Union["IndexIterScalar", Tuple[Hashable, ...]]],
+        ),
+        list,
+    )