@@ -97,9 +97,14 @@ class ellipsis(Enum):
97
97
98
98
Ellipsis = ellipsis .Ellipsis
99
99
100
- SparseIndexKind = Literal ["integer" , "block" ]
100
+ from scipy .sparse import spmatrix
101
+
102
+ from pandas ._typing import (
103
+ FillnaOptions ,
104
+ NumpySorter ,
105
+ )
101
106
102
- from pandas . _typing import NumpySorter
107
+ SparseIndexKind = Literal [ "integer" , "block" ]
103
108
104
109
from pandas import Series
105
110
@@ -141,7 +146,7 @@ def _get_fill(arr: SparseArray) -> np.ndarray:
141
146
142
147
def _sparse_array_op (
143
148
left : SparseArray , right : SparseArray , op : Callable , name : str
144
- ) -> Any :
149
+ ) -> SparseArray :
145
150
"""
146
151
Perform a binary operation between two arrays.
147
152
@@ -229,7 +234,9 @@ def _sparse_array_op(
229
234
return _wrap_result (name , result , index , fill , dtype = result_dtype )
230
235
231
236
232
- def _wrap_result (name , data , sparse_index , fill_value , dtype : Dtype | None = None ):
237
+ def _wrap_result (
238
+ name : str , data , sparse_index , fill_value , dtype : Dtype | None = None
239
+ ) -> SparseArray :
233
240
"""
234
241
wrap op result to have correct dtype
235
242
"""
@@ -500,7 +507,7 @@ def _simple_new(
500
507
return new
501
508
502
509
@classmethod
503
- def from_spmatrix (cls , data ) :
510
+ def from_spmatrix (cls : type [ SparseArrayT ] , data : spmatrix ) -> SparseArrayT :
504
511
"""
505
512
Create a SparseArray from a scipy.sparse matrix.
506
513
@@ -690,7 +697,12 @@ def isna(self):
690
697
dtype = SparseDtype (bool , self ._null_fill_value )
691
698
return type (self )._simple_new (isna (self .sp_values ), self .sp_index , dtype )
692
699
693
- def fillna (self , value = None , method = None , limit = None ):
700
+ def fillna (
701
+ self : SparseArrayT ,
702
+ value = None ,
703
+ method : FillnaOptions | None = None ,
704
+ limit : int | None = None ,
705
+ ) -> SparseArrayT :
694
706
"""
695
707
Fill missing values with `value`.
696
708
@@ -745,7 +757,7 @@ def fillna(self, value=None, method=None, limit=None):
745
757
746
758
return self ._simple_new (new_values , self ._sparse_index , new_dtype )
747
759
748
- def shift (self , periods : int = 1 , fill_value = None ):
760
+ def shift (self : SparseArrayT , periods : int = 1 , fill_value = None ) -> SparseArrayT :
749
761
750
762
if not len (self ) or periods == 0 :
751
763
return self .copy ()
@@ -791,7 +803,7 @@ def _first_fill_value_loc(self):
791
803
diff = indices [1 :] - indices [:- 1 ]
792
804
return np .searchsorted (diff , 2 ) + 1
793
805
794
- def unique (self ) :
806
+ def unique (self : SparseArrayT ) -> SparseArrayT :
795
807
uniques = list (algos .unique (self .sp_values ))
796
808
fill_loc = self ._first_fill_value_loc ()
797
809
if fill_loc >= 0 :
@@ -802,17 +814,15 @@ def _values_for_factorize(self):
802
814
# Still override this for hash_pandas_object
803
815
return np .asarray (self ), self .fill_value
804
816
805
- def factorize (self , na_sentinel = - 1 ):
817
+ def factorize (self , na_sentinel : int = - 1 ) -> tuple [ np . ndarray , SparseArray ] :
806
818
# Currently, ExtensionArray.factorize -> Tuple[ndarray, EA]
807
819
# The sparsity on this is backwards from what Sparse would want. Want
808
820
# ExtensionArray.factorize -> Tuple[EA, EA]
809
821
# Given that we have to return a dense array of codes, why bother
810
822
# implementing an efficient factorize?
811
823
codes , uniques = algos .factorize (np .asarray (self ), na_sentinel = na_sentinel )
812
- # error: Incompatible types in assignment (expression has type "SparseArray",
813
- # variable has type "Union[ndarray, Index]")
814
- uniques = SparseArray (uniques , dtype = self .dtype ) # type: ignore[assignment]
815
- return codes , uniques
824
+ uniques_sp = SparseArray (uniques , dtype = self .dtype )
825
+ return codes , uniques_sp
816
826
817
827
def value_counts (self , dropna : bool = True ) -> Series :
818
828
"""
@@ -930,8 +940,8 @@ def _get_val_at(self, loc):
930
940
return val
931
941
932
942
def take (
933
- self , indices , * , allow_fill : bool = False , fill_value = None
934
- ) -> SparseArray :
943
+ self : SparseArrayT , indices , * , allow_fill : bool = False , fill_value = None
944
+ ) -> SparseArrayT :
935
945
if is_scalar (indices ):
936
946
raise ValueError (f"'indices' must be an array, not a scalar '{ indices } '." )
937
947
indices = np .asarray (indices , dtype = np .int32 )
@@ -1222,7 +1232,7 @@ def astype(self, dtype: AstypeArg | None = None, copy: bool = True):
1222
1232
sp_values , self .sp_index , dtype # type: ignore[arg-type]
1223
1233
)
1224
1234
1225
- def map (self , mapper ):
1235
+ def map (self : SparseArrayT , mapper ) -> SparseArrayT :
1226
1236
"""
1227
1237
Map categories using input correspondence (dict, Series, or function).
1228
1238
@@ -1274,7 +1284,7 @@ def map(self, mapper):
1274
1284
1275
1285
return type (self )(sp_values , sparse_index = self .sp_index , fill_value = fill_value )
1276
1286
1277
- def to_dense (self ):
1287
+ def to_dense (self ) -> np . ndarray :
1278
1288
"""
1279
1289
Convert SparseArray to a NumPy array.
1280
1290
@@ -1407,7 +1417,7 @@ def sum(self, axis: int = 0, min_count: int = 0, *args, **kwargs) -> Scalar:
1407
1417
return na_value_for_dtype (self .dtype .subtype , compat = False )
1408
1418
return sp_sum + self .fill_value * nsparse
1409
1419
1410
- def cumsum (self , axis = 0 , * args , ** kwargs ):
1420
+ def cumsum (self , axis : int = 0 , * args , ** kwargs ) -> SparseArray :
1411
1421
"""
1412
1422
Cumulative sum of non-NA/null values.
1413
1423
0 commit comments