Implement decorators (#26)

src/interpreted/interpreter.py

@@ -472,6 +472,18 @@ def visit_ImportFrom(self, node: ImportFrom) -> None:
     def visit_FunctionDef(self, node: FunctionDef) -> None:
         parent_scope = self.scope
         function = UserFunction(node, parent_scope, self.globals)
+
+        decorators = reversed(node.decorators)
+
+        for decorator_node in decorators:
+            decorator = self.visit(decorator_node.value)
+
+            if not isinstance(decorator, Function):
+                object_type = decorator.__class__.__name__
+                raise InterpreterError(f"{object_type!r} object is not callable")
+
+            function = decorator.call(self, [function])
+
         self.scope.set(node.name, function)
 
     def visit_Assign(self, node: Assign) -> None:

src/interpreted/nodes.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Literal, Union
 
 BINOP = Literal["+", "-", "*", "/", "**", "@", "%", "^", "&"]
@@ -117,6 +117,7 @@ class FunctionDef(Statement):
     name: str
     params: list[str]
     body: list[Statement]
+    decorators: list[Decorator] = field(default_factory=list)
 
 
 @dataclass
@@ -194,3 +195,7 @@ class ImportFrom(Statement):
 @dataclass
 class Module:
     body: list[Statement]
+
+@dataclass
+class Decorator(Node):
+    value: Expression

src/interpreted/parser.py

@@ -15,6 +15,7 @@
     Compare,
     Constant,
     Continue,
+    Decorator,
     Dict,
     Expression,
     ExprStmt,
@@ -177,6 +178,11 @@ def expect_op(self, op: str) -> None:
             token = self.peek()
             raise ParseError(f"Expected '{op}', found '{token.string}'", self.index)
 
+    def expect_name(self, name: str) -> None:
+        if not self.match_name(name):
+            token = self.peek()
+            raise ParseError(f"Expected '{name}', found '{token.string}'", self.index)
+
     def parse(self) -> Module:
         statements: list[Statement] = []
         while not self.parsed:
@@ -190,10 +196,25 @@ def parse_statement(self) -> Statement:
         while self.match_type(TokenType.NEWLINE):
             pass
 
+        if self.peek().string == "@":
+            decorators = self.parse_decorators()
+            self.expect_name("def")
+            function_def = self.parse_function_def()
+            function_def.decorators = decorators
+            return function_def
         if self.match_name("def", "if", "for", "while"):
             return self.parse_multiline_statement()
         return self.parse_single_line_statement()
 
+    def parse_decorators(self) -> list[Decorator]:
+        decorators = []
+        while self.match_op("@"):
+            expression = self.parse_expression()
+            self.expect(TokenType.NEWLINE)
+            decorators.append(Decorator(value=expression))
+
+        return decorators
+
     def parse_multiline_statement(self) -> FunctionDef | For | If | While:
         keyword = self.current().string
         if keyword == "def":

tests/interpreted_test.py

@@ -95,6 +95,61 @@ def foo():
             """,
             "10\n",
         ),
+        (
+            """\
+            def foo(func):    
+                print('inside decorator')
+                return func
+
+            @foo
+            def xyz():
+                print('inside xyz')
+
+            xyz()
+            """,
+            "inside decorator\ninside xyz\n",
+        ),
+        (
+            """\
+            def decorator_foo(func):
+                print('Inside decorator foo')
+                return func
+
+            def ab5(func):  
+                print('Inside decorator bar')
+                return func
+
+            @decorator_foo
+            @ab5
+            def xyz():
+                print('Inside xyz')
+
+            xyz()
+            """,
+            "Inside decorator bar\nInside decorator foo\nInside xyz\n",
+        ),
+        (
+            """\
+            def decorator_foo(func):
+                print('Inside decorator foo')
+                return func
+
+            def ab5(func):  
+                print('Inside decorator bar')
+                def wrapper():
+                    print('Inside wrapper')
+                    return func()
+                return wrapper
+
+            @decorator_foo
+            @ab5
+            def xyz():
+                print('Inside xyz')
+
+            xyz()
+            """,
+            "Inside decorator bar\nInside decorator foo\nInside wrapper\nInside xyz\n",
+        ),
     ),
 )
 def test_interpret(source, output) -> None: