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