From 8906ac3db81a24565148f3e42e06a8ab09ab6397 Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Fri, 29 May 2026 11:25:11 +0200
Subject: [PATCH] feat(parser): add return statements

---
 gen/python.py              |  4 ++++
 midas/ast/printer.py       |  5 +++++
 midas/ast/python.py        | 11 +++++++++++
 midas/cli/highlighter.py   |  2 ++
 midas/parser/python.py     |  7 +++++++
 midas/resolver/resolver.py |  4 ++++
 6 files changed, 33 insertions(+)

diff --git a/gen/python.py b/gen/python.py
index 5e865ad..76b68a4 100644
--- a/gen/python.py
+++ b/gen/python.py
@@ -69,6 +69,10 @@ class AssignStmt:
     value: Expr
 
 
+class ReturnStmt:
+    value: Optional[Expr]
+
+
 ###<
 
 
diff --git a/midas/ast/printer.py b/midas/ast/printer.py
index e3ecde9..c923551 100644
--- a/midas/ast/printer.py
+++ b/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():
diff --git a/midas/ast/python.py b/midas/ast/python.py
index 48f4643..7c22b2f 100644
--- a/midas/ast/python.py
+++ b/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 #
 ###############
diff --git a/midas/cli/highlighter.py b/midas/cli/highlighter.py
index f4801bb..48d71d7 100644
--- a/midas/cli/highlighter.py
+++ b/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: ...
diff --git a/midas/parser/python.py b/midas/parser/python.py
index 892725a..36d6efb 100644
--- a/midas/parser/python.py
+++ b/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
diff --git a/midas/resolver/resolver.py b/midas/resolver/resolver.py
index 7b2b4d0..d63d893 100644
--- a/midas/resolver/resolver.py
+++ b/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)