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)
24import pandas.core.common as com
25from pandas.core.computation.common import (
26 ensure_decoded,
27 result_type_many,
28)
29from pandas.core.computation.scope import DEFAULT_GLOBALS
31from pandas.io.formats.printing import (
32 pprint_thing,
33 pprint_thing_encoded,
34)
36if TYPE_CHECKING:
37 from collections.abc import (
38 Callable,
39 Iterable,
40 Iterator,
41 )
43REDUCTIONS = ("sum", "prod", "min", "max")
45_unary_math_ops = (
46 "sin",
47 "cos",
48 "tan",
49 "exp",
50 "log",
51 "expm1",
52 "log1p",
53 "sqrt",
54 "sinh",
55 "cosh",
56 "tanh",
57 "arcsin",
58 "arccos",
59 "arctan",
60 "arccosh",
61 "arcsinh",
62 "arctanh",
63 "abs",
64 "log10",
65 "floor",
66 "ceil",
67)
68_binary_math_ops = ("arctan2",)
70MATHOPS = _unary_math_ops + _binary_math_ops
73LOCAL_TAG = "__pd_eval_local_"
76class Term:
77 def __new__(cls, name, env, side=None, encoding=None):
78 klass = Constant if not isinstance(name, str) else cls
79 supr_new = super(Term, klass).__new__
80 return supr_new(klass)
82 is_local: bool
84 def __init__(self, name, env, side=None, encoding=None) -> None:
85 # name is a str for Term, but may be something else for subclasses
86 self._name = name
87 self.env = env
88 self.side = side
89 tname = str(name)
90 self.is_local = tname.startswith(LOCAL_TAG) or tname in DEFAULT_GLOBALS
91 self._value = self._resolve_name()
92 self.encoding = encoding
94 @property
95 def local_name(self) -> str:
96 return self.name.replace(LOCAL_TAG, "")
98 def __repr__(self) -> str:
99 return pprint_thing(self.name)
101 def __call__(self, *args, **kwargs):
102 return self.value
104 def evaluate(self, *args, **kwargs) -> Term:
105 return self
107 def _resolve_name(self):
108 local_name = str(self.local_name)
109 is_local = self.is_local
110 if local_name in self.env.scope and isinstance(
111 self.env.scope[local_name], type
112 ):
113 is_local = False
115 res = self.env.resolve(local_name, is_local=is_local)
116 self.update(res)
118 if hasattr(res, "ndim") and isinstance(res.ndim, int) and res.ndim > 2:
119 raise NotImplementedError(
120 "N-dimensional objects, where N > 2, are not supported with eval"
121 )
122 return res
124 def update(self, value) -> None:
125 """
126 search order for local (i.e., @variable) variables:
128 scope, key_variable
129 [('locals', 'local_name'),
130 ('globals', 'local_name'),
131 ('locals', 'key'),
132 ('globals', 'key')]
133 """
134 key = self.name
136 # if it's a variable name (otherwise a constant)
137 if isinstance(key, str):
138 self.env.swapkey(self.local_name, key, new_value=value)
140 self.value = value
142 @property
143 def is_scalar(self) -> bool:
144 return is_scalar(self._value)
146 @property
147 def type(self):
148 try:
149 # potentially very slow for large, mixed dtype frames
150 return self._value.values.dtype
151 except AttributeError:
152 try:
153 # ndarray
154 return self._value.dtype
155 except AttributeError:
156 # scalar
157 return type(self._value)
159 return_type = type
161 @property
162 def raw(self) -> str:
163 return f"{type(self).__name__}(name={self.name!r}, type={self.type})"
165 @property
166 def is_datetime(self) -> bool:
167 try:
168 t = self.type.type
169 except AttributeError:
170 t = self.type
172 return issubclass(t, (datetime, np.datetime64))
174 @property
175 def value(self):
176 return self._value
178 @value.setter
179 def value(self, new_value) -> None:
180 self._value = new_value
182 @property
183 def name(self):
184 return self._name
186 @property
187 def ndim(self) -> int:
188 return self._value.ndim
191class Constant(Term):
192 def _resolve_name(self):
193 return self._name
195 @property
196 def name(self):
197 return self.value
199 def __repr__(self) -> str:
200 # in python 2 str() of float
201 # can truncate shorter than repr()
202 return repr(self.name)
205_bool_op_map = {"not": "~", "and": "&", "or": "|"}
208class Op:
209 """
210 Hold an operator of arbitrary arity.
211 """
213 op: str
215 def __init__(self, op: str, operands: Iterable[Term | Op], encoding=None) -> None:
216 self.op = _bool_op_map.get(op, op)
217 self.operands = operands
218 self.encoding = encoding
220 def __iter__(self) -> Iterator:
221 return iter(self.operands)
223 def __repr__(self) -> str:
224 """
225 Print a generic n-ary operator and its operands using infix notation.
226 """
227 # recurse over the operands
228 parened = (f"({pprint_thing(opr)})" for opr in self.operands)
229 return pprint_thing(f" {self.op} ".join(parened))
231 @property
232 def return_type(self):
233 # clobber types to bool if the op is a boolean operator
234 if self.op in (CMP_OPS_SYMS + BOOL_OPS_SYMS):
235 return np.bool_
236 return result_type_many(*(term.type for term in com.flatten(self)))
238 @property
239 def has_invalid_return_type(self) -> bool:
240 types = self.operand_types
241 obj_dtype_set = frozenset([np.dtype("object")])
242 return self.return_type == object and types - obj_dtype_set
244 @property
245 def operand_types(self):
246 return frozenset(term.type for term in com.flatten(self))
248 @property
249 def is_scalar(self) -> bool:
250 return all(operand.is_scalar for operand in self.operands)
252 @property
253 def is_datetime(self) -> bool:
254 try:
255 t = self.return_type.type
256 except AttributeError:
257 t = self.return_type
259 return issubclass(t, (datetime, np.datetime64))
262def _in(x, y):
263 """
264 Compute the vectorized membership of ``x in y`` if possible, otherwise
265 use Python.
266 """
267 try:
268 return x.isin(y)
269 except AttributeError:
270 if is_list_like(x):
271 try:
272 return y.isin(x)
273 except AttributeError:
274 pass
275 return x in y
278def _not_in(x, y):
279 """
280 Compute the vectorized membership of ``x not in y`` if possible,
281 otherwise use Python.
282 """
283 try:
284 return ~x.isin(y)
285 except AttributeError:
286 if is_list_like(x):
287 try:
288 return ~y.isin(x)
289 except AttributeError:
290 pass
291 return x not in y
294CMP_OPS_SYMS = (">", "<", ">=", "<=", "==", "!=", "in", "not in")
295_cmp_ops_funcs = (
296 operator.gt,
297 operator.lt,
298 operator.ge,
299 operator.le,
300 operator.eq,
301 operator.ne,
302 _in,
303 _not_in,
304)
305_cmp_ops_dict = dict(zip(CMP_OPS_SYMS, _cmp_ops_funcs, strict=True))
307BOOL_OPS_SYMS = ("&", "|", "and", "or")
308_bool_ops_funcs = (operator.and_, operator.or_, operator.and_, operator.or_)
309_bool_ops_dict = dict(zip(BOOL_OPS_SYMS, _bool_ops_funcs, strict=True))
311ARITH_OPS_SYMS = ("+", "-", "*", "/", "**", "//", "%")
312_arith_ops_funcs = (
313 operator.add,
314 operator.sub,
315 operator.mul,
316 operator.truediv,
317 operator.pow,
318 operator.floordiv,
319 operator.mod,
320)
321_arith_ops_dict = dict(zip(ARITH_OPS_SYMS, _arith_ops_funcs, strict=True))
323_binary_ops_dict = {}
325for d in (_cmp_ops_dict, _bool_ops_dict, _arith_ops_dict):
326 _binary_ops_dict.update(d)
329def is_term(obj) -> bool:
330 return isinstance(obj, Term)
333class BinOp(Op):
334 """
335 Hold a binary operator and its operands.
337 Parameters
338 ----------
339 op : str
340 lhs : Term or Op
341 rhs : Term or Op
342 """
344 def __init__(self, op: str, lhs, rhs) -> None:
345 super().__init__(op, (lhs, rhs))
346 self.lhs = lhs
347 self.rhs = rhs
349 self._disallow_scalar_only_bool_ops()
351 self.convert_values()
353 try:
354 self.func = _binary_ops_dict[op]
355 except KeyError as err:
356 # has to be made a list for python3
357 keys = list(_binary_ops_dict.keys())
358 raise ValueError(
359 f"Invalid binary operator {op!r}, valid operators are {keys}"
360 ) from err
362 def __call__(self, env):
363 """
364 Recursively evaluate an expression in Python space.
366 Parameters
367 ----------
368 env : Scope
370 Returns
371 -------
372 object
373 The result of an evaluated expression.
374 """
375 # recurse over the left/right nodes
376 left = self.lhs(env)
377 right = self.rhs(env)
379 return self.func(left, right)
381 def evaluate(self, env, engine: str, parser, term_type, eval_in_python):
382 """
383 Evaluate a binary operation *before* being passed to the engine.
385 Parameters
386 ----------
387 env : Scope
388 engine : str
389 parser : str
390 term_type : type
391 eval_in_python : list
393 Returns
394 -------
395 term_type
396 The "pre-evaluated" expression as an instance of ``term_type``
397 """
398 if engine == "python":
399 res = self(env)
400 else:
401 # recurse over the left/right nodes
403 left = self.lhs.evaluate(
404 env,
405 engine=engine,
406 parser=parser,
407 term_type=term_type,
408 eval_in_python=eval_in_python,
409 )
411 right = self.rhs.evaluate(
412 env,
413 engine=engine,
414 parser=parser,
415 term_type=term_type,
416 eval_in_python=eval_in_python,
417 )
419 # base cases
420 if self.op in eval_in_python:
421 res = self.func(left.value, right.value)
422 else:
423 from pandas.core.computation.eval import eval
425 res = eval(self, local_dict=env, engine=engine, parser=parser)
427 name = env.add_tmp(res)
428 return term_type(name, env=env)
430 def convert_values(self) -> None:
431 """
432 Convert datetimes to a comparable value in an expression.
433 """
435 def stringify(value):
436 encoder: Callable
437 if self.encoding is not None:
438 encoder = partial(pprint_thing_encoded, encoding=self.encoding)
439 else:
440 encoder = pprint_thing
441 return encoder(value)
443 lhs, rhs = self.lhs, self.rhs
445 if is_term(lhs) and lhs.is_datetime and is_term(rhs) and rhs.is_scalar:
446 v = rhs.value
447 if isinstance(v, (int, float)):
448 v = stringify(v)
449 v = Timestamp(ensure_decoded(v))
450 if v.tz is not None:
451 v = v.tz_convert("UTC")
452 self.rhs.update(v)
454 if is_term(rhs) and rhs.is_datetime and is_term(lhs) and lhs.is_scalar:
455 v = lhs.value
456 if isinstance(v, (int, float)):
457 v = stringify(v)
458 v = Timestamp(ensure_decoded(v))
459 if v.tz is not None:
460 v = v.tz_convert("UTC")
461 self.lhs.update(v)
463 def _disallow_scalar_only_bool_ops(self) -> None:
464 rhs = self.rhs
465 lhs = self.lhs
467 # GH#24883 unwrap dtype if necessary to ensure we have a type object
468 rhs_rt = rhs.return_type
469 rhs_rt = getattr(rhs_rt, "type", rhs_rt)
470 lhs_rt = lhs.return_type
471 lhs_rt = getattr(lhs_rt, "type", lhs_rt)
472 if (
473 (lhs.is_scalar or rhs.is_scalar)
474 and self.op in _bool_ops_dict
475 and (
476 not (
477 issubclass(rhs_rt, (bool, np.bool_))
478 and issubclass(lhs_rt, (bool, np.bool_))
479 )
480 )
481 ):
482 raise NotImplementedError("cannot evaluate scalar only bool ops")
485UNARY_OPS_SYMS = ("+", "-", "~", "not")
486_unary_ops_funcs = (operator.pos, operator.neg, operator.invert, operator.invert)
487_unary_ops_dict = dict(zip(UNARY_OPS_SYMS, _unary_ops_funcs, strict=True))
490class UnaryOp(Op):
491 """
492 Hold a unary operator and its operands.
494 Parameters
495 ----------
496 op : str
497 The token used to represent the operator.
498 operand : Term or Op
499 The Term or Op operand to the operator.
501 Raises
502 ------
503 ValueError
504 * If no function associated with the passed operator token is found.
505 """
507 def __init__(self, op: Literal["+", "-", "~", "not"], operand) -> None:
508 super().__init__(op, (operand,))
509 self.operand = operand
511 try:
512 self.func = _unary_ops_dict[op]
513 except KeyError as err:
514 raise ValueError(
515 f"Invalid unary operator {op!r}, valid operators are {UNARY_OPS_SYMS}"
516 ) from err
518 def __call__(self, env) -> MathCall:
519 operand = self.operand(env)
520 # error: Cannot call function of unknown type
521 return self.func(operand) # type: ignore[operator]
523 def __repr__(self) -> str:
524 return pprint_thing(f"{self.op}({self.operand})")
526 @property
527 def return_type(self) -> np.dtype:
528 operand = self.operand
529 if operand.return_type == np.dtype("bool"):
530 return np.dtype("bool")
531 if isinstance(operand, Op) and (
532 operand.op in _cmp_ops_dict or operand.op in _bool_ops_dict
533 ):
534 return np.dtype("bool")
535 return np.dtype("int")
538class MathCall(Op):
539 def __init__(self, func, args) -> None:
540 super().__init__(func.name, args)
541 self.func = func
543 def __call__(self, env):
544 # error: "Op" not callable
545 operands = [op(env) for op in self.operands] # type: ignore[operator]
546 return self.func.func(*operands)
548 def __repr__(self) -> str:
549 operands = map(str, self.operands)
550 return pprint_thing(f"{self.op}({','.join(operands)})")
553class FuncNode:
554 def __init__(self, name: str) -> None:
555 if name not in MATHOPS:
556 raise ValueError(f'"{name}" is not a supported function')
557 self.name = name
558 self.func = getattr(np, name)
560 def __call__(self, *args) -> MathCall:
561 return MathCall(self, args)