feat(gen): add Assertion class

midas/generator/collector.py

@@ -1,15 +1,43 @@
 import ast
+from dataclasses import dataclass
+from typing import Callable
 
 import midas.ast.python as p
 
+AssertionBuilder = Callable[..., ast.expr]
+
+
+@dataclass
+class Assertion:
+    bound_expr: p.Expr
+    inputs: list[p.Expr]
+    builder: AssertionBuilder
+    message: str
+
+    def is_bound_to(self, expr: p.Expr) -> bool:
+        return expr == self.bound_expr
+
 
 class AssertionCollector:
     def __init__(self):
-        self.assertions: list[tuple[p.Expr, list[ast.expr]]] = []
+        self.assertions: list[Assertion] = []
         self.definitions: dict[str, ast.stmt] = {}
 
-    def add(self, assertion):
-        self.assertions.append(assertion)
+    def add(
+        self,
+        bound_expr: p.Expr,
+        inputs: list[p.Expr],
+        builder: AssertionBuilder,
+        message: str,
+    ):
+        self.assertions.append(
+            Assertion(
+                bound_expr=bound_expr,
+                inputs=inputs,
+                builder=builder,
+                message=message,
+            )
+        )
 
     def define(self, name: str, stmt: ast.stmt):
         if name not in self.definitions:
@@ -18,5 +46,8 @@ class AssertionCollector:
     def get_definitions(self) -> list[ast.stmt]:
         return list(self.definitions.values())
 
-    def get_assertions(self) -> list[tuple[p.Expr, list[ast.expr]]]:
+    def get_assertions(self) -> list[Assertion]:
         return self.assertions
+
+    def get_assertions_for(self, expr: p.Expr) -> list[Assertion]:
+        return list(filter(lambda a: a.is_bound_to(expr), self.assertions))