HEL/midas
1
0
Fork 0
Code Issues Pull Requests Actions Releases Activity

feat(gen): materialize assertions from collector

Browse Source
This commit is contained in:
Louis Heredero
2026-07-02 17:10:18 +02:00
parent 47b2dfdd73
commit 8df01afd8c
2 changed files with 85 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
midas/generator/collector.py
View File

@@ -39,6 +39,12 @@ class AssertionCollector:
) )
) )
def remove(self, assertion: Assertion):
try:
self.assertions.remove(assertion)
except ValueError:
pass
def define(self, name: str, stmt: ast.stmt): def define(self, name: str, stmt: ast.stmt):
if name not in self.definitions: if name not in self.definitions:
self.definitions[name] = stmt self.definitions[name] = stmt

116
midas/generator/generator.py
View File

@@ -61,6 +61,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
self._alias_count: int = 0 self._alias_count: int = 0
self._predicate_count: int = 0 self._predicate_count: int = 0
self._scopes: list[Scope] = [] self._scopes: list[Scope] = []
self._aliases: list[tuple[p.Expr, ast.expr]] = []
self._constraint_generator: ConstraintGenerator = ConstraintGenerator(types) self._constraint_generator: ConstraintGenerator = ConstraintGenerator(types)
self._constraints: list[tuple[m.Expr, ast.expr]] = [] self._constraints: list[tuple[m.Expr, ast.expr]] = []
@@ -73,7 +74,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def generate_ast(self, typed_ast: TypedAST) -> ast.AST: def generate_ast(self, typed_ast: TypedAST) -> ast.AST:
self._typed_ast = typed_ast self._typed_ast = typed_ast
body: list[ast.stmt] = self._visit_body(typed_ast.stmts) body: list[ast.stmt] = self._visit_body(typed_ast.stmts, can_be_empty=True)
predicates: list[ast.stmt] = self._constraint_generator.get_definitions() predicates: list[ast.stmt] = self._constraint_generator.get_definitions()
body = predicates + body body = predicates + body
@@ -131,39 +132,48 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
output: str = ast.unparse(module) output: str = ast.unparse(module)
out_path.write_text(output) out_path.write_text(output)
def convert(self, expr: p.Expr) -> ast.expr:
for expr2, alias in self._aliases:
if expr2 == expr:
return alias
assertions = self._typed_ast.assertions.get_assertions_for(expr)
if len(assertions) != 0:
return self._apply_assertions(expr, assertions)
return expr.accept(self)
def visit_binary_expr(self, expr: p.BinaryExpr) -> ast.expr: def visit_binary_expr(self, expr: p.BinaryExpr) -> ast.expr:
return ast.BinOp( return ast.BinOp(
left=expr.left.accept(self), left=self.convert(expr.left),
op=expr.operator, op=expr.operator,
right=expr.right.accept(self), right=self.convert(expr.right),
) )
def visit_compare_expr(self, expr: p.CompareExpr) -> ast.expr: def visit_compare_expr(self, expr: p.CompareExpr) -> ast.expr:
return ast.Compare( return ast.Compare(
left=expr.left.accept(self), left=self.convert(expr.left),
ops=[expr.operator], ops=[expr.operator],
comparators=[expr.right.accept(self)], comparators=[self.convert(expr.right)],
) )
def visit_unary_expr(self, expr: p.UnaryExpr) -> ast.expr: def visit_unary_expr(self, expr: p.UnaryExpr) -> ast.expr:
return ast.UnaryOp( return ast.UnaryOp(
op=expr.operator, op=expr.operator,
operand=expr.right.accept(self), operand=self.convert(expr.right),
) )
def visit_call_expr(self, expr: p.CallExpr) -> ast.expr: def visit_call_expr(self, expr: p.CallExpr) -> ast.expr:
return ast.Call( return ast.Call(
func=expr.callee.accept(self), func=self.convert(expr.callee),
args=[arg.accept(self) for arg in expr.arguments], args=[self.convert(arg) for arg in expr.arguments],
keywords=[ keywords=[
ast.keyword(arg=name, value=arg.accept(self)) ast.keyword(arg=name, value=self.convert(arg))
for name, arg in expr.keywords.items() for name, arg in expr.keywords.items()
], ],
) )
def visit_get_expr(self, expr: p.GetExpr) -> ast.expr: def visit_get_expr(self, expr: p.GetExpr) -> ast.expr:
return ast.Attribute( return ast.Attribute(
value=expr.object.accept(self), value=self.convert(expr.object),
attr=expr.name, attr=expr.name,
) )
@@ -176,16 +186,16 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def visit_logical_expr(self, expr: p.LogicalExpr) -> ast.expr: def visit_logical_expr(self, expr: p.LogicalExpr) -> ast.expr:
return ast.BoolOp( return ast.BoolOp(
op=expr.operator, op=expr.operator,
values=[expr.left.accept(self), expr.right.accept(self)], values=[self.convert(expr.left), self.convert(expr.right)],
) )
def visit_cast_expr(self, expr: p.CastExpr) -> ast.expr: def visit_cast_expr(self, expr: p.CastExpr) -> ast.expr:
expr2: ast.expr = expr.expr.accept(self) expr2: ast.expr = self.convert(expr.expr)
if expr in self._typed_ast.evaluated_casts or expr.unsafe: if expr in self._typed_ast.evaluated_casts or expr.unsafe:
return expr2 return expr2
alias: ast.expr = self._make_alias(expr2) alias: ast.expr = self._make_alias(expr.expr, expr2)
type: Type = self._get_expr_type(expr) type: Type = self._get_expr_type(expr)
asserts: list[ast.stmt] = self._make_cast_asserts(expr.location, alias, type) asserts: list[ast.stmt] = self._make_cast_asserts(expr.location, alias, type)
@@ -196,38 +206,38 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def visit_ternary_expr(self, expr: p.TernaryExpr) -> ast.expr: def visit_ternary_expr(self, expr: p.TernaryExpr) -> ast.expr:
return ast.IfExp( return ast.IfExp(
test=expr.test.accept(self), test=self.convert(expr.test),
body=expr.if_true.accept(self), body=self.convert(expr.if_true),
orelse=expr.if_false.accept(self), orelse=self.convert(expr.if_false),
) )
def visit_list_expr(self, expr: p.ListExpr) -> ast.expr: def visit_list_expr(self, expr: p.ListExpr) -> ast.expr:
return ast.List( return ast.List(
elts=[item.accept(self) for item in expr.items], elts=[self.convert(item) for item in expr.items],
) )
def visit_dict_expr(self, expr: p.DictExpr) -> ast.expr: def visit_dict_expr(self, expr: p.DictExpr) -> ast.expr:
return ast.Dict( return ast.Dict(
keys=[key.accept(self) if key is not None else None for key in expr.keys], keys=[self.convert(key) if key is not None else None for key in expr.keys],
values=[value.accept(self) for value in expr.values], values=[self.convert(value) for value in expr.values],
) )
def visit_subscript_expr(self, expr: p.SubscriptExpr) -> ast.expr: def visit_subscript_expr(self, expr: p.SubscriptExpr) -> ast.expr:
return ast.Subscript( return ast.Subscript(
value=expr.object.accept(self), value=self.convert(expr.object),
slice=expr.index.accept(self), slice=self.convert(expr.index),
) )
def visit_slice_expr(self, expr: p.SliceExpr) -> ast.expr: def visit_slice_expr(self, expr: p.SliceExpr) -> ast.expr:
return ast.Slice( return ast.Slice(
lower=expr.lower.accept(self) if expr.lower is not None else None, lower=self.convert(expr.lower) if expr.lower is not None else None,
upper=expr.upper.accept(self) if expr.upper is not None else None, upper=self.convert(expr.upper) if expr.upper is not None else None,
step=expr.step.accept(self) if expr.step is not None else None, step=self.convert(expr.step) if expr.step is not None else None,
) )
def visit_tuple_expr(self, expr: p.TupleExpr) -> ast.expr: def visit_tuple_expr(self, expr: p.TupleExpr) -> ast.expr:
return ast.Tuple( return ast.Tuple(
elts=[item.accept(self) for item in expr.items], elts=[self.convert(item) for item in expr.items],
) )
def visit_raw_expr(self, expr: p.RawExpr) -> ast.expr: def visit_raw_expr(self, expr: p.RawExpr) -> ast.expr:
@@ -235,7 +245,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> ast.stmt: def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> ast.stmt:
return ast.Expr( return ast.Expr(
value=stmt.expr.accept(self), value=self.convert(stmt.expr),
) )
def visit_function(self, stmt: p.Function) -> ast.stmt: def visit_function(self, stmt: p.Function) -> ast.stmt:
@@ -248,12 +258,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
kwonlyargs=[ast.arg(arg=arg.name) for arg in stmt.kwonlyargs], kwonlyargs=[ast.arg(arg=arg.name) for arg in stmt.kwonlyargs],
kwarg=None, kwarg=None,
defaults=[ defaults=[
arg.default.accept(self) self.convert(arg.default)
for arg in stmt.posonlyargs + stmt.args for arg in stmt.posonlyargs + stmt.args
if arg.default is not None if arg.default is not None
], ],
kw_defaults=[ kw_defaults=[
arg.default.accept(self) if arg.default is not None else None self.convert(arg.default) if arg.default is not None else None
for arg in stmt.kwonlyargs for arg in stmt.kwonlyargs
], ],
), ),
@@ -267,20 +277,20 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def visit_assign_stmt(self, stmt: p.AssignStmt) -> ast.stmt: def visit_assign_stmt(self, stmt: p.AssignStmt) -> ast.stmt:
return ast.Assign( return ast.Assign(
targets=[target.accept(self) for target in stmt.targets], targets=[self.convert(target) for target in stmt.targets],
value=stmt.value.accept(self), value=self.convert(stmt.value),
) )
def visit_return_stmt(self, stmt: p.ReturnStmt) -> ast.stmt: def visit_return_stmt(self, stmt: p.ReturnStmt) -> ast.stmt:
return ast.Return( return ast.Return(
value=stmt.value.accept(self) if stmt.value is not None else None, value=self.convert(stmt.value) if stmt.value is not None else None,
) )
def visit_if_stmt(self, stmt: p.IfStmt) -> ast.stmt: def visit_if_stmt(self, stmt: p.IfStmt) -> ast.stmt:
return ast.If( return ast.If(
test=stmt.test.accept(self), test=self.convert(stmt.test),
body=self._visit_body(stmt.body), body=self._visit_body(stmt.body),
orelse=self._visit_body(stmt.orelse), orelse=self._visit_body(stmt.orelse, can_be_empty=True),
) )
def visit_pass(self, stmt: p.Pass) -> ast.stmt: def visit_pass(self, stmt: p.Pass) -> ast.stmt:
@@ -288,8 +298,8 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def visit_for_stmt(self, stmt: p.ForStmt) -> ast.stmt: def visit_for_stmt(self, stmt: p.ForStmt) -> ast.stmt:
return ast.For( return ast.For(
target=stmt.target.accept(self), target=self.convert(stmt.target),
iter=stmt.iterator.accept(self), iter=self.convert(stmt.iterator),
body=self._visit_body(stmt.body), body=self._visit_body(stmt.body),
orelse=[], orelse=[],
) )
@@ -297,7 +307,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def visit_raw_stmt(self, stmt: p.RawStmt) -> ast.stmt: def visit_raw_stmt(self, stmt: p.RawStmt) -> ast.stmt:
return stmt.stmt return stmt.stmt
def _visit_body(self, stmts: list[p.Stmt]) -> list[ast.stmt]: def _visit_body(
self, stmts: list[p.Stmt], can_be_empty: bool = False
) -> list[ast.stmt]:
generated: list[ast.stmt] = [] generated: list[ast.stmt] = []
for stmt in stmts: for stmt in stmts:
scope = Scope() scope = Scope()
@@ -315,9 +327,11 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
# Remove redundant pass statements # Remove redundant pass statements
if len(generated) > 1: if len(generated) > 1:
generated = [stmt for stmt in generated if not isinstance(stmt, ast.Pass)] generated = [stmt for stmt in generated if not isinstance(stmt, ast.Pass)]
if len(generated) == 0 and not can_be_empty:
generated = [ast.Pass()]
return generated return generated
def _make_alias(self, expr: ast.expr) -> ast.expr: def _make_alias(self, node: p.Expr, expr: ast.expr) -> ast.expr:
name: str = f"__midas_a{self._alias_count}__" name: str = f"__midas_a{self._alias_count}__"
alias = ast.Name(id=name) alias = ast.Name(id=name)
self._alias_count += 1 self._alias_count += 1
@@ -328,6 +342,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
value=expr, value=expr,
) )
) )
self._aliases.append((node, alias))
return alias return alias
def _build_assert(self, expr: ast.expr, message: str | ast.expr) -> ast.stmt: def _build_assert(self, expr: ast.expr, message: str | ast.expr) -> ast.stmt:
@@ -642,3 +657,30 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
body=body, body=body,
orelse=[], orelse=[],
) )
def _convert_assertion(self, assertion: Assertion) -> ast.stmt:
inputs: list[ast.expr] = []
for input in assertion.inputs:
converted: ast.expr = self.convert(input)
alias: ast.expr = self._make_alias(input, converted)
inputs.append(alias)
test: ast.expr = assertion.builder(*inputs)
location: Location = assertion.bound_expr.location
loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
return self._build_assert(
test, f"{loc_str}: AssertionError: {assertion.message}"
)
def _apply_assertions(self, expr: p.Expr, assertions: list[Assertion]) -> ast.expr:
for assertion in assertions:
assert_stmt: ast.stmt
assert_stmt = self._convert_assertion(assertion)
self._add_assert(assert_stmt)
# Mutating list in frozen dataclass
# Not ideal but easiest way to avoid duplicate assertions
self._typed_ast.assertions.remove(assertion)
return expr.accept(self)