From ad2fabf471539d518e130a0c06e98aa6f78caf15 Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Wed, 1 Jul 2026 22:32:13 +0200
Subject: [PATCH] feat(checker): add assertion collector

---
 midas/checker/frame_methods.py |  5 +++++
 midas/checker/python.py        |  2 ++
 midas/generator/collector.py   | 22 ++++++++++++++++++++++
 3 files changed, 29 insertions(+)
 create mode 100644 midas/generator/collector.py

diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py
index 0f85493..2afadfd 100644
--- a/midas/checker/frame_methods.py
+++ b/midas/checker/frame_methods.py
@@ -18,6 +18,7 @@ from midas.checker.types import (
     UnknownType,
     unfold_type,
 )
+from midas.generator.collector import AssertionCollector
 
 if TYPE_CHECKING:
     from midas.checker.python import PythonTyper, TypedExpr
@@ -77,6 +78,10 @@ class MethodRegistry(metaclass=_MethodRegistryMeta):
     def dispatcher(self) -> CallDispatcher[p.Expr]:
         return self.typer.dispatcher
 
+    @property
+    def assertions(self) -> AssertionCollector:
+        return self.typer.assertions
+
     def call(
         self,
         method: str,
diff --git a/midas/checker/python.py b/midas/checker/python.py
index d74cbce..70cb4c9 100644
--- a/midas/checker/python.py
+++ b/midas/checker/python.py
@@ -36,6 +36,7 @@ from midas.checker.types import (
     Variance,
     unfold_type,
 )
+from midas.generator.collector import AssertionCollector
 from midas.parser.python import PythonParser
 from midas.utils import TypedAST
 
@@ -87,6 +88,7 @@ class PythonTyper(
         self.dispatcher: CallDispatcher[p.Expr] = CallDispatcher[p.Expr](
             self.types, self.reporter
         )
+        self.assertions: AssertionCollector = AssertionCollector()
 
     def set_reporter(self, reporter: FileReporter):
         self.reporter = reporter
diff --git a/midas/generator/collector.py b/midas/generator/collector.py
new file mode 100644
index 0000000..38906fc
--- /dev/null
+++ b/midas/generator/collector.py
@@ -0,0 +1,22 @@
+import ast
+
+import midas.ast.python as p
+
+
+class AssertionCollector:
+    def __init__(self):
+        self.assertions: list[tuple[p.Expr, list[ast.expr]]] = []
+        self.definitions: dict[str, ast.stmt] = {}
+
+    def add(self, assertion):
+        self.assertions.append(assertion)
+
+    def define(self, name: str, stmt: ast.stmt):
+        if name not in self.definitions:
+            self.definitions[name] = stmt
+
+    def get_definitions(self) -> list[ast.stmt]:
+        return list(self.definitions.values())
+
+    def get_assertions(self) -> list[tuple[p.Expr, list[ast.expr]]]:
+        return self.assertions