From 47b2dfdd7307d9eb89d2b341ee7d4082e7c7ebaf Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Thu, 2 Jul 2026 17:09:50 +0200
Subject: [PATCH] feat(gen): add assertion collector to TypedAST

---
 midas/checker/python.py      | 1 +
 midas/generator/generator.py | 2 ++
 midas/utils.py               | 2 ++
 3 files changed, 5 insertions(+)

diff --git a/midas/checker/python.py b/midas/checker/python.py
index 374f736..a42115f 100644
--- a/midas/checker/python.py
+++ b/midas/checker/python.py
@@ -115,6 +115,7 @@ class PythonTyper(
             stmts=stmts,
             judgements=self.judgements,
             evaluated_casts=self.evaluated_casts,
+            assertions=self.assertions,
         )
 
     def judge(self, expr: p.Expr, type: Type):
diff --git a/midas/generator/generator.py b/midas/generator/generator.py
index c2f46ee..a8b3145 100644
--- a/midas/generator/generator.py
+++ b/midas/generator/generator.py
@@ -30,6 +30,7 @@ from midas.checker.types import (
     UnitType,
     UnknownType,
 )
+from midas.generator.collector import Assertion, AssertionCollector
 from midas.generator.constraints import ConstraintGenerator
 from midas.generator.stubs import StubsGenerator
 from midas.utils import TypedAST
@@ -55,6 +56,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
             stmts=[],
             judgements=[],
             evaluated_casts=[],
+            assertions=AssertionCollector(),
         )
         self._alias_count: int = 0
         self._predicate_count: int = 0
diff --git a/midas/utils.py b/midas/utils.py
index 0ac90cf..c225a89 100644
--- a/midas/utils.py
+++ b/midas/utils.py
@@ -3,6 +3,7 @@ from typing import Any, Callable, Optional
 
 import midas.ast.python as p
 from midas.checker.types import Type
+from midas.generator.collector import AssertionCollector
 
 AllowRepeat = Callable[[object], bool]
 
@@ -63,3 +64,4 @@ class TypedAST:
     stmts: list[p.Stmt]
     judgements: list[tuple[p.Expr, Type]]
     evaluated_casts: list[p.CastExpr]
+    assertions: AssertionCollector