engines.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. """
  2. Engine classes for :func:`~pandas.eval`
  3. """
  4. from __future__ import annotations
  5. import abc
  6. from typing import TYPE_CHECKING
  7. from pandas.errors import NumExprClobberingError
  8. from pandas.core.computation.align import (
  9. align_terms,
  10. reconstruct_object,
  11. )
  12. from pandas.core.computation.ops import (
  13. MATHOPS,
  14. REDUCTIONS,
  15. )
  16. from pandas.io.formats import printing
  17. if TYPE_CHECKING:
  18. from pandas.core.computation.expr import Expr
  19. _ne_builtins = frozenset(MATHOPS + REDUCTIONS)
  20. def _check_ne_builtin_clash(expr: Expr) -> None:
  21. """
  22. Attempt to prevent foot-shooting in a helpful way.
  23. Parameters
  24. ----------
  25. expr : Expr
  26. Terms can contain
  27. """
  28. names = expr.names
  29. overlap = names & _ne_builtins
  30. if overlap:
  31. s = ", ".join([repr(x) for x in overlap])
  32. raise NumExprClobberingError(
  33. f'Variables in expression "{expr}" overlap with builtins: ({s})'
  34. )
  35. class AbstractEngine(metaclass=abc.ABCMeta):
  36. """Object serving as a base class for all engines."""
  37. has_neg_frac = False
  38. def __init__(self, expr) -> None:
  39. self.expr = expr
  40. self.aligned_axes = None
  41. self.result_type = None
  42. def convert(self) -> str:
  43. """
  44. Convert an expression for evaluation.
  45. Defaults to return the expression as a string.
  46. """
  47. return printing.pprint_thing(self.expr)
  48. def evaluate(self) -> object:
  49. """
  50. Run the engine on the expression.
  51. This method performs alignment which is necessary no matter what engine
  52. is being used, thus its implementation is in the base class.
  53. Returns
  54. -------
  55. object
  56. The result of the passed expression.
  57. """
  58. if not self._is_aligned:
  59. self.result_type, self.aligned_axes = align_terms(self.expr.terms)
  60. # make sure no names in resolvers and locals/globals clash
  61. res = self._evaluate()
  62. return reconstruct_object(
  63. self.result_type, res, self.aligned_axes, self.expr.terms.return_type
  64. )
  65. @property
  66. def _is_aligned(self) -> bool:
  67. return self.aligned_axes is not None and self.result_type is not None
  68. @abc.abstractmethod
  69. def _evaluate(self):
  70. """
  71. Return an evaluated expression.
  72. Parameters
  73. ----------
  74. env : Scope
  75. The local and global environment in which to evaluate an
  76. expression.
  77. Notes
  78. -----
  79. Must be implemented by subclasses.
  80. """
  81. class NumExprEngine(AbstractEngine):
  82. """NumExpr engine class"""
  83. has_neg_frac = True
  84. def _evaluate(self):
  85. import numexpr as ne
  86. # convert the expression to a valid numexpr expression
  87. s = self.convert()
  88. env = self.expr.env
  89. scope = env.full_scope
  90. _check_ne_builtin_clash(self.expr)
  91. return ne.evaluate(s, local_dict=scope)
  92. class PythonEngine(AbstractEngine):
  93. """
  94. Evaluate an expression in Python space.
  95. Mostly for testing purposes.
  96. """
  97. has_neg_frac = False
  98. def evaluate(self):
  99. return self.expr()
  100. def _evaluate(self) -> None:
  101. pass
  102. ENGINES: dict[str, type[AbstractEngine]] = {
  103. "numexpr": NumExprEngine,
  104. "python": PythonEngine,
  105. }