Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1""" 

2Shared methods for Index subclasses backed by ExtensionArray. 

3""" 

4from typing import List 

5 

6import numpy as np 

7 

8from pandas.compat.numpy import function as nv 

9from pandas.util._decorators import Appender, cache_readonly 

10 

11from pandas.core.dtypes.common import ( 

12 ensure_platform_int, 

13 is_dtype_equal, 

14 is_object_dtype, 

15) 

16from pandas.core.dtypes.generic import ABCSeries 

17 

18from pandas.core.arrays import ExtensionArray 

19from pandas.core.indexers import deprecate_ndim_indexing 

20from pandas.core.indexes.base import Index 

21from pandas.core.ops import get_op_result_name 

22 

23 

24def inherit_from_data(name: str, delegate, cache: bool = False, wrap: bool = False): 

25 """ 

26 Make an alias for a method of the underlying ExtensionArray. 

27 

28 Parameters 

29 ---------- 

30 name : str 

31 Name of an attribute the class should inherit from its EA parent. 

32 delegate : class 

33 cache : bool, default False 

34 Whether to convert wrapped properties into cache_readonly 

35 wrap : bool, default False 

36 Whether to wrap the inherited result in an Index. 

37 

38 Returns 

39 ------- 

40 attribute, method, property, or cache_readonly 

41 """ 

42 

43 attr = getattr(delegate, name) 

44 

45 if isinstance(attr, property): 

46 if cache: 

47 

48 def cached(self): 

49 return getattr(self._data, name) 

50 

51 cached.__name__ = name 

52 cached.__doc__ = attr.__doc__ 

53 method = cache_readonly(cached) 

54 

55 else: 

56 

57 def fget(self): 

58 result = getattr(self._data, name) 

59 if wrap: 

60 if isinstance(result, type(self._data)): 

61 return type(self)._simple_new(result, name=self.name) 

62 return Index(result, name=self.name) 

63 return result 

64 

65 def fset(self, value): 

66 setattr(self._data, name, value) 

67 

68 fget.__name__ = name 

69 fget.__doc__ = attr.__doc__ 

70 

71 method = property(fget, fset) 

72 

73 elif not callable(attr): 

74 # just a normal attribute, no wrapping 

75 method = attr 

76 

77 else: 

78 

79 def method(self, *args, **kwargs): 

80 result = attr(self._data, *args, **kwargs) 

81 if wrap: 

82 if isinstance(result, type(self._data)): 

83 return type(self)._simple_new(result, name=self.name) 

84 return Index(result, name=self.name) 

85 return result 

86 

87 method.__name__ = name 

88 method.__doc__ = attr.__doc__ 

89 return method 

90 

91 

92def inherit_names(names: List[str], delegate, cache: bool = False, wrap: bool = False): 

93 """ 

94 Class decorator to pin attributes from an ExtensionArray to a Index subclass. 

95 

96 Parameters 

97 ---------- 

98 names : List[str] 

99 delegate : class 

100 cache : bool, default False 

101 wrap : bool, default False 

102 Whether to wrap the inherited result in an Index. 

103 """ 

104 

105 def wrapper(cls): 

106 for name in names: 

107 meth = inherit_from_data(name, delegate, cache=cache, wrap=wrap) 

108 setattr(cls, name, meth) 

109 

110 return cls 

111 

112 return wrapper 

113 

114 

115def _make_wrapped_comparison_op(opname): 

116 """ 

117 Create a comparison method that dispatches to ``._data``. 

118 """ 

119 

120 def wrapper(self, other): 

121 if isinstance(other, ABCSeries): 

122 # the arrays defer to Series for comparison ops but the indexes 

123 # don't, so we have to unwrap here. 

124 other = other._values 

125 

126 other = _maybe_unwrap_index(other) 

127 

128 op = getattr(self._data, opname) 

129 return op(other) 

130 

131 wrapper.__name__ = opname 

132 return wrapper 

133 

134 

135def make_wrapped_arith_op(opname): 

136 def method(self, other): 

137 if ( 

138 isinstance(other, Index) 

139 and is_object_dtype(other.dtype) 

140 and type(other) is not Index 

141 ): 

142 # We return NotImplemented for object-dtype index *subclasses* so they have 

143 # a chance to implement ops before we unwrap them. 

144 # See https://github.com/pandas-dev/pandas/issues/31109 

145 return NotImplemented 

146 meth = getattr(self._data, opname) 

147 result = meth(_maybe_unwrap_index(other)) 

148 return _wrap_arithmetic_op(self, other, result) 

149 

150 method.__name__ = opname 

151 return method 

152 

153 

154def _wrap_arithmetic_op(self, other, result): 

155 if result is NotImplemented: 

156 return NotImplemented 

157 

158 if isinstance(result, tuple): 

159 # divmod, rdivmod 

160 assert len(result) == 2 

161 return ( 

162 _wrap_arithmetic_op(self, other, result[0]), 

163 _wrap_arithmetic_op(self, other, result[1]), 

164 ) 

165 

166 if not isinstance(result, Index): 

167 # Index.__new__ will choose appropriate subclass for dtype 

168 result = Index(result) 

169 

170 res_name = get_op_result_name(self, other) 

171 result.name = res_name 

172 return result 

173 

174 

175def _maybe_unwrap_index(obj): 

176 """ 

177 If operating against another Index object, we need to unwrap the underlying 

178 data before deferring to the DatetimeArray/TimedeltaArray/PeriodArray 

179 implementation, otherwise we will incorrectly return NotImplemented. 

180 

181 Parameters 

182 ---------- 

183 obj : object 

184 

185 Returns 

186 ------- 

187 unwrapped object 

188 """ 

189 if isinstance(obj, Index): 

190 return obj._data 

191 return obj 

192 

193 

194class ExtensionIndex(Index): 

195 """ 

196 Index subclass for indexes backed by ExtensionArray. 

197 """ 

198 

199 _data: ExtensionArray 

200 

201 __eq__ = _make_wrapped_comparison_op("__eq__") 

202 __ne__ = _make_wrapped_comparison_op("__ne__") 

203 __lt__ = _make_wrapped_comparison_op("__lt__") 

204 __gt__ = _make_wrapped_comparison_op("__gt__") 

205 __le__ = _make_wrapped_comparison_op("__le__") 

206 __ge__ = _make_wrapped_comparison_op("__ge__") 

207 

208 def __getitem__(self, key): 

209 result = self._data[key] 

210 if isinstance(result, type(self._data)): 

211 return type(self)(result, name=self.name) 

212 

213 # Includes cases where we get a 2D ndarray back for MPL compat 

214 deprecate_ndim_indexing(result) 

215 return result 

216 

217 def __iter__(self): 

218 return self._data.__iter__() 

219 

220 @property 

221 def _ndarray_values(self) -> np.ndarray: 

222 return self._data._ndarray_values 

223 

224 @Appender(Index.dropna.__doc__) 

225 def dropna(self, how="any"): 

226 if how not in ("any", "all"): 

227 raise ValueError(f"invalid how option: {how}") 

228 

229 if self.hasnans: 

230 return self._shallow_copy(self._data[~self._isnan]) 

231 return self._shallow_copy() 

232 

233 def repeat(self, repeats, axis=None): 

234 nv.validate_repeat(tuple(), dict(axis=axis)) 

235 result = self._data.repeat(repeats, axis=axis) 

236 return self._shallow_copy(result) 

237 

238 @Appender(Index.take.__doc__) 

239 def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs): 

240 nv.validate_take(tuple(), kwargs) 

241 indices = ensure_platform_int(indices) 

242 

243 taken = self._assert_take_fillable( 

244 self._data, 

245 indices, 

246 allow_fill=allow_fill, 

247 fill_value=fill_value, 

248 na_value=self._na_value, 

249 ) 

250 return type(self)(taken, name=self.name) 

251 

252 def unique(self, level=None): 

253 if level is not None: 

254 self._validate_index_level(level) 

255 

256 result = self._data.unique() 

257 return self._shallow_copy(result) 

258 

259 def _get_unique_index(self, dropna=False): 

260 if self.is_unique and not dropna: 

261 return self 

262 

263 result = self._data.unique() 

264 if dropna and self.hasnans: 

265 result = result[~result.isna()] 

266 return self._shallow_copy(result) 

267 

268 @Appender(Index.map.__doc__) 

269 def map(self, mapper, na_action=None): 

270 # Try to run function on index first, and then on elements of index 

271 # Especially important for group-by functionality 

272 try: 

273 result = mapper(self) 

274 

275 # Try to use this result if we can 

276 if isinstance(result, np.ndarray): 

277 result = Index(result) 

278 

279 if not isinstance(result, Index): 

280 raise TypeError("The map function must return an Index object") 

281 return result 

282 except Exception: 

283 return self.astype(object).map(mapper) 

284 

285 @Appender(Index.astype.__doc__) 

286 def astype(self, dtype, copy=True): 

287 if is_dtype_equal(self.dtype, dtype) and copy is False: 

288 # Ensure that self.astype(self.dtype) is self 

289 return self 

290 

291 new_values = self._data.astype(dtype, copy=copy) 

292 

293 # pass copy=False because any copying will be done in the 

294 # _data.astype call above 

295 return Index(new_values, dtype=new_values.dtype, name=self.name, copy=False)