From 74f51f361a4d21886b2a34a400ba5fea0a2cd6e1 Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Mon, 15 Jun 2026 14:16:10 +0200
Subject: [PATCH] feat(checker): make checker return TypedAST

---
 midas/checker/checker.py |  7 ++++---
 midas/checker/python.py  |  5 ++++-
 midas/utils.py           | 10 ++++++++++
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/midas/checker/checker.py b/midas/checker/checker.py
index c26f0aa..c64a282 100644
--- a/midas/checker/checker.py
+++ b/midas/checker/checker.py
@@ -6,6 +6,7 @@ from midas.checker.midas import MidasTyper
 from midas.checker.python import PythonTyper
 from midas.checker.registry import TypesRegistry
 from midas.checker.reporter import Reporter
+from midas.utils import TypedAST
 
 
 class TypeChecker:
@@ -23,12 +24,12 @@ class TypeChecker:
     def import_midas_source(self, source: str, path: Optional[str] = None):
         self.midas_typer.process(source, path)
 
-    def type_check(self, path: Path):
+    def type_check(self, path: Path) -> TypedAST:
         source: str = path.read_text()
         return self.type_check_source(source, path=str(path))
 
-    def type_check_source(self, source: str, path: Optional[str] = None):
-        self.python_typer.process(source, path)
+    def type_check_source(self, source: str, path: Optional[str] = None) -> TypedAST:
+        return self.python_typer.process(source, path)
 
     @property
     def diagnostics(self) -> list[Diagnostic]:
diff --git a/midas/checker/python.py b/midas/checker/python.py
index 93ac111..d667164 100644
--- a/midas/checker/python.py
+++ b/midas/checker/python.py
@@ -19,6 +19,7 @@ from midas.checker.types import (
     unfold_type,
 )
 from midas.parser.python import PythonParser
+from midas.utils import TypedAST
 
 TypedExpr = tuple[p.Expr, Type]
 
@@ -60,7 +61,7 @@ class PythonTyper(
         self.locals: dict[p.Expr, int] = {}
         self.judgements: list[tuple[p.Expr, Type]] = []
 
-    def process(self, source: str, path: Optional[str]):
+    def process(self, source: str, path: Optional[str]) -> TypedAST:
         self.reporter = self.reporter.for_file(path)
 
         tree: ast.Module = ast.parse(source, filename=path or "<unknown>")
@@ -75,6 +76,8 @@ class PythonTyper(
 
         self.check(stmts)
 
+        return TypedAST(stmts=stmts, judgements=self.judgements)
+
     def type_of(self, expr: p.Expr) -> Type:
         """Evaluate the type of an expression
 
diff --git a/midas/utils.py b/midas/utils.py
index bf10b95..f10bdb3 100644
--- a/midas/utils.py
+++ b/midas/utils.py
@@ -1,5 +1,9 @@
+from dataclasses import dataclass
 from typing import Any, Callable, Optional
 
+import midas.ast.python as p
+from midas.checker.types import Type
+
 AllowRepeat = Callable[[object], bool]
 
 
@@ -52,3 +56,9 @@ class UniversalJSONDumper:
                 }
             case _:
                 raise ValueError(f"Unsupported value: {obj}")
+
+
+@dataclass(frozen=True, kw_only=True)
+class TypedAST:
+    stmts: list[p.Stmt]
+    judgements: list[tuple[p.Expr, Type]]