Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Operator classes for eval.
3"""
5from __future__ import annotations
7from datetime import datetime
8from functools import partial
9import operator
10from typing import (
11 TYPE_CHECKING,
12 Literal,
13)
15import numpy as np
17from pandas._libs.tslibs import Timestamp
19from pandas.core.dtypes.common import (
20 is_list_like,
21 is_scalar,
22)
23from pandas.core.dtypes.dtypes import (
24 IntervalDtype,
25 PeriodDtype,
26)
28import pandas.core.common as com
29from pandas.core.computation.common import (
30 ensure_decoded,
31 result_type_many,
32)
33from pandas.core.computation.scope import DEFAULT_GLOBALS
35from pandas.io.formats.printing import (
36 pprint_thing,
37 pprint_thing_encoded,
38)
40if TYPE_CHECKING:
41 from collections.abc import (
42 Callable,
43 Iterable,
44 Iterator,
45 )
47REDUCTIONS = ("sum", "prod", "min", "max")
49_unary_math_ops = (
50 "sin",
51 "cos",
52 "tan",
53 "exp",
54 "log",
55 "expm1",
56 "log1p",
57 "sqrt",
58 "sinh",
59 "cosh",
60 "tanh",
61 "arcsin",
62 "arccos",
63 "arctan",
64 "arccosh",
65 "arcsinh",
66 "arctanh",
67 "abs",
68 "log10",
69 "floor",
70 "ceil",
71)
72_binary_math_ops = ("arctan2",)
74MATHOPS = _unary_math_ops + _binary_math_ops
77LOCAL_TAG = "__pd_eval_local_"
80class Term:
81 def __new__(cls, name, env, side=None, encoding=None):
82 klass = Constant if not isinstance(name, str) else cls
83 supr_new = super(Term, klass).__new__
84 return supr_new(klass)
86 is_local: bool
88 def __init__(self, name, env, side=None, encoding=None) -> None:
89 # name is a str for Term, but may be something else for subclasses
90 self._name = name
91 self.env = env
92 self.side = side
93 tname = str(name)
94 self.is_local = tname.startswith(LOCAL_TAG) or tname in DEFAULT_GLOBALS
95 self._value = self._resolve_name()
96 self.encoding = encoding
98 @property
99 def local_name(self) -> str:
100 return self.name.replace(LOCAL_TAG, "")
102 def __repr__(self) -> str:
103 return pprint_thing(self.name)
105 def __call__(self, *args, **kwargs):
106 return self.value
108 def evaluate(self, *args, **kwargs) -> Term:
109 return self
111 def _resolve_name(self):
112 local_name = str(self.local_name)
113 is_local = self.is_local
114 if (
115 not str(self.name).startswith(LOCAL_TAG)
116 and local_name in self.env.scope
117 and isinstance(self.env.scope[local_name], type)
118 ):
119 is_local = False
121 res = self.env.resolve(local_name, is_local=is_local)
122 self.update(res)
124 if hasattr(res, "ndim") and isinstance(res.ndim, int) and res.ndim > 2:
125 raise NotImplementedError(
126 "N-dimensional objects, where N > 2, are not supported with eval"
127 )
128 return res
130 def update(self, value) -> None:
131 """
132 search order for local (i.e., @variable) variables:
134 scope, key_variable
135 [('locals', 'local_name'),
136 ('globals', 'local_name'),
137 ('locals', 'key'),
138 ('globals', 'key')]
139 """
140 key = self.name
142 # if it's a variable name (otherwise a constant)
143 if isinstance(key, str):
144 self.env.swapkey(self.local_name, key, new_value=value)
146 self.value = value
148 @property
149 def is_scalar(self) -> bool:
150 return is_scalar(self._value)
152 @property
153 def type(self):
154 try:
155 # GH#35247 avoid .values which boxes PeriodArray/IntervalArray
156 dtype = self._value.dtype
157 except AttributeError:
158 try:
159 return self._value.values.dtype # DataFrame
160 except AttributeError:
161 return type(self._value) # scalar
163 # Match .values.dtype behavior for types that external_values boxes
164 if isinstance(dtype, (PeriodDtype, IntervalDtype)):
165 return np.dtype("object")
166 return dtype
168 return_type = type
170 @property
171 def raw(self) -> str:
172 return f"{type(self).__name__}(name={self.name!r}, type={self.type})"
174 @property
175 def is_datetime(self) -> bool:
176 try:
177 t = self.type.type
178 except AttributeError:
179 t = self.type
181 return issubclass(t, (datetime, np.datetime64))
183 @property
184 def value(self):
185 return self._value
187 @value.setter
188 def value(self, new_value) -> None:
189 self._value = new_value
191 @property
192 def name(self):
193 return self._name
195 @property
196 def ndim(self) -> int:
197 return self._value.ndim
200class Constant(Term):
201 def _resolve_name(self):
202 return self._name
204 @property
205 def name(self):
206 return self.value
208 def __repr__(self) -> str:
209 return repr(self.name)
212_bool_op_map = {"not": "~", "and": "&", "or": "|"}
215class Op:
216 """
217 Hold an operator of arbitrary arity.
218 """
220 op: str
222 def __init__(self, op: str, operands: Iterable[Term | Op], encoding=None) -> None:
223 self.op = _bool_op_map.get(op, op)
224 self.operands = operands
225 self.encoding = encoding
227 def __iter__(self) -> Iterator[Term | Op]:
228 return iter(self.operands)
230 def __repr__(self) -> str:
231 """
232 Print a generic n-ary operator and its operands using infix notation.
233 """
234 # recurse over the operands
235 parened = (f"({pprint_thing(opr)})" for opr in self.operands)
236 return pprint_thing(f" {self.op} ".join(parened))
238 @property
239 def return_type(self):
240 # clobber types to bool if the op is a boolean operator
241 if self.op in (CMP_OPS_SYMS + BOOL_OPS_SYMS):
242 return np.bool_
243 return result_type_many(*(term.type for term in com.flatten(self)))
245 @property
246 def has_invalid_return_type(self) -> bool:
247 types = self.operand_types
248 obj_dtype_set = frozenset([np.dtype("object")])
249 return self.return_type == object and types - obj_dtype_set
251 @property
252 def operand_types(self):
253 return frozenset(term.type for term in com.flatten(self))
255 @property
256 def is_scalar(self) -> bool:
257 return all(operand.is_scalar for operand in self.operands)
259 @property
260 def is_datetime(self) -> bool:
261 try:
262 t = self.return_type.type
263 except AttributeError:
264 t = self.return_type
266 return issubclass(t, (datetime, np.datetime64))
269def _in(x, y):
270 """
271 Compute the vectorized membership of ``x in y`` if possible, otherwise
272 use Python.
273 """
274 try:
275 return x.isin(y)
276 except AttributeError:
277 if is_list_like(x):
278 try:
279 return y.isin(x)
280 except AttributeError:
281 pass
282 return x in y
285def _not_in(x, y):
286 """
287 Compute the vectorized membership of ``x not in y`` if possible,
288 otherwise use Python.
289 """
290 try:
291 return ~x.isin(y)
292 except AttributeError:
293 if is_list_like(x):
294 try:
295 return ~y.isin(x)
296 except AttributeError:
297 pass
298 return x not in y
301CMP_OPS_SYMS = (">", "<", ">=", "<=", "==", "!=", "in", "not in")
302_cmp_ops_funcs = (
303 operator.gt,
304 operator.lt,
305 operator.ge,
306 operator.le,
307 operator.eq,
308 operator.ne,
309 _in,
310 _not_in,
311)
312_cmp_ops_dict = dict(zip(CMP_OPS_SYMS, _cmp_ops_funcs, strict=True))
314BOOL_OPS_SYMS = ("&", "|", "and", "or")
315_bool_ops_funcs = (operator.and_, operator.or_, operator.and_, operator.or_)
316_bool_ops_dict = dict(zip(BOOL_OPS_SYMS, _bool_ops_funcs, strict=True))
318ARITH_OPS_SYMS = ("+", "-", "*", "/", "**", "//", "%")
319_arith_ops_funcs = (
320 operator.add,
321 operator.sub,
322 operator.mul,
323 operator.truediv,
324 operator.pow,
325 operator.floordiv,
326 operator.mod,
327)
328_arith_ops_dict = dict(zip(ARITH_OPS_SYMS, _arith_ops_funcs, strict=True))
330_binary_ops_dict = {}
332for d in (_cmp_ops_dict, _bool_ops_dict, _arith_ops_dict):
333 _binary_ops_dict.update(d)
336def is_term(obj) -> bool:
337 return isinstance(obj, Term)
340class BinOp(Op):
341 """
342 Hold a binary operator and its operands.
344 Parameters
345 ----------
346 op : str
347 lhs : Term or Op
348 rhs : Term or Op
349 """
351 def __init__(self, op: str, lhs, rhs) -> None:
352 super().__init__(op, (lhs, rhs))
353 self.lhs = lhs
354 self.rhs = rhs
356 self._disallow_scalar_only_bool_ops()
358 self.convert_values()
360 try:
361 self.func = _binary_ops_dict[op]
362 except KeyError as err:
363 # has to be made a list for python3
364 keys = list(_binary_ops_dict.keys())
365 raise ValueError(
366 f"Invalid binary operator {op!r}, valid operators are {keys}"
367 ) from err
369 def __call__(self, env):
370 """
371 Recursively evaluate an expression in Python space.
373 Parameters
374 ----------
375 env : Scope
377 Returns
378 -------
379 object
380 The result of an evaluated expression.
381 """
382 # recurse over the left/right nodes
383 left = self.lhs(env)
384 right = self.rhs(env)
386 return self.func(left, right)
388 def evaluate(self, env, engine: str, parser, term_type, eval_in_python):
389 """
390 Evaluate a binary operation *before* being passed to the engine.
392 Parameters
393 ----------
394 env : Scope
395 engine : str
396 parser : str
397 term_type : type
398 eval_in_python : list
400 Returns
401 -------
402 term_type
403 The "pre-evaluated" expression as an instance of ``term_type``
404 """
405 if engine == "python":
406 res = self(env)
407 else:
408 # recurse over the left/right nodes
410 left = self.lhs.evaluate(
411 env,
412 engine=engine,
413 parser=parser,
414 term_type=term_type,
415 eval_in_python=eval_in_python,
416 )
418 right = self.rhs.evaluate(
419 env,
420 engine=engine,
421 parser=parser,
422 term_type=term_type,
423 eval_in_python=eval_in_python,
424 )
426 # base cases
427 if self.op in eval_in_python:
428 res = self.func(left.value, right.value)
429 else:
430 from pandas.core.computation.eval import eval
432 res = eval(self, local_dict=env, engine=engine, parser=parser)
434 name = env.add_tmp(res)
435 return term_type(name, env=env)
437 def convert_values(self) -> None:
438 """
439 Convert datetimes to a comparable value in an expression.
440 """
442 def stringify(value):
443 encoder: Callable[[object], str | bytes]
444 if self.encoding is not None:
445 encoder = partial(pprint_thing_encoded, encoding=self.encoding)
446 else:
447 encoder = pprint_thing
448 return encoder(value)
450 lhs, rhs = self.lhs, self.rhs
452 if is_term(lhs) and lhs.is_datetime and is_term(rhs) and rhs.is_scalar:
453 v = rhs.value
454 if isinstance(v, (int, float)):
455 v = stringify(v)
456 v = Timestamp(ensure_decoded(v))
457 if v.tz is not None:
458 v = v.tz_convert("UTC")
459 self.rhs.update(v)
461 if is_term(rhs) and rhs.is_datetime and is_term(lhs) and lhs.is_scalar:
462 v = lhs.value
463 if isinstance(v, (int, float)):
464 v = stringify(v)
465 v = Timestamp(ensure_decoded(v))
466 if v.tz is not None:
467 v = v.tz_convert("UTC")
468 self.lhs.update(v)
470 def _disallow_scalar_only_bool_ops(self) -> None:
471 rhs = self.rhs
472 lhs = self.lhs
474 # GH#24883 unwrap dtype if necessary to ensure we have a type object
475 rhs_rt = rhs.return_type
476 rhs_rt = getattr(rhs_rt, "type", rhs_rt)
477 lhs_rt = lhs.return_type
478 lhs_rt = getattr(lhs_rt, "type", lhs_rt)
479 if (
480 (lhs.is_scalar or rhs.is_scalar)
481 and self.op in _bool_ops_dict
482 and (
483 not (
484 issubclass(rhs_rt, (bool, np.bool_))
485 and issubclass(lhs_rt, (bool, np.bool_))
486 )
487 )
488 ):
489 raise NotImplementedError("cannot evaluate scalar only bool ops")
492UNARY_OPS_SYMS = ("+", "-", "~", "not")
493_unary_ops_funcs = (operator.pos, operator.neg, operator.invert, operator.invert)
494_unary_ops_dict = dict(zip(UNARY_OPS_SYMS, _unary_ops_funcs, strict=True))
497class UnaryOp(Op):
498 """
499 Hold a unary operator and its operands.
501 Parameters
502 ----------
503 op : str
504 The token used to represent the operator.
505 operand : Term or Op
506 The Term or Op operand to the operator.
508 Raises
509 ------
510 ValueError
511 * If no function associated with the passed operator token is found.
512 """
514 def __init__(self, op: Literal["+", "-", "~", "not"], operand) -> None:
515 super().__init__(op, (operand,))
516 self.operand = operand
518 try:
519 self.func = _unary_ops_dict[op]
520 except KeyError as err:
521 raise ValueError(
522 f"Invalid unary operator {op!r}, valid operators are {UNARY_OPS_SYMS}"
523 ) from err
525 def __call__(self, env) -> MathCall:
526 operand = self.operand(env)
527 # error: Cannot call function of unknown type
528 return self.func(operand) # type: ignore[operator]
530 def __repr__(self) -> str:
531 return pprint_thing(f"{self.op}({self.operand})")
533 @property
534 def return_type(self) -> np.dtype:
535 operand = self.operand
536 if operand.return_type == np.dtype("bool"):
537 return np.dtype("bool")
538 if isinstance(operand, Op) and (
539 operand.op in _cmp_ops_dict or operand.op in _bool_ops_dict
540 ):
541 return np.dtype("bool")
542 return np.dtype("int")
545class MathCall(Op):
546 def __init__(self, func, args) -> None:
547 super().__init__(func.name, args)
548 self.func = func
550 def __call__(self, env):
551 # error: "Op" not callable
552 operands = [op(env) for op in self.operands] # type: ignore[operator]
553 return self.func.func(*operands)
555 def __repr__(self) -> str:
556 operands = map(str, self.operands)
557 return pprint_thing(f"{self.op}({','.join(operands)})")
560class FuncNode:
561 def __init__(self, name: str) -> None:
562 if name not in MATHOPS:
563 raise ValueError(f'"{name}" is not a supported function')
564 self.name = name
565 self.func = getattr(np, name)
567 def __call__(self, *args) -> MathCall:
568 return MathCall(self, args)