feat(parser): add return statements

gen/python.py

@@ -69,6 +69,10 @@ class AssignStmt:
     value: Expr
 
 
+class ReturnStmt:
+    value: Optional[Expr]
+
+
 ###<
 
 

midas/ast/printer.py

@@ -449,6 +449,11 @@ class PythonAstPrinter(
             with self._child_level(single=True):
                 stmt.value.accept(self)
 
+    def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
+        self._write_line("ReturnStmt")
+        with self._child_level():
+            self._write_optional_child("value", stmt.value, last=True)
+
     def visit_binary_expr(self, expr: p.BinaryExpr) -> None:
         self._write_line("BinaryExpr")
         with self._child_level():

midas/ast/python.py

@@ -100,6 +100,9 @@ class Stmt(ABC):
         @abstractmethod
         def visit_assign_stmt(self, stmt: AssignStmt) -> T: ...
 
+        @abstractmethod
+        def visit_return_stmt(self, stmt: ReturnStmt) -> T: ...
+
 
 @dataclass(frozen=True)
 class ExpressionStmt(Stmt):
@@ -150,6 +153,14 @@ class AssignStmt(Stmt):
         return visitor.visit_assign_stmt(self)
 
 
+@dataclass(frozen=True)
+class ReturnStmt(Stmt):
+    value: Optional[Expr]
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_return_stmt(self)
+
+
 ###############
 # Expressions #
 ###############

midas/cli/highlighter.py

@@ -153,6 +153,8 @@ class PythonHighlighter(
 
     def visit_assign_stmt(self, stmt: p.AssignStmt) -> None: ...
 
+    def visit_return_stmt(self, stmt: p.ReturnStmt) -> None: ...
+
     def visit_binary_expr(self, expr: p.BinaryExpr) -> None: ...
 
     def visit_compare_expr(self, expr: p.CompareExpr) -> None: ...

midas/parser/python.py

@@ -19,6 +19,7 @@ from midas.ast.python import (
     LiteralExpr,
     LogicalExpr,
     MidasType,
+    ReturnStmt,
     Stmt,
     TypeAssign,
     UnaryExpr,
@@ -70,6 +71,12 @@ class PythonParser:
                     expr=self.parse_expr(expr),
                 )
 
+            case ast.Return(value=value):
+                return ReturnStmt(
+                    location=location,
+                    value=self.parse_expr(value) if value is not None else None,
+                )
+
             case _:
                 print(f"Unsupported statement: {ast.unparse(node)}")
                 return None

midas/resolver/resolver.py

@@ -72,6 +72,10 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
                 case _:
                     raise Exception(f"Unsupported assignment to {target}")
 
+    def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
+        if stmt.value is not None:
+            self.resolve(stmt.value)
+
     def visit_binary_expr(self, expr: p.BinaryExpr) -> None:
         self.resolve(expr.left)
         self.resolve(expr.right)