Implement Decorators (#177)

compiler/stmt.py

@@ -309,9 +309,16 @@ class StatementVisitor(ast.NodeVisitor):
     self.writer.write_label(loop.end_label)
 
   def visit_FunctionDef(self, node):
-    self._write_py_context(node.lineno)
+    self._write_py_context(node.lineno + len(node.decorator_list))
     func = self.expr_visitor.visit_function_inline(node)
     self.block.bind_var(self.writer, node.name, func.expr)
+    while node.decorator_list:
+      decorator = node.decorator_list.pop()
+      wrapped = ast.Name(node.name, ast.Load)
+      decorated = ast.Call(decorator, [wrapped], [], None, None)
+      target = ast.Assign([wrapped], decorated)
+      target.lineno = node.lineno + len(node.decorator_list)
+      self.visit_Assign(target)
 
   def visit_Global(self, node):
     self._write_py_context(node.lineno)

compiler/stmt_test.py

@@ -212,6 +212,26 @@ class StatementVisitorTest(unittest.TestCase):
         util.ParseError, "'continue' not in loop",
         _ParseAndVisit, 'for i in (1,):\n  pass\nelse:\n  continue')
 
+  def testFunctionDecorator(self):
+    self.assertEqual((0, '<b>foo</b>\n'), _GrumpRun(textwrap.dedent("""\
+        def bold(fn):
+          return lambda: '<b>' + fn() + '</b>'
+        @bold
+        def foo():
+          return 'foo'
+        print foo()""")))
+
+  def testFunctionDecoratorWithArg(self):
+    self.assertEqual((0, '<b id=red>foo</b>\n'), _GrumpRun(textwrap.dedent("""\
+        def tag(name):
+          def bold(fn):
+            return lambda: '<b id=' + name + '>' + fn() + '</b>'
+          return bold
+        @tag('red')
+        def foo():
+          return 'foo'
+        print foo()""")))
+
   def testFunctionDef(self):
     self.assertEqual((0, 'bar baz\n'), _GrumpRun(textwrap.dedent("""\
         def foo(a, b):