Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/arrays/_mixins.py: 31%

1from __future__ import annotations 

2 

3from functools import wraps 

4from typing import ( 

5 TYPE_CHECKING, 

6 Any, 

7 Literal, 

8 Self, 

9 cast, 

10 overload, 

11) 

12 

13import numpy as np 

14 

15from pandas._libs import lib 

16from pandas._libs.arrays import NDArrayBacked 

17from pandas._libs.tslibs import is_supported_dtype 

18from pandas._typing import ( 

19 ArrayLike, 

20 AxisInt, 

21 Dtype, 

22 F, 

23 FillnaOptions, 

24 PositionalIndexer2D, 

25 PositionalIndexerTuple, 

26 ScalarIndexer, 

27 SequenceIndexer, 

28 Shape, 

29 TakeIndexer, 

30 npt, 

31) 

32from pandas.errors import AbstractMethodError 

33from pandas.util._decorators import doc 

34from pandas.util._validators import ( 

35 validate_bool_kwarg, 

36 validate_insert_loc, 

37) 

38 

39from pandas.core.dtypes.common import pandas_dtype 

40from pandas.core.dtypes.dtypes import ( 

41 DatetimeTZDtype, 

42 ExtensionDtype, 

43 PeriodDtype, 

44) 

45from pandas.core.dtypes.missing import array_equivalent 

46 

47from pandas.core import missing 

48from pandas.core.algorithms import ( 

49 take, 

50 unique, 

51 value_counts_internal as value_counts, 

52) 

53from pandas.core.array_algos.quantile import quantile_with_mask 

54from pandas.core.array_algos.transforms import shift 

55from pandas.core.arrays.base import ExtensionArray 

56from pandas.core.construction import extract_array 

57from pandas.core.indexers import ( 

58 check_array_indexer, 

59 getitem_returns_view, 

60) 

61from pandas.core.sorting import nargminmax 

62 

63if TYPE_CHECKING: 

64 from collections.abc import Sequence 

65 

66 from pandas._typing import ( 

67 NumpySorter, 

68 NumpyValueArrayLike, 

69 ) 

70 

71 from pandas import Series 

72 

73 

74def ravel_compat(meth: F) -> F: 

75 """ 

76 Decorator to ravel a 2D array before passing it to a cython operation, 

77 then reshape the result to our own shape. 

78 """ 

79 

80 @wraps(meth) 

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

82 if self.ndim == 1: 

83 return meth(self, *args, **kwargs) 

84 

85 flags = self._ndarray.flags 

86 flat = self.ravel("K") 

87 result = meth(flat, *args, **kwargs) 

88 order = "F" if flags.f_contiguous else "C" 

89 return result.reshape(self.shape, order=order) 

90 

91 return cast(F, method) 

92 

93 

94class NDArrayBackedExtensionArray(NDArrayBacked, ExtensionArray): 

95 """ 

96 ExtensionArray that is backed by a single NumPy ndarray. 

97 """ 

98 

99 _ndarray: np.ndarray 

100 

101 # scalar used to denote NA value inside our self._ndarray, e.g. -1 

102 # for Categorical, iNaT for Period. Outside of object dtype, 

103 # self.isna() should be exactly locations in self._ndarray with 

104 # _internal_fill_value. 

105 _internal_fill_value: Any 

106 

107 def _box_func(self, x): 

108 """ 

109 Wrap numpy type in our dtype.type if necessary. 

110 """ 

111 return x 

112 

113 def _validate_scalar(self, value): 

114 # used by NDArrayBackedExtensionIndex.insert 

115 raise AbstractMethodError(self) 

116 

117 # ------------------------------------------------------------------------ 

118 

119 @overload 

120 def view(self) -> Self: ... 

121 

122 @overload 

123 def view(self, dtype: Dtype | None = ...) -> ArrayLike: ... 

124 

125 def view(self, dtype: Dtype | None = None) -> ArrayLike: 

126 # We handle datetime64, datetime64tz, timedelta64, and period 

127 # dtypes here. Everything else we pass through to the underlying 

128 # ndarray. 

129 if dtype is None or dtype is self.dtype: 

130 return self._from_backing_data(self._ndarray) 

131 

132 if isinstance(dtype, type): 

133 # we sometimes pass non-dtype objects, e.g np.ndarray; 

134 # pass those through to the underlying ndarray 

135 return self._ndarray.view(dtype) 

136 

137 dtype = pandas_dtype(dtype) 

138 arr = self._ndarray 

139 

140 if isinstance(dtype, PeriodDtype): 

141 cls = dtype.construct_array_type() 

142 return cls(arr.view("i8"), dtype=dtype) 

143 elif isinstance(dtype, DatetimeTZDtype): 

144 dt_cls = dtype.construct_array_type() 

145 dt64_values = arr.view(f"M8[{dtype.unit}]") 

146 return dt_cls._simple_new(dt64_values, dtype=dtype) 

147 elif lib.is_np_dtype(dtype, "M") and is_supported_dtype(dtype): 

148 from pandas.core.arrays import DatetimeArray 

149 

150 dt64_values = arr.view(dtype) 

151 return DatetimeArray._simple_new(dt64_values, dtype=dtype) 

152 elif lib.is_np_dtype(dtype, "m") and is_supported_dtype(dtype): 

153 from pandas.core.arrays import TimedeltaArray 

154 

155 td64_values = arr.view(dtype) 

156 return TimedeltaArray._simple_new(td64_values, dtype=dtype) 

157 # error: Argument "dtype" to "view" of "ndarray" has incompatible type 

158 # "ExtensionDtype | dtype[Any]"; expected "dtype[Any] | _HasDType[dtype[Any]]" 

159 return arr.view(dtype=dtype) # type: ignore[arg-type] 

160 

161 def take( 

162 self, 

163 indices: TakeIndexer, 

164 *, 

165 allow_fill: bool = False, 

166 fill_value: Any = None, 

167 axis: AxisInt = 0, 

168 ) -> Self: 

169 if allow_fill: 

170 fill_value = self._validate_scalar(fill_value) 

171 

172 new_data = take( 

173 self._ndarray, 

174 indices, 

175 allow_fill=allow_fill, 

176 fill_value=fill_value, 

177 axis=axis, 

178 ) 

179 return self._from_backing_data(new_data) 

180 

181 # ------------------------------------------------------------------------ 

182 

183 def equals(self, other) -> bool: 

184 if type(self) is not type(other): 

185 return False 

186 if self.dtype != other.dtype: 

187 return False 

188 return bool(array_equivalent(self._ndarray, other._ndarray, dtype_equal=True)) 

189 

190 @classmethod 

191 def _from_factorized(cls, values, original): 

192 assert values.dtype == original._ndarray.dtype 

193 return original._from_backing_data(values) 

194 

195 def _values_for_argsort(self) -> np.ndarray: 

196 return self._ndarray 

197 

198 def _values_for_factorize(self): 

199 return self._ndarray, self._internal_fill_value 

200 

201 def _hash_pandas_object( 

202 self, *, encoding: str, hash_key: str, categorize: bool 

203 ) -> npt.NDArray[np.uint64]: 

204 from pandas.core.util.hashing import hash_array 

205 

206 values = self._ndarray 

207 return hash_array( 

208 values, encoding=encoding, hash_key=hash_key, categorize=categorize 

209 ) 

210 

211 def _cast_pointwise_result(self, values: ArrayLike) -> ArrayLike: 

212 values = np.asarray(values, dtype=object) 

213 return lib.maybe_convert_objects(values, convert_non_numeric=True) 

214 

215 # Signature of "argmin" incompatible with supertype "ExtensionArray" 

216 def argmin(self, axis: AxisInt = 0, skipna: bool = True): # type: ignore[override] 

217 # override base class by adding axis keyword 

218 validate_bool_kwarg(skipna, "skipna") 

219 if not skipna and self._hasna: 

220 raise ValueError("Encountered an NA value with skipna=False") 

221 return nargminmax(self, "argmin", axis=axis) 

222 

223 # Signature of "argmax" incompatible with supertype "ExtensionArray" 

224 def argmax(self, axis: AxisInt = 0, skipna: bool = True): # type: ignore[override] 

225 # override base class by adding axis keyword 

226 validate_bool_kwarg(skipna, "skipna") 

227 if not skipna and self._hasna: 

228 raise ValueError("Encountered an NA value with skipna=False") 

229 return nargminmax(self, "argmax", axis=axis) 

230 

231 def unique(self) -> Self: 

232 new_data = unique(self._ndarray) 

233 return self._from_backing_data(new_data) 

234 

235 @classmethod 

236 @doc(ExtensionArray._concat_same_type) 

237 def _concat_same_type( 

238 cls, 

239 to_concat: Sequence[Self], 

240 axis: AxisInt = 0, 

241 ) -> Self: 

242 if not lib.dtypes_all_equal([x.dtype for x in to_concat]): 

243 dtypes = {str(x.dtype) for x in to_concat} 

244 raise ValueError("to_concat must have the same dtype", dtypes) 

245 

246 return super()._concat_same_type(to_concat, axis=axis) 

247 

248 @doc(ExtensionArray.searchsorted) 

249 def searchsorted( 

250 self, 

251 value: NumpyValueArrayLike | ExtensionArray, 

252 side: Literal["left", "right"] = "left", 

253 sorter: NumpySorter | None = None, 

254 ) -> npt.NDArray[np.intp] | np.intp: 

255 npvalue = self._validate_setitem_value(value) 

256 return self._ndarray.searchsorted(npvalue, side=side, sorter=sorter) 

257 

258 @doc(ExtensionArray.shift) 

259 def shift(self, periods: int = 1, fill_value=None) -> Self: 

260 # NB: shift is always along axis=0 

261 axis = 0 

262 fill_value = self._validate_scalar(fill_value) 

263 new_values = shift(self._ndarray, periods, axis, fill_value) 

264 

265 return self._from_backing_data(new_values) 

266 

267 def __setitem__(self, key, value) -> None: 

268 if self._readonly: 

269 raise ValueError("Cannot modify read-only array") 

270 

271 key = check_array_indexer(self, key) 

272 value = self._validate_setitem_value(value) 

273 self._ndarray[key] = value 

274 

275 def _validate_setitem_value(self, value): 

276 return value 

277 

278 @overload 

279 def __getitem__(self, key: ScalarIndexer) -> Any: ... 

280 

281 @overload 

282 def __getitem__( 

283 self, 

284 key: SequenceIndexer | PositionalIndexerTuple, 

285 ) -> Self: ... 

286 

287 def __getitem__( 

288 self, 

289 key: PositionalIndexer2D, 

290 ) -> Self | Any: 

291 if lib.is_integer(key): 

292 # fast-path 

293 result = self._ndarray[key] 

294 if self.ndim == 1: 

295 return self._box_func(result) 

296 result = self._from_backing_data(result) 

297 if getitem_returns_view(self, key): 

298 result._readonly = self._readonly 

299 return result 

300 

301 # error: Incompatible types in assignment (expression has type "ExtensionArray", 

302 # variable has type "Union[int, slice, ndarray]") 

303 key = extract_array(key, extract_numpy=True) # type: ignore[assignment] 

304 key = check_array_indexer(self, key) 

305 result = self._ndarray[key] 

306 if lib.is_scalar(result): 

307 return self._box_func(result) 

308 

309 result = self._from_backing_data(result) 

310 if getitem_returns_view(self, key): 

311 result._readonly = self._readonly 

312 return result 

313 

314 def _pad_or_backfill( 

315 self, 

316 *, 

317 method: FillnaOptions, 

318 limit: int | None = None, 

319 limit_area: Literal["inside", "outside"] | None = None, 

320 copy: bool = True, 

321 ) -> Self: 

322 mask = self.isna() 

323 if mask.any(): 

324 # (for now) when self.ndim == 2, we assume axis=0 

325 func = missing.get_fill_func(method, ndim=self.ndim) 

326 

327 npvalues = self._ndarray.T 

328 if copy: 

329 npvalues = npvalues.copy() 

330 func(npvalues, limit=limit, limit_area=limit_area, mask=mask.T) 

331 npvalues = npvalues.T 

332 

333 if copy: 

334 new_values = self._from_backing_data(npvalues) 

335 else: 

336 new_values = self 

337 

338 elif copy: 

339 new_values = self.copy() 

340 else: 

341 new_values = self 

342 return new_values 

343 

344 @doc(ExtensionArray.fillna) 

345 def fillna(self, value, limit: int | None = None, copy: bool = True) -> Self: 

346 mask = self.isna() 

347 if limit is not None and limit < len(self): 

348 # mypy doesn't like that mask can be an EA which need not have `cumsum` 

349 modify = mask.cumsum() > limit # type: ignore[union-attr] 

350 if modify.any(): 

351 # Only copy mask if necessary 

352 mask = mask.copy() 

353 mask[modify] = False 

354 # error: Argument 2 to "check_value_size" has incompatible type 

355 # "ExtensionArray"; expected "ndarray" 

356 value = missing.check_value_size( 

357 value, 

358 mask, # type: ignore[arg-type] 

359 len(self), 

360 ) 

361 

362 if mask.any(): 

363 # fill with value 

364 if copy: 

365 new_values = self.copy() 

366 else: 

367 new_values = self[:] 

368 new_values[mask] = value 

369 else: 

370 # We validate the fill_value even if there is nothing to fill 

371 self._validate_setitem_value(value) 

372 

373 if not copy: 

374 new_values = self[:] 

375 else: 

376 new_values = self.copy() 

377 return new_values 

378 

379 # ------------------------------------------------------------------------ 

380 # Reductions 

381 

382 def _wrap_reduction_result(self, axis: AxisInt | None, result) -> Any: 

383 if axis is None or self.ndim == 1: 

384 return self._box_func(result) 

385 return self._from_backing_data(result) 

386 

387 # ------------------------------------------------------------------------ 

388 # __array_function__ methods 

389 

390 def _putmask(self, mask: npt.NDArray[np.bool_], value) -> None: 

391 """ 

392 Analogue to np.putmask(self, mask, value) 

393 

394 Parameters 

395 ---------- 

396 mask : np.ndarray[bool] 

397 value : scalar or listlike 

398 

399 Raises 

400 ------ 

401 TypeError 

402 If value cannot be cast to self.dtype. 

403 """ 

404 value = self._validate_setitem_value(value) 

405 

406 np.putmask(self._ndarray, mask, value) 

407 

408 def _where(self: Self, mask: npt.NDArray[np.bool_], value) -> Self: 

409 """ 

410 Analogue to np.where(mask, self, value) 

411 

412 Parameters 

413 ---------- 

414 mask : np.ndarray[bool] 

415 value : scalar or listlike 

416 

417 Raises 

418 ------ 

419 TypeError 

420 If value cannot be cast to self.dtype. 

421 """ 

422 value = self._validate_setitem_value(value) 

423 

424 res_values = np.where(mask, self._ndarray, value) 

425 if res_values.dtype != self._ndarray.dtype: 

426 raise AssertionError( 

427 # GH#56410 

428 "Something has gone wrong, please report a bug at " 

429 "github.com/pandas-dev/pandas/" 

430 ) 

431 return self._from_backing_data(res_values) 

432 

433 # ------------------------------------------------------------------------ 

434 # Index compat methods 

435 

436 def insert(self, loc: int, item) -> Self: 

437 """ 

438 Make new ExtensionArray inserting new item at location. Follows 

439 Python list.append semantics for negative values. 

440 

441 Parameters 

442 ---------- 

443 loc : int 

444 item : object 

445 

446 Returns 

447 ------- 

448 type(self) 

449 """ 

450 loc = validate_insert_loc(loc, len(self)) 

451 

452 code = self._validate_scalar(item) 

453 

454 new_vals = np.concatenate( 

455 ( 

456 self._ndarray[:loc], 

457 np.asarray([code], dtype=self._ndarray.dtype), 

458 self._ndarray[loc:], 

459 ) 

460 ) 

461 return self._from_backing_data(new_vals) 

462 

463 # ------------------------------------------------------------------------ 

464 # Additional array methods 

465 # These are not part of the EA API, but we implement them because 

466 # pandas assumes they're there. 

467 

468 def value_counts(self, dropna: bool = True) -> Series: 

469 """ 

470 Return a Series containing counts of unique values. 

471 

472 Parameters 

473 ---------- 

474 dropna : bool, default True 

475 Don't include counts of NA values. 

476 

477 Returns 

478 ------- 

479 Series 

480 """ 

481 if self.ndim != 1: 

482 raise NotImplementedError 

483 

484 from pandas import ( 

485 Index, 

486 Series, 

487 ) 

488 

489 if dropna: 

490 # error: Unsupported operand type for ~ ("ExtensionArray") 

491 values = self[~self.isna()]._ndarray # type: ignore[operator] 

492 else: 

493 values = self._ndarray 

494 

495 result = value_counts(values, sort=False, dropna=dropna) 

496 

497 index_arr = self._from_backing_data(np.asarray(result.index._data)) 

498 index = Index(index_arr, name=result.index.name, copy=False) 

499 return Series(result._values, index=index, name=result.name, copy=False) 

500 

501 def _quantile( 

502 self, 

503 qs: npt.NDArray[np.float64], 

504 interpolation: str, 

505 ) -> Self: 

506 # TODO: disable for Categorical if not ordered? 

507 

508 mask = np.asarray(self.isna()) 

509 arr = self._ndarray 

510 fill_value = self._internal_fill_value 

511 

512 res_values = quantile_with_mask(arr, mask, fill_value, qs, interpolation) 

513 if res_values.dtype == self._ndarray.dtype: 

514 return self._from_backing_data(res_values) 

515 else: 

516 # e.g. test_quantile_empty we are empty integer dtype and res_values 

517 # has floating dtype 

518 # TODO: technically __init__ isn't defined here. 

519 # Should we raise NotImplementedError and handle this on NumpyEA? 

520 return type(self)(res_values) # type: ignore[call-arg] 

521 

522 # ------------------------------------------------------------------------ 

523 # numpy-like methods 

524 

525 @classmethod 

526 def _empty(cls, shape: Shape, dtype: ExtensionDtype) -> Self: 

527 """ 

528 Analogous to np.empty(shape, dtype=dtype) 

529 

530 Parameters 

531 ---------- 

532 shape : tuple[int] 

533 dtype : ExtensionDtype 

534 """ 

535 # The base implementation uses a naive approach to find the dtype 

536 # for the backing ndarray 

537 arr = cls._from_sequence([], dtype=dtype) 

538 backing = np.empty(shape, dtype=arr._ndarray.dtype) 

539 return arr._from_backing_data(backing)