Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/strings/object_array.py: 22%

1from __future__ import annotations 

2 

3import functools 

4import re 

5import textwrap 

6from typing import ( 

7 TYPE_CHECKING, 

8 Literal, 

9 cast, 

10) 

11import unicodedata 

12 

13import numpy as np 

14 

15from pandas._libs import lib 

16import pandas._libs.missing as libmissing 

17import pandas._libs.ops as libops 

18from pandas.util._validators import validate_na_arg 

19 

20from pandas.core.dtypes.common import pandas_dtype 

21from pandas.core.dtypes.missing import isna 

22 

23if TYPE_CHECKING: 

24 from collections.abc import ( 

25 Callable, 

26 Sequence, 

27 ) 

28 

29 from pandas._typing import ( 

30 NpDtype, 

31 Scalar, 

32 ) 

33 

34 

35class ObjectStringArrayMixin: 

36 """ 

37 String Methods operating on object-dtype ndarrays. 

38 """ 

39 

40 def __len__(self) -> int: 

41 # For typing, _str_map relies on the object being sized. 

42 raise NotImplementedError 

43 

44 def _str_getitem(self, key): 

45 if isinstance(key, slice): 

46 return self._str_slice(start=key.start, stop=key.stop, step=key.step) 

47 else: 

48 return self._str_get(key) 

49 

50 def _str_map( 

51 self, 

52 f, 

53 na_value=lib.no_default, 

54 dtype: NpDtype | None = None, 

55 convert: bool = True, 

56 ): 

57 """ 

58 Map a callable over valid elements of the array. 

59 

60 Parameters 

61 ---------- 

62 f : Callable 

63 A function to call on each non-NA element. 

64 na_value : Scalar, optional 

65 The value to set for NA values. Might also be used for the 

66 fill value if the callable `f` raises an exception. 

67 This defaults to ``self.dtype.na_value`` which is ``np.nan`` 

68 for object-dtype and Categorical and ``pd.NA`` for StringArray. 

69 dtype : Dtype, optional 

70 The dtype of the result array. 

71 convert : bool, default True 

72 Whether to call `maybe_convert_objects` on the resulting ndarray 

73 """ 

74 if dtype is None: 

75 dtype = np.dtype("object") 

76 if na_value is lib.no_default: 

77 na_value = self.dtype.na_value # type: ignore[attr-defined] 

78 

79 if not len(self): 

80 return np.array([], dtype=dtype) 

81 

82 arr = np.asarray(self, dtype=object) 

83 mask = isna(arr) 

84 map_convert = convert and not np.all(mask) 

85 try: 

86 result = lib.map_infer_mask( 

87 arr, f, mask.view(np.uint8), convert=map_convert 

88 ) 

89 except (TypeError, AttributeError) as err: 

90 # Reraise the exception if callable `f` got wrong number of args. 

91 # The user may want to be warned by this, instead of getting NaN 

92 p_err = ( 

93 r"((takes)|(missing)) (?(2)from \d+ to )?\d+ " 

94 r"(?(3)required )positional arguments?" 

95 ) 

96 

97 if len(err.args) >= 1 and re.search(p_err, err.args[0]): 

98 # FIXME: this should be totally avoidable 

99 raise err 

100 

101 def g(x): 

102 # This type of fallback behavior can be removed once 

103 # we remove object-dtype .str accessor. 

104 try: 

105 return f(x) 

106 except (TypeError, AttributeError): 

107 return na_value 

108 

109 return self._str_map(g, na_value=na_value, dtype=dtype) 

110 if not isinstance(result, np.ndarray): 

111 return result 

112 if na_value is not np.nan: 

113 np.putmask(result, mask, na_value) 

114 if convert and result.dtype == object: 

115 result = lib.maybe_convert_objects(result) 

116 return result 

117 

118 def _str_count(self, pat, flags: int = 0): 

119 regex = re.compile(pat, flags=flags) 

120 f = lambda x: len(regex.findall(x)) 

121 return self._str_map(f, dtype="int64") 

122 

123 def _str_pad( 

124 self, 

125 width: int, 

126 side: Literal["left", "right", "both"] = "left", 

127 fillchar: str = " ", 

128 ): 

129 if side == "left": 

130 f = lambda x: x.rjust(width, fillchar) 

131 elif side == "right": 

132 f = lambda x: x.ljust(width, fillchar) 

133 elif side == "both": 

134 f = lambda x: x.center(width, fillchar) 

135 else: # pragma: no cover 

136 raise ValueError("Invalid side") 

137 return self._str_map(f) 

138 

139 def _str_contains( 

140 self, 

141 pat, 

142 case: bool = True, 

143 flags: int = 0, 

144 na=lib.no_default, 

145 regex: bool = True, 

146 ): 

147 validate_na_arg(na, name="na") 

148 if regex: 

149 if not case: 

150 flags |= re.IGNORECASE 

151 

152 pat = re.compile(pat, flags=flags) 

153 

154 f = lambda x: pat.search(x) is not None 

155 elif case: 

156 f = lambda x: pat in x 

157 else: 

158 upper_pat = pat.upper() 

159 f = lambda x: upper_pat in x.upper() 

160 return self._str_map(f, na, dtype=np.dtype("bool")) 

161 

162 def _str_startswith(self, pat, na=lib.no_default): 

163 validate_na_arg(na, name="na") 

164 f = lambda x: x.startswith(pat) 

165 return self._str_map(f, na_value=na, dtype=np.dtype(bool)) 

166 

167 def _str_endswith(self, pat, na=lib.no_default): 

168 validate_na_arg(na, name="na") 

169 f = lambda x: x.endswith(pat) 

170 return self._str_map(f, na_value=na, dtype=np.dtype(bool)) 

171 

172 def _str_replace( 

173 self, 

174 pat: str | re.Pattern, 

175 repl: str | Callable, 

176 n: int = -1, 

177 case: bool = True, 

178 flags: int = 0, 

179 regex: bool = True, 

180 ): 

181 if case is False: 

182 # add case flag, if provided 

183 flags |= re.IGNORECASE 

184 

185 if regex or flags or callable(repl): 

186 if not isinstance(pat, re.Pattern): 

187 if regex is False: 

188 pat = re.escape(pat) 

189 pat = re.compile(pat, flags=flags) 

190 

191 n = n if n >= 0 else 0 

192 f = lambda x: pat.sub(repl=repl, string=x, count=n) 

193 else: 

194 f = lambda x: x.replace(pat, repl, n) 

195 

196 return self._str_map(f, dtype=str) 

197 

198 def _str_repeat(self, repeats: int | Sequence[int]): 

199 if lib.is_integer(repeats): 

200 rint = cast(int, repeats) 

201 

202 def scalar_rep(x): 

203 try: 

204 return bytes.__mul__(x, rint) 

205 except TypeError: 

206 return str.__mul__(x, rint) 

207 

208 return self._str_map(scalar_rep, dtype=str) 

209 else: 

210 from pandas.core.arrays.string_ import BaseStringArray 

211 

212 def rep(x, r): 

213 if x is libmissing.NA: 

214 return x 

215 try: 

216 return bytes.__mul__(x, r) 

217 except TypeError: 

218 return str.__mul__(x, r) 

219 

220 result = libops.vec_binop( 

221 np.asarray(self), 

222 np.asarray(repeats, dtype=object), 

223 rep, 

224 ) 

225 if not isinstance(self, BaseStringArray): 

226 return result 

227 # Not going through map, so we have to do this here. 

228 return type(self)._from_sequence(result, dtype=self.dtype) 

229 

230 def _str_match( 

231 self, 

232 pat: str | re.Pattern, 

233 case: bool = True, 

234 flags: int = 0, 

235 na: Scalar | lib.NoDefault = lib.no_default, 

236 ): 

237 if not case: 

238 flags |= re.IGNORECASE 

239 

240 if isinstance(pat, re.Pattern): 

241 # We need to check that flags matches pat.flags. 

242 # pat.flags will have re.U regardless, so we need to add it here 

243 # before checking for a match 

244 flags = flags | re.U 

245 

246 if flags != pat.flags: 

247 raise ValueError("Cannot pass flags that do not match pat.flags") 

248 regex = pat 

249 else: 

250 regex = re.compile(pat, flags=flags) 

251 

252 f = lambda x: regex.match(x) is not None 

253 return self._str_map(f, na_value=na, dtype=np.dtype(bool)) 

254 

255 def _str_fullmatch( 

256 self, 

257 pat: str | re.Pattern, 

258 case: bool = True, 

259 flags: int = 0, 

260 na: Scalar | lib.NoDefault = lib.no_default, 

261 ): 

262 if not case: 

263 flags |= re.IGNORECASE 

264 

265 regex = re.compile(pat, flags=flags) 

266 

267 f = lambda x: regex.fullmatch(x) is not None 

268 return self._str_map(f, na_value=na, dtype=np.dtype(bool)) 

269 

270 def _str_encode(self, encoding, errors: str = "strict"): 

271 f = lambda x: x.encode(encoding, errors=errors) 

272 return self._str_map(f, dtype=object) 

273 

274 def _str_find(self, sub, start: int = 0, end=None): 

275 return self._str_find_(sub, start, end, side="left") 

276 

277 def _str_rfind(self, sub, start: int = 0, end=None): 

278 return self._str_find_(sub, start, end, side="right") 

279 

280 def _str_find_(self, sub, start, end, side): 

281 if side == "left": 

282 method = "find" 

283 elif side == "right": 

284 method = "rfind" 

285 else: # pragma: no cover 

286 raise ValueError("Invalid side") 

287 

288 if end is None: 

289 f = lambda x: getattr(x, method)(sub, start) 

290 else: 

291 f = lambda x: getattr(x, method)(sub, start, end) 

292 return self._str_map(f, dtype="int64") 

293 

294 def _str_findall(self, pat, flags: int = 0): 

295 regex = re.compile(pat, flags=flags) 

296 return self._str_map(regex.findall, dtype="object") 

297 

298 def _str_get(self, i): 

299 def f(x): 

300 if isinstance(x, dict): 

301 return x.get(i) 

302 elif len(x) > i >= -len(x): 

303 return x[i] 

304 return self.dtype.na_value # type: ignore[attr-defined] 

305 

306 return self._str_map(f) 

307 

308 def _str_index(self, sub, start: int = 0, end=None): 

309 if end: 

310 f = lambda x: x.index(sub, start, end) 

311 else: 

312 f = lambda x: x.index(sub, start, end) 

313 return self._str_map(f, dtype="int64") 

314 

315 def _str_rindex(self, sub, start: int = 0, end=None): 

316 if end: 

317 f = lambda x: x.rindex(sub, start, end) 

318 else: 

319 f = lambda x: x.rindex(sub, start, end) 

320 return self._str_map(f, dtype="int64") 

321 

322 def _str_join(self, sep: str): 

323 return self._str_map(sep.join) 

324 

325 def _str_partition(self, sep: str, expand): 

326 result = self._str_map(lambda x: x.partition(sep), dtype="object") 

327 return result 

328 

329 def _str_rpartition(self, sep: str, expand): 

330 return self._str_map(lambda x: x.rpartition(sep), dtype="object") 

331 

332 def _str_len(self): 

333 return self._str_map(len, dtype="int64") 

334 

335 def _str_slice(self, start=None, stop=None, step=None): 

336 obj = slice(start, stop, step) 

337 return self._str_map(lambda x: x[obj]) 

338 

339 def _str_slice_replace(self, start=None, stop=None, repl=None): 

340 if repl is None: 

341 repl = "" 

342 

343 def f(x): 

344 if x[start:stop] == "": 

345 local_stop = start 

346 else: 

347 local_stop = stop 

348 y = "" 

349 if start is not None: 

350 y += x[:start] 

351 y += repl 

352 if stop is not None: 

353 y += x[local_stop:] 

354 return y 

355 

356 return self._str_map(f) 

357 

358 def _str_split( 

359 self, 

360 pat: str | re.Pattern | None = None, 

361 n=-1, 

362 expand: bool = False, 

363 regex: bool | None = None, 

364 ): 

365 if pat is None: 

366 if n is None or n == 0: 

367 n = -1 

368 f = lambda x: x.split(pat, n) 

369 else: 

370 new_pat: str | re.Pattern 

371 if regex is True or isinstance(pat, re.Pattern): 

372 new_pat = re.compile(pat) 

373 elif regex is False: 

374 new_pat = pat 

375 # regex is None so link to old behavior #43563 

376 elif len(pat) == 1: 

377 new_pat = pat 

378 else: 

379 new_pat = re.compile(pat) 

380 

381 if isinstance(new_pat, re.Pattern): 

382 if n is None or n == -1: 

383 n = 0 

384 f = lambda x: new_pat.split(x, maxsplit=n) 

385 else: 

386 if n is None or n == 0: 

387 n = -1 

388 f = lambda x: x.split(pat, n) 

389 return self._str_map(f, dtype=object) 

390 

391 def _str_rsplit(self, pat=None, n=-1): 

392 if n is None or n == 0: 

393 n = -1 

394 f = lambda x: x.rsplit(pat, n) 

395 return self._str_map(f, dtype="object") 

396 

397 def _str_translate(self, table): 

398 return self._str_map(lambda x: x.translate(table)) 

399 

400 def _str_wrap(self, width: int, **kwargs): 

401 kwargs["width"] = width 

402 tw = textwrap.TextWrapper(**kwargs) 

403 return self._str_map(lambda s: "\n".join(tw.wrap(s))) 

404 

405 def _str_get_dummies(self, sep: str = "|", dtype: NpDtype | None = None): 

406 from pandas import Series 

407 

408 if dtype is None: 

409 dtype = np.int64 

410 arr = Series(self).fillna("") 

411 try: 

412 arr = sep + arr + sep 

413 except (TypeError, NotImplementedError): 

414 arr = sep + arr.astype(str) + sep 

415 

416 tags: set[str] = set() 

417 for ts in Series(arr, copy=False).str.split(sep): 

418 tags.update(ts) 

419 tags2 = sorted(tags - {""}) 

420 

421 _dtype = pandas_dtype(dtype) 

422 dummies_dtype: NpDtype 

423 if isinstance(_dtype, np.dtype): 

424 dummies_dtype = _dtype 

425 else: 

426 dummies_dtype = np.bool_ 

427 dummies = np.empty((len(arr), len(tags2)), dtype=dummies_dtype, order="F") 

428 

429 def _isin(test_elements: str, element: str) -> bool: 

430 return element in test_elements 

431 

432 for i, t in enumerate(tags2): 

433 pat = sep + t + sep 

434 dummies[:, i] = lib.map_infer( 

435 arr.to_numpy(), functools.partial(_isin, element=pat) 

436 ) 

437 return dummies, tags2 

438 

439 def _str_upper(self): 

440 return self._str_map(lambda x: x.upper()) 

441 

442 def _str_isalnum(self): 

443 return self._str_map(str.isalnum, dtype="bool") 

444 

445 def _str_isalpha(self): 

446 return self._str_map(str.isalpha, dtype="bool") 

447 

448 def _str_isascii(self): 

449 return self._str_map(str.isascii, dtype="bool") 

450 

451 def _str_isdecimal(self): 

452 return self._str_map(str.isdecimal, dtype="bool") 

453 

454 def _str_isdigit(self): 

455 return self._str_map(str.isdigit, dtype="bool") 

456 

457 def _str_islower(self): 

458 return self._str_map(str.islower, dtype="bool") 

459 

460 def _str_isnumeric(self): 

461 return self._str_map(str.isnumeric, dtype="bool") 

462 

463 def _str_isspace(self): 

464 return self._str_map(str.isspace, dtype="bool") 

465 

466 def _str_istitle(self): 

467 return self._str_map(str.istitle, dtype="bool") 

468 

469 def _str_isupper(self): 

470 return self._str_map(str.isupper, dtype="bool") 

471 

472 def _str_capitalize(self): 

473 return self._str_map(str.capitalize) 

474 

475 def _str_casefold(self): 

476 return self._str_map(str.casefold) 

477 

478 def _str_title(self): 

479 return self._str_map(str.title) 

480 

481 def _str_swapcase(self): 

482 return self._str_map(str.swapcase) 

483 

484 def _str_lower(self): 

485 return self._str_map(str.lower) 

486 

487 def _str_normalize(self, form): 

488 f = lambda x: unicodedata.normalize(form, x) 

489 return self._str_map(f) 

490 

491 def _str_strip(self, to_strip=None): 

492 return self._str_map(lambda x: x.strip(to_strip)) 

493 

494 def _str_lstrip(self, to_strip=None): 

495 return self._str_map(lambda x: x.lstrip(to_strip)) 

496 

497 def _str_rstrip(self, to_strip=None): 

498 return self._str_map(lambda x: x.rstrip(to_strip)) 

499 

500 def _str_removeprefix(self, prefix: str): 

501 return self._str_map(lambda x: x.removeprefix(prefix)) 

502 

503 def _str_removesuffix(self, suffix: str): 

504 return self._str_map(lambda x: x.removesuffix(suffix)) 

505 

506 def _str_extract(self, pat: str, flags: int = 0, expand: bool = True): 

507 regex = re.compile(pat, flags=flags) 

508 na_value = self.dtype.na_value # type: ignore[attr-defined] 

509 

510 if not expand: 

511 

512 def g(x): 

513 m = regex.search(x) 

514 return m.groups()[0] if m else na_value 

515 

516 return self._str_map(g, convert=False) 

517 

518 empty_row = [na_value] * regex.groups 

519 

520 def f(x): 

521 if not isinstance(x, str): 

522 return empty_row 

523 m = regex.search(x) 

524 if m: 

525 return [na_value if item is None else item for item in m.groups()] 

526 else: 

527 return empty_row 

528 

529 return [f(val) for val in np.asarray(self)] 

530 

531 def _str_zfill(self, width: int): 

532 return self._str_map(lambda x: x.zfill(width))