From 8304760fe08fe28e40d192aeb3457a663895d458 Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Thu, 28 May 2026 15:26:33 +0200
Subject: [PATCH] fix(parser): add function body and all_args property

---
 gen/python.py          | 7 ++++++-
 midas/ast/python.py    | 7 ++++++-
 midas/parser/python.py | 9 +++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/gen/python.py b/gen/python.py
index 0aadd57..5e865ad 100644
--- a/gen/python.py
+++ b/gen/python.py
@@ -46,13 +46,18 @@ class Function:
     args: list[Argument]
     kwonlyargs: list[Argument]
     returns: Optional[MidasType]
+    body: list[Stmt]
 
     @dataclass(frozen=True, kw_only=True)
     class Argument:
         location: Optional[Location] = None
-        name: Optional[str]
+        name: str
         type: Optional[MidasType]
 
+    @property
+    def all_args(self) -> list[Argument]:
+        return self.posonlyargs + self.args + self.kwonlyargs
+
 
 class TypeAssign:
     name: str
diff --git a/midas/ast/python.py b/midas/ast/python.py
index d4fc032..96c218e 100644
--- a/midas/ast/python.py
+++ b/midas/ast/python.py
@@ -116,13 +116,18 @@ class Function(Stmt):
     args: list[Argument]
     kwonlyargs: list[Argument]
     returns: Optional[MidasType]
+    body: list[Stmt]
 
     @dataclass(frozen=True, kw_only=True)
     class Argument:
         location: Optional[Location] = None
-        name: Optional[str]
+        name: str
         type: Optional[MidasType]
 
+    @property
+    def all_args(self) -> list[Argument]:
+        return self.posonlyargs + self.args + self.kwonlyargs
+
     def accept(self, visitor: Stmt.Visitor[T]) -> T:
         return visitor.visit_function(self)
 
diff --git a/midas/parser/python.py b/midas/parser/python.py
index 4b6a3f1..470f666 100644
--- a/midas/parser/python.py
+++ b/midas/parser/python.py
@@ -128,11 +128,19 @@ class PythonParser:
                     kwonlyargs=kwonlyargs,
                 ),
                 returns=returns,
+                body=raw_body,
             ):
 
                 def parse_args(args_list: list[ast.arg]) -> list[Function.Argument]:
                     return [self._parse_function_argument(arg) for arg in args_list]
 
+                body: list[Stmt] = []
+                for stmt in raw_body:
+                    stmts = self.parse_stmt(stmt)
+                    if isinstance(stmts, Stmt):
+                        body.append(stmts)
+                    elif stmts is not None:
+                        body.extend(stmts)
                 return Function(
                     location=loc,
                     name=name,
@@ -140,6 +148,7 @@ class PythonParser:
                     args=parse_args(args),
                     kwonlyargs=parse_args(kwonlyargs),
                     returns=self._parse_type(returns) if returns is not None else None,
+                    body=body,
                 )
             case _:
                 print(f"Unsupported function definition: {ast.unparse(node)}")