TYP: SparseIndexKind (pandas-dev#43536)

Author: mzeitlin11

pandas/core/arrays/sparse/array.py

@@ -97,6 +97,8 @@ class ellipsis(Enum):
 
     Ellipsis = ellipsis.Ellipsis
 
+    SparseIndexKind = Literal["integer", "block"]
+
     from pandas._typing import NumpySorter
 
     from pandas import Series
@@ -334,7 +336,7 @@ def __init__(
         sparse_index=None,
         index=None,
         fill_value=None,
-        kind="integer",
+        kind: SparseIndexKind = "integer",
         dtype: Dtype | None = None,
         copy: bool = False,
     ):
@@ -624,7 +626,7 @@ def fill_value(self, value):
         self._dtype = SparseDtype(self.dtype.subtype, value)
 
     @property
-    def kind(self) -> str:
+    def kind(self) -> SparseIndexKind:
         """
         The kind of sparse index for this array. One of {'integer', 'block'}.
         """
@@ -1676,7 +1678,10 @@ def _formatter(self, boxed=False):
 
 
 def make_sparse(
-    arr: np.ndarray, kind="block", fill_value=None, dtype: NpDtype | None = None
+    arr: np.ndarray,
+    kind: SparseIndexKind = "block",
+    fill_value=None,
+    dtype: NpDtype | None = None,
 ):
     """
     Convert ndarray to sparse format
@@ -1735,8 +1740,17 @@ def make_sparse(
     return sparsified_values, index, fill_value
 
 
-def make_sparse_index(length, indices, kind) -> SparseIndex:
+@overload
+def make_sparse_index(length: int, indices, kind: Literal["block"]) -> BlockIndex:
+    ...
+
+
+@overload
+def make_sparse_index(length: int, indices, kind: Literal["integer"]) -> IntIndex:
+    ...
+
 
+def make_sparse_index(length: int, indices, kind: SparseIndexKind) -> SparseIndex:
     index: SparseIndex
     if kind == "block" or isinstance(kind, BlockIndex):
         locs, lens = splib.get_blocks(indices)