Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/col.py: 29%

1from __future__ import annotations 

2 

3from collections.abc import ( 

4 Callable, 

5 Hashable, 

6 Sequence, 

7) 

8from typing import ( 

9 TYPE_CHECKING, 

10 Any, 

11 NoReturn, 

12) 

13 

14from pandas.util._decorators import set_module 

15 

16if TYPE_CHECKING: 

17 from pandas import ( 

18 DataFrame, 

19 Series, 

20 ) 

21 

22 

23# Used only for generating the str repr of expressions. 

24_OP_SYMBOLS = { 

25 "__add__": "+", 

26 "__radd__": "+", 

27 "__sub__": "-", 

28 "__rsub__": "-", 

29 "__mul__": "*", 

30 "__rmul__": "*", 

31 "__truediv__": "/", 

32 "__rtruediv__": "/", 

33 "__floordiv__": "//", 

34 "__rfloordiv__": "//", 

35 "__mod__": "%", 

36 "__rmod__": "%", 

37 "__ge__": ">=", 

38 "__gt__": ">", 

39 "__le__": "<=", 

40 "__lt__": "<", 

41 "__eq__": "==", 

42 "__ne__": "!=", 

43 "__and__": "&", 

44 "__rand__": "&", 

45 "__or__": "|", 

46 "__ror__": "|", 

47 "__xor__": "^", 

48 "__rxor__": "^", 

49} 

50 

51 

52def _parse_args(df: DataFrame, *args: Any) -> tuple[Any, ...]: 

53 # Parse `args`, evaluating any expressions we encounter. 

54 return tuple( 

55 x._eval_expression(df) if isinstance(x, Expression) else x for x in args 

56 ) 

57 

58 

59def _parse_kwargs(df: DataFrame, **kwargs: Any) -> dict[str, Any]: 

60 # Parse `kwargs`, evaluating any expressions we encounter. 

61 return { 

62 key: val._eval_expression(df) if isinstance(val, Expression) else val 

63 for key, val in kwargs.items() 

64 } 

65 

66 

67def _pretty_print_args_kwargs(*args: Any, **kwargs: Any) -> str: 

68 inputs_repr = ", ".join(repr(arg) for arg in args) 

69 kwargs_repr = ", ".join(f"{k}={v!r}" for k, v in kwargs.items()) 

70 

71 all_args = [] 

72 if inputs_repr: 

73 all_args.append(inputs_repr) 

74 if kwargs_repr: 

75 all_args.append(kwargs_repr) 

76 

77 return ", ".join(all_args) 

78 

79 

80@set_module("pandas.api.typing") 

81class Expression: 

82 """ 

83 Class representing a deferred expression evaluated against a DataFrame. 

84 

85 Expressions are initially created via ``pd.col`` and can be composed 

86 using arithmetic, comparison, and any method or attribute that the 

87 underlying object supports. This includes NumPy ufuncs, indexing, and 

88 the boolean operators ``&``, ``|``, and ``~``. 

89 

90 Expressions are solely intended for use in expression-based APIs, such 

91 as :meth:`DataFrame.assign <pandas.DataFrame.assign>` and 

92 :meth:`DataFrame.loc <pandas.DataFrame.loc>`, where they are 

93 evaluated against the DataFrame they are passed to. 

94 

95 This is not meant to be instantiated directly. Instead, use :meth:`pandas.col`. 

96 

97 See Also 

98 -------- 

99 pandas.col : Create a deferred column reference. 

100 

101 Notes 

102 ----- 

103 The attributes that are available depend on the expression and where it 

104 is used. Attribute access is itself deferred: it is applied to the 

105 result of evaluating the expression, so an attribute that is not 

106 supported by that result raises when the expression is evaluated. 

107 

108 Examples 

109 -------- 

110 >>> df = pd.DataFrame({"a": [1, 2], "b": [3, 4]}) 

111 >>> df.assign(c=pd.col("a") + pd.col("b")) 

112 a b c 

113 0 1 3 4 

114 1 2 4 6 

115 """ 

116 

117 def __init__( 

118 self, 

119 func: Callable[[DataFrame], Any], 

120 repr_str: str, 

121 needs_parenthese: bool = False, 

122 ) -> None: 

123 self._func = func 

124 self._repr_str = repr_str 

125 self._needs_parentheses = needs_parenthese 

126 

127 def _eval_expression(self, df: DataFrame) -> Any: 

128 return self._func(df) 

129 

130 def _with_op( 

131 self, op: str, other: Any, repr_str: str, needs_parentheses: bool = True 

132 ) -> Expression: 

133 if isinstance(other, Expression): 

134 return Expression( 

135 lambda df: getattr(self._eval_expression(df), op)( 

136 other._eval_expression(df) 

137 ), 

138 repr_str, 

139 needs_parenthese=needs_parentheses, 

140 ) 

141 else: 

142 return Expression( 

143 lambda df: getattr(self._eval_expression(df), op)(other), 

144 repr_str, 

145 needs_parenthese=needs_parentheses, 

146 ) 

147 

148 def _maybe_wrap_parentheses(self, other: Any) -> tuple[str, str]: 

149 if self._needs_parentheses: 

150 self_repr = f"({self!r})" 

151 else: 

152 self_repr = f"{self!r}" 

153 if isinstance(other, Expression) and other._needs_parentheses: 

154 other_repr = f"({other!r})" 

155 else: 

156 other_repr = f"{other!r}" 

157 return self_repr, other_repr 

158 

159 # Binary ops 

160 def __add__(self, other: Any) -> Expression: 

161 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

162 return self._with_op("__add__", other, f"{self_repr} + {other_repr}") 

163 

164 def __radd__(self, other: Any) -> Expression: 

165 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

166 return self._with_op("__radd__", other, f"{other_repr} + {self_repr}") 

167 

168 def __sub__(self, other: Any) -> Expression: 

169 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

170 return self._with_op("__sub__", other, f"{self_repr} - {other_repr}") 

171 

172 def __rsub__(self, other: Any) -> Expression: 

173 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

174 return self._with_op("__rsub__", other, f"{other_repr} - {self_repr}") 

175 

176 def __mul__(self, other: Any) -> Expression: 

177 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

178 return self._with_op("__mul__", other, f"{self_repr} * {other_repr}") 

179 

180 def __rmul__(self, other: Any) -> Expression: 

181 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

182 return self._with_op("__rmul__", other, f"{other_repr} * {self_repr}") 

183 

184 def __matmul__(self, other: Any) -> Expression: 

185 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

186 return self._with_op("__matmul__", other, f"{self_repr} @ {other_repr}") 

187 

188 def __rmatmul__(self, other: Any) -> Expression: 

189 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

190 return self._with_op("__rmatmul__", other, f"{other_repr} @ {self_repr}") 

191 

192 def __pow__(self, other: Any) -> Expression: 

193 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

194 return self._with_op("__pow__", other, f"{self_repr} ** {other_repr}") 

195 

196 def __rpow__(self, other: Any) -> Expression: 

197 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

198 return self._with_op("__rpow__", other, f"{other_repr} ** {self_repr}") 

199 

200 def __truediv__(self, other: Any) -> Expression: 

201 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

202 return self._with_op("__truediv__", other, f"{self_repr} / {other_repr}") 

203 

204 def __rtruediv__(self, other: Any) -> Expression: 

205 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

206 return self._with_op("__rtruediv__", other, f"{other_repr} / {self_repr}") 

207 

208 def __floordiv__(self, other: Any) -> Expression: 

209 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

210 return self._with_op("__floordiv__", other, f"{self_repr} // {other_repr}") 

211 

212 def __rfloordiv__(self, other: Any) -> Expression: 

213 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

214 return self._with_op("__rfloordiv__", other, f"{other_repr} // {self_repr}") 

215 

216 def __ge__(self, other: Any) -> Expression: 

217 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

218 return self._with_op("__ge__", other, f"{self_repr} >= {other_repr}") 

219 

220 def __gt__(self, other: Any) -> Expression: 

221 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

222 return self._with_op("__gt__", other, f"{self_repr} > {other_repr}") 

223 

224 def __le__(self, other: Any) -> Expression: 

225 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

226 return self._with_op("__le__", other, f"{self_repr} <= {other_repr}") 

227 

228 def __lt__(self, other: Any) -> Expression: 

229 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

230 return self._with_op("__lt__", other, f"{self_repr} < {other_repr}") 

231 

232 def __eq__(self, other: object) -> Expression: # type: ignore[override] 

233 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

234 return self._with_op("__eq__", other, f"{self_repr} == {other_repr}") 

235 

236 def __ne__(self, other: object) -> Expression: # type: ignore[override] 

237 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

238 return self._with_op("__ne__", other, f"{self_repr} != {other_repr}") 

239 

240 def __mod__(self, other: Any) -> Expression: 

241 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

242 return self._with_op("__mod__", other, f"{self_repr} % {other_repr}") 

243 

244 def __rmod__(self, other: Any) -> Expression: 

245 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

246 return self._with_op("__rmod__", other, f"{other_repr} % {self_repr}") 

247 

248 # Logical ops 

249 def __and__(self, other: Any) -> Expression: 

250 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

251 return self._with_op("__and__", other, f"{self_repr} & {other_repr}") 

252 

253 def __rand__(self, other: Any) -> Expression: 

254 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

255 return self._with_op("__rand__", other, f"{other_repr} & {self_repr}") 

256 

257 def __or__(self, other: Any) -> Expression: 

258 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

259 return self._with_op("__or__", other, f"{self_repr} | {other_repr}") 

260 

261 def __ror__(self, other: Any) -> Expression: 

262 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

263 return self._with_op("__ror__", other, f"{other_repr} | {self_repr}") 

264 

265 def __xor__(self, other: Any) -> Expression: 

266 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

267 return self._with_op("__xor__", other, f"{self_repr} ^ {other_repr}") 

268 

269 def __rxor__(self, other: Any) -> Expression: 

270 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

271 return self._with_op("__rxor__", other, f"{other_repr} ^ {self_repr}") 

272 

273 def __invert__(self) -> Expression: 

274 return Expression( 

275 lambda df: ~self._eval_expression(df), 

276 f"~{self._repr_str}", 

277 needs_parenthese=True, 

278 ) 

279 

280 def __neg__(self) -> Expression: 

281 if self._needs_parentheses: 

282 repr_str = f"-({self._repr_str})" 

283 else: 

284 repr_str = f"-{self._repr_str}" 

285 return Expression( 

286 lambda df: -self._eval_expression(df), 

287 repr_str, 

288 needs_parenthese=True, 

289 ) 

290 

291 def __pos__(self) -> Expression: 

292 if self._needs_parentheses: 

293 repr_str = f"+({self._repr_str})" 

294 else: 

295 repr_str = f"+{self._repr_str}" 

296 return Expression( 

297 lambda df: +self._eval_expression(df), 

298 repr_str, 

299 needs_parenthese=True, 

300 ) 

301 

302 def __abs__(self) -> Expression: 

303 return Expression( 

304 lambda df: abs(self._eval_expression(df)), 

305 f"abs({self._repr_str})", 

306 needs_parenthese=True, 

307 ) 

308 

309 def __array_ufunc__( 

310 self, ufunc: Callable[..., Any], method: str, *inputs: Any, **kwargs: Any 

311 ) -> Expression: 

312 def func(df: DataFrame) -> Any: 

313 parsed_inputs = _parse_args(df, *inputs) 

314 parsed_kwargs = _parse_kwargs(df, *kwargs) 

315 return ufunc(*parsed_inputs, **parsed_kwargs) 

316 

317 args_str = _pretty_print_args_kwargs(*inputs, **kwargs) 

318 repr_str = f"{ufunc.__name__}({args_str})" 

319 

320 return Expression(func, repr_str) 

321 

322 def __getitem__(self, item: Any) -> Expression: 

323 return self._with_op( 

324 "__getitem__", item, f"{self!r}[{item!r}]", needs_parentheses=True 

325 ) 

326 

327 def _call_with_func(self, func: Callable[..., Any], **kwargs: Any) -> Expression: 

328 def wrapped(df: DataFrame) -> Any: 

329 parsed_kwargs = _parse_kwargs(df, **kwargs) 

330 return func(**parsed_kwargs) 

331 

332 args_str = _pretty_print_args_kwargs(**kwargs) 

333 repr_str = func.__name__ + "(" + args_str + ")" 

334 

335 return Expression(wrapped, repr_str) 

336 

337 def __call__(self, *args: Any, **kwargs: Any) -> Expression: 

338 def func(df: DataFrame, *args: Any, **kwargs: Any) -> Any: 

339 parsed_args = _parse_args(df, *args) 

340 parsed_kwargs = _parse_kwargs(df, **kwargs) 

341 return self._eval_expression(df)(*parsed_args, **parsed_kwargs) 

342 

343 args_str = _pretty_print_args_kwargs(*args, **kwargs) 

344 repr_str = f"{self._repr_str}({args_str})" 

345 return Expression(lambda df: func(df, *args, **kwargs), repr_str) 

346 

347 def __getattr__(self, name: str, /) -> Any: 

348 repr_str = f"{self!r}" 

349 if self._needs_parentheses: 

350 repr_str = f"({repr_str})" 

351 repr_str += f".{name}" 

352 return Expression(lambda df: getattr(self._eval_expression(df), name), repr_str) 

353 

354 def case_when(self, caselist: Sequence[tuple[Any, Any]]) -> Expression: 

355 """ 

356 Create an expression that evaluates :meth:`Series.case_when` in a DataFrame 

357 context. 

358 

359 This is intended to enable patterns like:: 

360 

361 df.assign(result=pd.col("a").case_when([(pd.col("b") > 0, 1)])) 

362 

363 where conditions/replacements may reference other columns via ``pd.col``. 

364 """ 

365 

366 def func(df: DataFrame) -> Any: 

367 ser = self._eval_expression(df) 

368 evaluated = [] 

369 for condition, replacement in caselist: 

370 if isinstance(condition, Expression): 

371 condition = condition._eval_expression(df) 

372 if isinstance(replacement, Expression): 

373 replacement = replacement._eval_expression(df) 

374 evaluated.append((condition, replacement)) 

375 return ser.case_when(evaluated) 

376 

377 # Keep repr compact; caselist may be large. 

378 repr_str = f"{self!r}.case_when(...)" 

379 return Expression(func, repr_str) 

380 

381 def __repr__(self) -> str: 

382 return self._repr_str or "Expr(...)" 

383 

384 # Unsupported ops 

385 def __bool__(self) -> NoReturn: 

386 raise TypeError("boolean value of an expression is ambiguous") 

387 

388 def __iter__(self) -> NoReturn: 

389 raise TypeError("Expression objects are not iterable") 

390 

391 def __copy__(self) -> NoReturn: 

392 raise TypeError("Expression objects are not copiable") 

393 

394 def __deepcopy__(self, memo: dict[int, Any] | None) -> NoReturn: 

395 raise TypeError("Expression objects are not copiable") 

396 

397 

398@set_module("pandas") 

399def col(col_name: Hashable) -> Expression: 

400 """ 

401 Generate deferred object representing a column of a DataFrame. 

402 

403 Any place which accepts ``lambda df: df[col_name]``, such as 

404 :meth:`DataFrame.assign` or :meth:`DataFrame.loc`, can also accept 

405 ``pd.col(col_name)``. 

406 

407 .. versionadded:: 3.0.0 

408 

409 Parameters 

410 ---------- 

411 col_name : Hashable 

412 Column name. 

413 

414 Returns 

415 ------- 

416 :class:`pandas.api.typing.Expression` 

417 A deferred object representing a column of a DataFrame. 

418 

419 See Also 

420 -------- 

421 DataFrame.query : Query columns of a dataframe using string expressions. 

422 DataFrame.select : Select a subset of columns from the DataFrame. 

423 

424 Examples 

425 -------- 

426 

427 You can use `col` in `assign`. 

428 

429 >>> df = pd.DataFrame({"name": ["beluga", "narwhal"], "speed": [100, 110]}) 

430 >>> df.assign(name_titlecase=pd.col("name").str.title()) 

431 name speed name_titlecase 

432 0 beluga 100 Beluga 

433 1 narwhal 110 Narwhal 

434 

435 You can also use it for filtering. 

436 

437 >>> df.loc[pd.col("speed") > 105] 

438 name speed 

439 1 narwhal 110 

440 """ 

441 if not isinstance(col_name, Hashable): 

442 msg = f"Expected Hashable, got: {type(col_name)}" 

443 raise TypeError(msg) 

444 

445 def func(df: DataFrame) -> Series: 

446 if col_name not in df.columns: 

447 columns_str = str(df.columns.tolist()) 

448 max_len = 90 

449 if len(columns_str) > max_len: 

450 columns_str = columns_str[:max_len] + "...]" 

451 

452 msg = ( 

453 f"Column '{col_name}' not found in given DataFrame.\n\n" 

454 f"Hint: did you mean one of {columns_str} instead?" 

455 ) 

456 raise ValueError(msg) 

457 return df[col_name] 

458 

459 return Expression(func, f"col({col_name!r})") 

460 

461 

462__all__ = ["Expression", "col"]