diff --git a/pandas/core/arrays/sparse/array.py b/pandas/core/arrays/sparse/array.py index 6ae216cd3263c..c0c13942ee850 100644 --- a/pandas/core/arrays/sparse/array.py +++ b/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'}. """ @@ -1630,7 +1632,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 @@ -1689,8 +1694,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)