engines.py 3.7 KB

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