ENH: allow in-line expression assignment with df.eval

TST: tests for local name overlaps

ENH: moved assign to visit_Assign from visit_Module

Author: jreback

doc/source/release.rst

@@ -72,7 +72,8 @@ Experimental Features
     ``numexpr`` behind the scenes. This results in large speedups for complicated
     expressions involving large DataFrames/Series.
   - :class:`~pandas.DataFrame` has a new :meth:`~pandas.DataFrame.eval` that
-    evaluates an expression in the context of the ``DataFrame``.
+    evaluates an expression in the context of the ``DataFrame``; allows
+    inline expression assignment
   - A :meth:`~pandas.DataFrame.query` method has been added that allows
     you to select elements of a ``DataFrame`` using a natural query syntax nearly
     identical to Python syntax.

pandas/computation/eval.py

@@ -203,4 +203,10 @@ def eval(expr, parser='pandas', engine='numexpr', truediv=True,
     eng = _engines[engine]
     eng_inst = eng(parsed_expr)
     ret = eng_inst.evaluate()
+
+    # assign if needed
+    if parsed_expr.assignee is not None and parsed_expr.assigner is not None:
+        parsed_expr.assignee[parsed_expr.assigner] = ret
+        return None
+
     return ret

pandas/computation/expr.py

@@ -21,6 +21,7 @@
                                     _arith_ops_syms, _unary_ops_syms, is_term)
 from pandas.computation.ops import _reductions, _mathops, _LOCAL_TAG
 from pandas.computation.ops import Op, BinOp, UnaryOp, Term, Constant, Div
+from pandas.computation.ops import UndefinedVariableError
 
 
 def _ensure_scope(level=2, global_dict=None, local_dict=None, resolvers=None,
@@ -417,6 +418,8 @@ def __init__(self, env, engine, parser, preparser=_preparse):
         self.engine = engine
         self.parser = parser
         self.preparser = preparser
+        self.assignee = None
+        self.assigner = None
 
     def visit(self, node, **kwargs):
         if isinstance(node, string_types):
@@ -575,9 +578,39 @@ def visit_Slice(self, node, **kwargs):
         return slice(lower, upper, step)
 
     def visit_Assign(self, node, **kwargs):
-        cmpr = ast.Compare(ops=[ast.Eq()], left=node.targets[0],
-                           comparators=[node.value])
-        return self.visit(cmpr)
+        """
+        support a single assignment node, like
+
+        c = a + b
+
+        set the assignee at the top level, must be a Name node which
+        might or might not exist in the resolvers
+
+        """
+
+        if len(node.targets) != 1:
+            raise SyntaxError('can only assign a single expression')
+        if not isinstance(node.targets[0], ast.Name):
+            raise SyntaxError('left hand side of an assignment must be a single name')
+
+        # we have no one to assign to
+        if not len(self.env.resolvers):
+            raise NotImplementedError
+
+        try:
+            assigner = self.visit(node.targets[0], **kwargs)
+        except (UndefinedVariableError):
+            assigner = node.targets[0].id
+
+        self.assigner = getattr(assigner,'name',assigner)
+        if self.assigner is None:
+            raise SyntaxError('left hand side of an assignment must be a single resolvable name')
+        try:
+            self.assignee = self.env.resolvers[0]
+        except:
+            raise ValueError('cannot create an assignee for this expression')
+
+        return self.visit(node.value, **kwargs)
 
     def visit_Attribute(self, node, **kwargs):
         attr = node.attr
@@ -669,7 +702,7 @@ def visitor(x, y):
         return reduce(visitor, operands)
 
 
-_python_not_supported = frozenset(['Assign', 'Dict', 'Call', 'BoolOp',
+_python_not_supported = frozenset(['Dict', 'Call', 'BoolOp',
                                    'In', 'NotIn'])
 _numexpr_supported_calls = frozenset(_reductions + _mathops)
 
@@ -712,6 +745,14 @@ def __init__(self, expr, engine='numexpr', parser='pandas', env=None,
         self.terms = self.parse()
         self.truediv = truediv
 
+    @property
+    def assigner(self):
+        return getattr(self._visitor,'assigner',None)
+
+    @property
+    def assignee(self):
+        return getattr(self._visitor,'assignee',None)
+
     def __call__(self):
         self.env.locals['truediv'] = self.truediv
         return self.terms(self.env)

pandas/computation/pytables.py

@@ -389,6 +389,11 @@ def visit_USub(self, node, **kwargs):
     def visit_Index(self, node, **kwargs):
         return self.visit(node.value).value
 
+    def visit_Assign(self, node, **kwargs):
+        cmpr = ast.Compare(ops=[ast.Eq()], left=node.targets[0],
+                           comparators=[node.value])
+        return self.visit(cmpr)
+
     def visit_Subscript(self, node, **kwargs):
         value = self.visit(node.value)
         slobj = self.visit(node.slice)

pandas/computation/tests/test_eval.py

@@ -24,6 +24,7 @@
 from pandas.computation.ops import (_binary_ops_dict, _unary_ops_dict,
                                     _special_case_arith_ops_syms,
                                     _arith_ops_syms, _bool_ops_syms)
+from pandas.computation.common import NameResolutionError
 import pandas.computation.expr as expr
 import pandas.util.testing as tm
 from pandas.util.testing import (assert_frame_equal, randbool,
@@ -1154,6 +1155,45 @@ def test_assignment_fails(self):
         self.assertRaises(NotImplementedError, self.eval, expr1,
                           local_dict={'df': df, 'df2': df2})
 
+    def test_assignment_column(self):
+        df = DataFrame(np.random.randn(5, 2), columns=list('ab'))
+        orig_df = df.copy()
+
+        # multiple assignees
+        self.assertRaises(SyntaxError, df.eval, 'd c = a + b')
+
+        # invalid assignees
+        self.assertRaises(SyntaxError, df.eval, 'd,c = a + b')
+        self.assertRaises(SyntaxError, df.eval, 'Timestamp("20131001") = a + b')
+
+        # single assignment - existing variable
+        expected = orig_df.copy()
+        expected['a'] = expected['a'] + expected['b']
+        df = orig_df.copy()
+        df.eval('a = a + b')
+        assert_frame_equal(df,expected)
+
+        # single assignment - new variable
+        expected = orig_df.copy()
+        expected['c'] = expected['a'] + expected['b']
+        df = orig_df.copy()
+        df.eval('c = a + b')
+        assert_frame_equal(df,expected)
+
+        # with a local name overlap
+        a = 1
+        df = orig_df.copy()
+        df.eval('a = 1 + b')
+        expected = orig_df.copy()
+        expected['a'] = 1 + expected['b']
+        assert_frame_equal(df,expected)
+
+        df = orig_df.copy()
+        def f():
+            a = 1
+            df.eval('a=a+b')
+        self.assertRaises(NameResolutionError, f)
+
     def test_basic_period_index_boolean_expression(self):
         df = mkdf(2, 2, data_gen_f=f, c_idx_type='p', r_idx_type='i')