ENH: add boolean operators

Author: cpcloud

pandas/computation/expr.py

@@ -151,8 +151,8 @@ def _filter_nodes(superclass, all_nodes=_all_nodes):
  'IfExp', 'DictComp',
  'SetComp', 'Repr',
  'Lambda', 'Set', 'In',
- 'NotIn', 'AST',
- 'Is', 'IsNot'])) -
+ 'NotIn', 'AST', 'Is',
+ 'IsNot'])) -
  _hacked_nodes)
 
 # we're adding a different assignment in some cases to be equality comparison
@@ -211,12 +211,12 @@ class BaseExprVisitor(ast.NodeVisitor):
  """
  binary_ops = _cmp_ops_syms + _bool_ops_syms + _arith_ops_syms
  binary_op_nodes = ('Gt', 'Lt', 'GtE', 'LtE', 'Eq', 'NotEq', 'BitAnd',
- 'BitOr', 'Add', 'Sub', 'Mult', 'Div', 'Pow', 'FloorDiv',
- 'Mod')
+ 'BitOr', 'And', 'Or', 'Add', 'Sub', 'Mult', 'Div',
+ 'Pow', 'FloorDiv', 'Mod')
  binary_op_nodes_map = dict(itertools.izip(binary_ops, binary_op_nodes))
 
  unary_ops = _unary_ops_syms
- unary_op_nodes = 'UAdd', 'USub', 'Invert'
+ unary_op_nodes = 'UAdd', 'USub', 'Invert', 'Not'
  unary_op_nodes_map = dict(itertools.izip(unary_ops, unary_op_nodes))
 
  def __init__(self, env, preparser=_preparse):
@@ -354,13 +354,31 @@ def visit_Compare(self, node, **kwargs):
  self.visit(comp, side='right'))
  return node
 
+ def visit_BoolOp(self, node, **kwargs):
+ op = self.visit(node.op)
+ def visitor(x, y):
+ try:
+ lhs = self.visit(x)
+ except TypeError:
+ lhs = x
+
+ try:
+ rhs = self.visit(y)
+ except TypeError:
+ rhs = y
+
+ return op(lhs, rhs)
 
-_python_not_supported = frozenset(['Assign', 'BoolOp', 'Not', 'Str', 'Slice',
- 'Index', 'Subscript', 'Tuple', 'List',
- 'Dict', 'Call'])
+ operands = node.values
+ return reduce(visitor, operands)
+
+
+_python_not_supported = frozenset(['Assign', 'Str', 'Slice', 'Index',
+ 'Subscript', 'Tuple', 'List', 'Dict',
+ 'Call'])
 _numexpr_supported_calls = frozenset(_reductions + _mathops)
 
-@disallow(_unsupported_nodes | _python_not_supported)
+@disallow((_unsupported_nodes | _python_not_supported) - _boolop_nodes)
 class PandasExprVisitor(BaseExprVisitor):
  def __init__(self, env, preparser=_preparse):
  super(PandasExprVisitor, self).__init__(env, preparser)

pandas/computation/ops.py

@@ -107,11 +107,15 @@ def _print_operand(opr):
  return opr.name if is_term(opr) else unicode(opr)
 
 
+def _get_op(op):
+ return {'not': '~', 'and': '&', 'or': '|'}.get(op, op)
+
+
 class Op(StringMixin):
  """Hold an operator of unknown arity
  """
  def __init__(self, op, operands, *args, **kwargs):
- self.op = op
+ self.op = _get_op(op)
  self.operands = operands
 
  def __iter__(self):
@@ -137,8 +141,8 @@ def return_type(self):
 _cmp_ops_funcs = op.gt, op.lt, op.ge, op.le, op.eq, op.ne
 _cmp_ops_dict = dict(zip(_cmp_ops_syms, _cmp_ops_funcs))
 
-_bool_ops_syms = '&', '|'
-_bool_ops_funcs = op.and_, op.or_
+_bool_ops_syms = '&', '|', 'and', 'or'
+_bool_ops_funcs = op.and_, op.or_, op.and_, op.or_
 _bool_ops_dict = dict(zip(_bool_ops_syms, _bool_ops_funcs))
 
 _arith_ops_syms = '+', '-', '*', '/', '**', '//', '%'
@@ -237,8 +241,8 @@ def __init__(self, lhs, rhs, *args, **kwargs):
  _cast_inplace(self.operands, np.float_)
 
 
-_unary_ops_syms = '+', '-', '~'
-_unary_ops_funcs = op.pos, op.neg, op.invert
+_unary_ops_syms = '+', '-', '~', 'not'
+_unary_ops_funcs = op.pos, op.neg, op.invert, op.invert
 _unary_ops_dict = dict(zip(_unary_ops_syms, _unary_ops_funcs))
 
 

pandas/computation/pytables.py

@@ -340,7 +340,6 @@ def prune(self, klass):
 
 _op_classes = {'unary': UnaryOp}
 
-
 class ExprVisitor(BaseExprVisitor):
  def __init__(self, env, **kwargs):
  super(ExprVisitor, self).__init__(env)
@@ -366,6 +365,7 @@ def visit_USub(self, node, **kwargs):
  def visit_Index(self, node, **kwargs):
  return self.visit(node.value).value
 
+
 class Expr(expr.Expr):
 
  """ hold a pytables like expression, comprised of possibly multiple 'terms'

pandas/computation/tests/test_eval.py

@@ -23,7 +23,7 @@
 from pandas.computation.expr import PythonExprVisitor, PandasExprVisitor
 from pandas.computation.ops import (_binary_ops_dict, _unary_ops_dict,
  _special_case_arith_ops_syms,
- _arith_ops_syms)
+ _arith_ops_syms, Constant)
 import pandas.computation.expr as expr
 from pandas.computation import pytables
 from pandas.computation.expressions import _USE_NUMEXPR
@@ -599,39 +599,6 @@ def test_is_expr():
  check_is_expr(engine)
 
 
-def check_not_fails(engine):
- x = True
- assert_raises(NotImplementedError, pd.eval, 'not x', engine=engine,
- local_dict={'x': x})
-
-
-def test_not_fails():
- for engine in _engines:
- check_not_fails(engine)
-
-
-def check_and_fails(engine):
- x, y = False, True
- assert_raises(NotImplementedError, pd.eval, 'x and y', engine=engine,
- local_dict={'x': x, 'y': y})
-
-
-def test_and_fails():
- for engine in _engines:
- check_and_fails(engine)
-
-
-def check_or_fails(engine):
- x, y = True, False
- assert_raises(NotImplementedError, pd.eval, 'x or y', engine=engine,
- local_dict={'x': x, 'y': y})
-
-
-def test_or_fails():
- for engine in _engines:
- check_or_fails(engine)
-
-
 _parsers = {'python': PythonExprVisitor, 'pytables': pytables.ExprVisitor,
  'pandas': PandasExprVisitor}
 
@@ -641,8 +608,9 @@ def check_disallowed_nodes(visitor):
  VisitorClass = _parsers[visitor]
  uns_ops = VisitorClass.unsupported_nodes
  inst = VisitorClass('x + 1')
+
  for ops in uns_ops:
- assert_raises(NotImplementedError, getattr(inst, ops), inst, ast.AST())
+ assert_raises(NotImplementedError, getattr(inst, ops))
 
 
 def test_disallowed_nodes():