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

fix(gen): improve cast assert message for columns

Browse Source
This commit is contained in:
Louis Heredero
2026-07-08 23:26:48 +02:00
parent 62612bd8db
commit 2df6bca948
1 changed files with 86 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

114
midas/generator/generator.py
View File

@@ -269,7 +269,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
alias: ast.expr = self._make_alias(expr.expr, 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, context=[]
)
for assert_ in asserts: for assert_ in asserts:
self._add_assert(assert_) self._add_assert(assert_)
@@ -526,7 +528,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
raise RuntimeError(f"Cannot get type judgement for {query}") raise RuntimeError(f"Cannot get type judgement for {query}")
def _make_cast_asserts( def _make_cast_asserts(
self, src_location: Location, expr: ast.expr, type: Type self,
src_location: Location,
expr: ast.expr,
type: Type,
*,
context: list[str],
) -> list[ast.stmt]: ) -> list[ast.stmt]:
"""Generate assertions for the given cast expression """Generate assertions for the given cast expression
@@ -535,6 +542,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
the source file the source file
expr (ast.expr): the expression being cast expr (ast.expr): the expression being cast
type (Type): the target type type (Type): the target type
context (list[str]): the current context
Returns: Returns:
list[ast.stmt]: the generated assertion statements list[ast.stmt]: the generated assertion statements
@@ -551,12 +559,16 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
args=[expr, ast.Name(id=name)], args=[expr, ast.Name(id=name)],
keywords=[], keywords=[],
), ),
self._make_cast_assert_message(src_location, expr, type), self._make_cast_assert_message(
src_location, expr, type, context=context
),
) )
] ]
case DerivedType(type=base): case DerivedType(type=base):
return self._make_cast_asserts(src_location, expr, base) return self._make_cast_asserts(
src_location, expr, base, context=context
)
case UnitType(): case UnitType():
return [ return [
@@ -568,19 +580,25 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
ast.Constant(value=None), ast.Constant(value=None),
], ],
), ),
self._make_cast_assert_message(src_location, expr, type), self._make_cast_assert_message(
src_location, expr, type, context=context
),
), ),
] ]
case AppliedType(body=body): case AppliedType(body=body):
return self._make_cast_asserts(src_location, expr, body) return self._make_cast_asserts(
src_location, expr, body, context=context
)
case ConstraintType(type=base, constraint=constraint): case ConstraintType(type=base, constraint=constraint):
asserts: list[ast.stmt] = self._make_cast_asserts( asserts: list[ast.stmt] = self._make_cast_asserts(
src_location, expr, base src_location, expr, base, context=context
) )
asserts.append( asserts.append(
self._make_constraint_assert(src_location, expr, constraint) self._make_constraint_assert(
src_location, expr, constraint, context=context
)
) )
return asserts return asserts
@@ -588,7 +606,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
# TODO: check with type from arguments / use call-site context # TODO: check with type from arguments / use call-site context
if bound is None: if bound is None:
return [] return []
return self._make_cast_asserts(src_location, expr, bound) return self._make_cast_asserts(
src_location, expr, bound, context=context
)
case TupleType(items=items): case TupleType(items=items):
asserts: list[ast.stmt] = [ asserts: list[ast.stmt] = [
@@ -598,13 +618,17 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
args=[expr, ast.Name(id="tuple")], args=[expr, ast.Name(id="tuple")],
keywords=[], keywords=[],
), ),
self._make_cast_assert_message(src_location, expr, type), self._make_cast_assert_message(
src_location, expr, type, context=context
),
), ),
] ]
assert isinstance(expr, ast.Tuple) assert isinstance(expr, ast.Tuple)
for item, item_type in zip(expr.elts, items): for item, item_type in zip(expr.elts, items):
asserts.extend( asserts.extend(
self._make_cast_asserts(src_location, item, item_type) self._make_cast_asserts(
src_location, item, item_type, context=context
)
) )
return asserts return asserts
@@ -618,7 +642,11 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
keywords=[], keywords=[],
), ),
self._make_cast_assert_message( self._make_cast_assert_message(
src_location, expr, type, ": Not a dataframe" src_location,
expr,
type,
context=context,
extra=": Not a dataframe",
), ),
), ),
] ]
@@ -634,7 +662,8 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
src_location, src_location,
expr, expr,
type, type,
f": Missing column {column.name}", context=context,
extra=f": Missing column '{column.name}'",
), ),
) )
) )
@@ -645,6 +674,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
value=expr, slice=ast.Constant(value=column.name) value=expr, slice=ast.Constant(value=column.name)
), ),
column.type, column.type,
context=context + [f"in column '{column.name}'"],
) )
) )
return asserts return asserts
@@ -659,12 +689,19 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
keywords=[], keywords=[],
), ),
self._make_cast_assert_message( self._make_cast_assert_message(
src_location, expr, type, ": Not a column" src_location,
expr,
type,
context=context,
extra=": Not a column",
), ),
), ),
] ]
inner_assert: Optional[ast.stmt] = self._make_column_inner_assert( inner_assert: Optional[ast.stmt] = self._make_column_inner_assert(
src_location, expr, type src_location,
expr,
type,
context,
) )
if inner_assert is not None: if inner_assert is not None:
asserts.append(inner_assert) asserts.append(inner_assert)
@@ -689,7 +726,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
location: Location, location: Location,
expr: ast.expr, expr: ast.expr,
type: Type, type: Type,
extra: Optional[str] = None, *,
context: list[str],
extra: str = "",
) -> ast.expr: ) -> ast.expr:
"""Build an AST node for a cast assertion message """Build an AST node for a cast assertion message
@@ -703,12 +742,14 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
source file source file
expr (ast.expr): the expression being cast expr (ast.expr): the expression being cast
type (Type): the target type type (Type): the target type
extra (Optional[str], optional): extra text to append at the end of context (list[str]): the current context
the message. Defaults to None. extra (str, optional): extra text to append at the end of
the message. Defaults to "".
Returns: Returns:
ast.expr: the generated message (as an f-string) ast.expr: the generated message (as an f-string)
""" """
context_str: str = "".join(map(lambda c: f", {c}", context))
loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}" loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
# f"file.py:L1:1: CastError: Cannot cast {type(expr).__name__} to Type" # f"file.py:L1:1: CastError: Cannot cast {type(expr).__name__} to Type"
return ast.JoinedStr( return ast.JoinedStr(
@@ -725,12 +766,17 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
), ),
conversion=-1, conversion=-1,
), ),
ast.Constant(f" to {type}{extra or ''}"), ast.Constant(f" to {type}{context_str}{extra}"),
] ]
) )
def _make_constraint_assert( def _make_constraint_assert(
self, src_location: Location, expr: ast.expr, constraint: m.Expr self,
src_location: Location,
expr: ast.expr,
constraint: m.Expr,
*,
context: list[str],
) -> ast.stmt: ) -> ast.stmt:
"""Build an assertion for the given constraint on the given expression """Build an assertion for the given constraint on the given expression
@@ -739,6 +785,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
source file source file
expr (ast.expr): the expression subject to `constraint` expr (ast.expr): the expression subject to `constraint`
constraint (m.Expr): the constraint applied on `expr` constraint (m.Expr): the constraint applied on `expr`
context (list[str]): the current context
Returns: Returns:
ast.stmt: the assert statement checking the constraint ast.stmt: the assert statement checking the constraint
@@ -750,11 +797,13 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
args=[expr], args=[expr],
keywords=[], keywords=[],
), ),
self._make_constraint_assert_message(src_location, constraint), self._make_constraint_assert_message(
src_location, constraint, context=context
),
) )
def _make_constraint_assert_message( def _make_constraint_assert_message(
self, location: Location, constraint: m.Expr self, location: Location, constraint: m.Expr, *, context: list[str]
) -> ast.expr: ) -> ast.expr:
"""Build an assert message for the given constraint """Build an assert message for the given constraint
@@ -762,16 +811,18 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
location (Location): the location of the cast expression in the location (Location): the location of the cast expression in the
source file source file
constraint (m.Expr): the constraint constraint (m.Expr): the constraint
context (list[str]): the current context
Returns: Returns:
ast.expr: the assert message ast.expr: the assert message
""" """
printer = MidasPrinter() printer = MidasPrinter()
constraint_str: str = printer.print(constraint) constraint_str: str = printer.print(constraint)
context_str: str = "".join(map(lambda c: f", {c}", context))
loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}" loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
# f"file.py:L1:1: ConstraintError: Value does not fit constraint 'v > 0'" # f"file.py:L1:1: ConstraintError: Value does not fit constraint 'v > 0'"
return ast.Constant( return ast.Constant(
f"{loc_str}: ConstraintError: Value does not fit constraint '{constraint_str}'" f"{loc_str}: ConstraintError: Value does not fit constraint '{constraint_str}'{context_str}"
) )
def _get_constraint(self, expr: m.Expr) -> ast.expr: def _get_constraint(self, expr: m.Expr) -> ast.expr:
@@ -880,7 +931,11 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
) )
def _make_column_inner_assert( def _make_column_inner_assert(
self, src_location: Location, column: ast.expr, type: ColumnType self,
src_location: Location,
column: ast.expr,
type: ColumnType,
context: list[str],
) -> Optional[ast.stmt]: ) -> Optional[ast.stmt]:
"""Build a for-loop checking the type of values inside a column """Build a for-loop checking the type of values inside a column
@@ -889,18 +944,21 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
source file source file
column (ast.expr): the column being cast column (ast.expr): the column being cast
type (ColumnType): the type of the column type (ColumnType): the type of the column
context (list[str]): the current context
Returns: Returns:
Optional[ast.stmt]: a for-loop checking the values, or `None` if no Optional[ast.stmt]: a for-loop checking the values, or `None` if no
assertions are necessary assertions are necessary
""" """
# TODO: improve message, maybe chain contexts
col: ast.expr = ast.Name(id="col") value: ast.expr = ast.Name(id="value")
body: list[ast.stmt] = self._make_cast_asserts(src_location, col, type.type) body: list[ast.stmt] = self._make_cast_asserts(
src_location, value, type.type, context=context
)
if len(body) == 0: if len(body) == 0:
return None return None
return ast.For( return ast.For(
target=col, target=value,
iter=column, iter=column,
body=body, body=body,
orelse=[], orelse=[],