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 column. 

84 

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

86 """ 

87 

88 def __init__( 

89 self, 

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

91 repr_str: str, 

92 needs_parenthese: bool = False, 

93 ) -> None: 

94 self._func = func 

95 self._repr_str = repr_str 

96 self._needs_parentheses = needs_parenthese 

97 

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

99 return self._func(df) 

100 

101 def _with_op( 

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

103 ) -> Expression: 

104 if isinstance(other, Expression): 

105 return Expression( 

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

107 other._eval_expression(df) 

108 ), 

109 repr_str, 

110 needs_parenthese=needs_parentheses, 

111 ) 

112 else: 

113 return Expression( 

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

115 repr_str, 

116 needs_parenthese=needs_parentheses, 

117 ) 

118 

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

120 if self._needs_parentheses: 

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

122 else: 

123 self_repr = f"{self!r}" 

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

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

126 else: 

127 other_repr = f"{other!r}" 

128 return self_repr, other_repr 

129 

130 # Binary ops 

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

132 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

134 

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

136 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

138 

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

140 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

142 

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

144 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

146 

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

148 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

150 

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

152 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

154 

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

156 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

158 

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

160 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

162 

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

164 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

166 

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

168 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

170 

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

172 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

174 

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

176 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

178 

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

180 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

182 

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

184 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

186 

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

188 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

190 

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

192 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

194 

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

196 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

198 

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

200 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

202 

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

204 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

206 

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

208 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

210 

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

212 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

214 

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

216 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

218 

219 # Logical ops 

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

221 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

223 

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

225 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

227 

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

229 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

231 

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

233 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

235 

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

237 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

239 

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

241 self_repr, other_repr = self._maybe_wrap_parentheses(other) 

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

243 

244 def __invert__(self) -> Expression: 

245 return Expression( 

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

247 f"~{self._repr_str}", 

248 needs_parenthese=True, 

249 ) 

250 

251 def __neg__(self) -> Expression: 

252 if self._needs_parentheses: 

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

254 else: 

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

256 return Expression( 

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

258 repr_str, 

259 needs_parenthese=True, 

260 ) 

261 

262 def __pos__(self) -> Expression: 

263 if self._needs_parentheses: 

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

265 else: 

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

267 return Expression( 

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

269 repr_str, 

270 needs_parenthese=True, 

271 ) 

272 

273 def __abs__(self) -> Expression: 

274 return Expression( 

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

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

277 needs_parenthese=True, 

278 ) 

279 

280 def __array_ufunc__( 

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

282 ) -> Expression: 

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

284 parsed_inputs = _parse_args(df, *inputs) 

285 parsed_kwargs = _parse_kwargs(df, *kwargs) 

286 return ufunc(*parsed_inputs, **parsed_kwargs) 

287 

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

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

290 

291 return Expression(func, repr_str) 

292 

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

294 return self._with_op( 

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

296 ) 

297 

298 def _call_with_func(self, func: Callable, **kwargs: Any) -> Expression: 

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

300 parsed_kwargs = _parse_kwargs(df, **kwargs) 

301 return func(**parsed_kwargs) 

302 

303 args_str = _pretty_print_args_kwargs(**kwargs) 

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

305 

306 return Expression(wrapped, repr_str) 

307 

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

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

310 parsed_args = _parse_args(df, *args) 

311 parsed_kwargs = _parse_kwargs(df, **kwargs) 

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

313 

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

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

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

317 

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

319 repr_str = f"{self!r}" 

320 if self._needs_parentheses: 

321 repr_str = f"({repr_str})" 

322 repr_str += f".{name}" 

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

324 

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

326 """ 

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

328 context. 

329 

330 This is intended to enable patterns like:: 

331 

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

333 

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

335 """ 

336 

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

338 ser = self._eval_expression(df) 

339 evaluated = [] 

340 for condition, replacement in caselist: 

341 if isinstance(condition, Expression): 

342 condition = condition._eval_expression(df) 

343 if isinstance(replacement, Expression): 

344 replacement = replacement._eval_expression(df) 

345 evaluated.append((condition, replacement)) 

346 return ser.case_when(evaluated) 

347 

348 # Keep repr compact; caselist may be large. 

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

350 return Expression(func, repr_str) 

351 

352 def __repr__(self) -> str: 

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

354 

355 # Unsupported ops 

356 def __bool__(self) -> NoReturn: 

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

358 

359 def __iter__(self) -> NoReturn: 

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

361 

362 def __copy__(self) -> NoReturn: 

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

364 

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

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

367 

368 

369@set_module("pandas") 

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

371 """ 

372 Generate deferred object representing a column of a DataFrame. 

373 

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

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

376 ``pd.col(col_name)``. 

377 

378 .. versionadded:: 3.0.0 

379 

380 Parameters 

381 ---------- 

382 col_name : Hashable 

383 Column name. 

384 

385 Returns 

386 ------- 

387 `pandas.api.typing.Expression` 

388 A deferred object representing a column of a DataFrame. 

389 

390 See Also 

391 -------- 

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

393 

394 Examples 

395 -------- 

396 

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

398 

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

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

401 name speed name_titlecase 

402 0 beluga 100 Beluga 

403 1 narwhal 110 Narwhal 

404 

405 You can also use it for filtering. 

406 

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

408 name speed 

409 1 narwhal 110 

410 """ 

411 if not isinstance(col_name, Hashable): 

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

413 raise TypeError(msg) 

414 

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

416 if col_name not in df.columns: 

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

418 max_len = 90 

419 if len(columns_str) > max_len: 

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

421 

422 msg = ( 

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

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

425 ) 

426 raise ValueError(msg) 

427 return df[col_name] 

428 

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

430 

431 

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