Ticket #113: ticket113.2.diff

genshi/template/tests/eval.py

@@ -379 +379 @@
             expr.evaluate({'something': Something()})
             self.fail('Expected UndefinedError')
         except UndefinedError, e:
+            self.assertEqual('<Something> has no member named "nil"', str(e))
             exc_type, exc_value, exc_traceback = sys.exc_info()
+            search_string = "<Expression 'something.nil'>"
             frame = exc_traceback.tb_next
-            frames = []
             while frame.tb_next:
                 frame = frame.tb_next
-                frames.append(frame)
-            self.assertEqual('<Something> has no member named "nil"', str(e))
-            self.assertEqual("<Expression 'something.nil'>",
-                             frames[-3].tb_frame.f_code.co_name)
-            self.assertEqual('index.html',
-                             frames[-3].tb_frame.f_code.co_filename)
-            self.assertEqual(50, frames[-3].tb_lineno)
+                code = frame.tb_frame.f_code
+                if code.co_name == search_string:
+                    break
+            else:
+                self.fail("never found the frame I was looking for")
+            self.assertEqual('index.html', code.co_filename)
+            self.assertEqual(50, frame.tb_lineno)
 
     def test_error_getitem_undefined_string(self):
         class Something(object):
@@ -402 +403 @@
             expr.evaluate({'something': Something()})
             self.fail('Expected UndefinedError')
         except UndefinedError, e:
+            self.assertEqual('<Something> has no member named "nil"', str(e))
             exc_type, exc_value, exc_traceback = sys.exc_info()
+            search_string = '''<Expression 'something["nil"]'>'''
             frame = exc_traceback.tb_next
-            frames = []
             while frame.tb_next:
                 frame = frame.tb_next
-                frames.append(frame)
-            self.assertEqual('<Something> has no member named "nil"', str(e))
-            self.assertEqual('''<Expression 'something["nil"]'>''',
-                             frames[-3].tb_frame.f_code.co_name)
-            self.assertEqual('index.html',
-                             frames[-3].tb_frame.f_code.co_filename)
-            self.assertEqual(50, frames[-3].tb_lineno)
+                code = frame.tb_frame.f_code
+                if code.co_name == search_string:
+                    break
+            else:
+                self.fail("never found the frame I was looking for")
+            self.assertEqual('index.html', code.co_filename)
+            self.assertEqual(50, frame.tb_lineno)
 
 
 class SuiteTestCase(unittest.TestCase):
@@ -431 +433 @@
         assert 'donothing' in data
         self.assertEqual(None, data['donothing']())
 
+    def test_def_with_multiple_statements(self):
+        suite = Suite("""def donothing():
+    if True:
+        return foo
+""")
+        data = {'foo': 'bar'}
+        suite.execute(data)
+        assert 'donothing' in data
+        self.assertEqual('bar', data['donothing']())
+
     def test_delete(self):
         suite = Suite("""foo = 42
 del foo
@@ -504 +516 @@
         suite.execute(data)
         self.assertEqual(4, data['x'])
 
+    def test_augmented_attribute_assignment(self):
+        suite = Suite("d['k'] += 42")
+        d = {"k": 1}
+        suite.execute({"d": d})
+        self.assertEqual(43, d["k"])
 
+    def test_local_augmented_assign(self):
+        Suite("x = 1; x += 42; assert x == 43").execute({})
+
+    def test_assign_in_list(self):
+        suite = Suite("[d['k']] = 'foo',; assert d['k'] == 'foo'")
+        d = {"k": "bar"}
+        suite.execute({"d": d})
+        self.assertEqual("foo", d["k"])
+
+    def test_exec(self):
+        suite = Suite("x = 1; exec d['k']; assert x == 42, x")
+        suite.execute({"d": {"k": "x = 42"}})
+
+    def test_return(self):
+        suite = Suite("""
+def f():
+    return v
+
+assert f() == 42
+""")
+        suite.execute({"v": 42})
+
+    def test_assign_to_dict_item(self):
+        suite = Suite("d['k'] = 'foo'")
+        data = {'d': {}}
+        suite.execute(data)
+        self.assertEqual('foo', data['d']['k'])
+
+    def test_assign_to_attribute(self):
+        class Something(object): pass
+        something = Something()
+        suite = Suite("obj.attr = 'foo'")
+        data = {"obj": something}
+        suite.execute(data)
+        self.assertEqual('foo', something.attr)
+
+    def test_delattr(self):
+        class Something(object):
+            def __init__(self):
+                self.attr = 'foo'
+        obj = Something()
+        Suite("del obj.attr").execute({'obj': obj})
+        self.failIf(hasattr(obj, 'attr'))
+
+    def test_delitem(self):
+        d = {'k': 'foo'}
+        Suite("del d['k']").execute({'d': d})
+        self.failIf('k' in d, `d`)
+
+
 def suite():
     suite = unittest.TestSuite()
     suite.addTest(doctest.DocTestSuite(Expression.__module__))

genshi/template/eval.py

@@ -140 +140 @@
     """Executes Python statements used in templates.
 
     >>> data = dict(test='Foo', items=[1, 2, 3], dict={'some': 'thing'})
-    >>> Suite('foo = dict.some').execute(data)
+    >>> Suite("foo = dict['some']").execute(data)
     >>> data['foo']
     'thing'
     """
@@ -361 +361 @@
     return parse(source, mode)
 
 def _compile(node, source=None, mode='eval', filename=None, lineno=-1):
-    tree = TemplateASTTransformer().visit(node)
+    xform = {'eval': ExpressionASTTransformer}.get(mode, TemplateASTTransformer)
+    tree = xform().visit(node)
     if isinstance(filename, unicode):
         # unicode file names not allowed for code objects
         filename = filename.encode('utf-8', 'replace')
@@ -397 +398 @@
     Every visitor method can be overridden to return an AST node that has been
     altered or replaced in some way.
     """
-    _visitors = {}
 
     def visit(self, node):
         if node is None:
             return None
-        v = self._visitors.get(node.__class__)
-        if not v:
-            v = getattr(self.__class__, 'visit%s' % node.__class__.__name__,
-                        self.__class__._visitDefault)
-            self._visitors[node.__class__] = v
-        return v(self, node)
+        if type(node) is tuple:
+            return tuple([self.visit(n) for n in node])
+        visitor = getattr(self, 'visit%s' % node.__class__.__name__,
+                          self._visitDefault)
+        return visitor(node)
 
     def _visitDefault(self, node):
         return node
@@ -475 +474 @@
         node.expr = self.visit(node.expr)
         return node
 
+    def visitAssAttr(self, node):
+        node.expr = self.visit(node.expr)
+        return node
+
+    def visitAugAssign(self, node):
+        node.node = self.visit(node.node)
+        node.expr = self.visit(node.expr)
+        return node
+
     def visitDecorators(self, node):
         node.nodes = [self.visit(x) for x in node.nodes]
         return node
 
+    def visitExec(self, node):
+        node.expr = self.visit(node.expr)
+        node.locals = self.visit(node.locals)
+        node.globals = self.visit(node.globals)
+        return node
+
     def visitFor(self, node):
         node.assign = self.visit(node.assign)
         node.list = self.visit(node.list)
@@ -503 +517 @@
         node.expr3 = self.visit(node.expr3)
         return node
 
+    def visitReturn(self, node):
+        node.value = self.visit(node.value)
+        return node
+
     def visitTryExcept(self, node):
         node.body = self.visit(node.body)
         node.handlers = self.visit(node.handlers)
@@ -536 +554 @@
         node.nodes = [self.visit(x) for x in node.nodes]
         return node
     visitAnd = visitOr = visitBitand = visitBitor = visitBitxor = _visitBoolOp
-    visitAssTuple = _visitBoolOp
+    visitAssTuple = visitAssList = _visitBoolOp
 
     def _visitBinOp(self, node):
         node.left = self.visit(node.left)
@@ -651 +669 @@
             self.locals[-1].add(node.name)
         return node
 
+    def visitAugAssign(self, node):
+        if isinstance(node.node, ast.Name):
+            name = node.node.name
+            node.node = ast.Subscript(ast.Name('data'), 'OP_APPLY',
+                                      [ast.Const(name)])
+            node.expr = self.visit(node.expr)
+            return ast.If([
+                (ast.Compare(ast.Const(name), [('in', ast.Name('data'))]),
+                 ast.Stmt([node]))],
+                ast.Stmt([ast.Raise(ast.CallFunc(ast.Name('UndefinedError'),
+                                                 [ast.Const(name)]),
+                                    None, None)]))
+        else:
+            return ASTTransformer.visitAugAssign(self, node)
+
     def visitClass(self, node):
         self.locals.append(set())
         node = ASTTransformer.visitClass(self, node)
@@ -675 +708 @@
         self.locals.pop()
         return node
 
-    def visitGetattr(self, node):
-        return ast.CallFunc(ast.Name('_lookup_attr'), [
-            ast.Name('data'), self.visit(node.expr),
-            ast.Const(node.attrname)
-        ])
-
     def visitLambda(self, node):
         self.locals.append(set(flatten(node.argnames)))
         node = ASTTransformer.visitLambda(self, node)
@@ -703 +730 @@
         func_args = [ast.Name('data'), ast.Const(node.name)]
         return ast.CallFunc(ast.Name('_lookup_name'), func_args)
 
+
+class ExpressionASTTransformer(TemplateASTTransformer):
+    """Concrete AST transformer that implements the AST transformations needed
+    for code embedded in templates.
+    """
+
+    def visitGetattr(self, node):
+        return ast.CallFunc(ast.Name('_lookup_attr'), [
+            ast.Name('data'), self.visit(node.expr),
+            ast.Const(node.attrname)
+        ])
+
     def visitSubscript(self, node):
         return ast.CallFunc(ast.Name('_lookup_item'), [
             ast.Name('data'), self.visit(node.expr),