feat(checker): add assertion collector

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,

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

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