fix(checker): correctly check length of frame/column

This commit is contained in:
2026-07-03 12:28:39 +02:00
parent ec80b1e92e
commit 4d3e3f44a1
2 changed files with 34 additions and 18 deletions

View File

@@ -355,6 +355,21 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
def _assert_same_length(self, call_expr: p.Expr, column1: p.Expr, column2: p.Expr):
func_name: str = "__midas_column_same_length__"
# Efficiently compute length
# https://stackoverflow.com/a/15943975/11109181
def len_of_col(col: ast.expr) -> ast.expr:
return ast.Call(
func=ast.Name(id="len"),
args=[
ast.Attribute(
value=col,
attr="index",
)
],
keywords=[],
)
self.assertions.define(
func_name,
ast.FunctionDef(
@@ -372,16 +387,10 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
body=[
ast.Return(
value=ast.Compare(
left=ast.Attribute(
value=ast.Name(id="column1"),
attr="size",
),
left=len_of_col(ast.Name(id="column1")),
ops=[ast.Eq()],
comparators=[
ast.Attribute(
value=ast.Name(id="column2"),
attr="size",
)
len_of_col(ast.Name(id="column2")),
],
)
)

View File

@@ -434,6 +434,21 @@ class FrameMethodRegistry(MethodRegistry[Call]):
def _assert_same_length(self, call_expr: p.Expr, frame1: p.Expr, frame2: p.Expr):
func_name: str = "__midas_frame_same_length__"
# Efficiently compute length
# https://stackoverflow.com/a/15943975/11109181
def len_of_df(df: ast.expr) -> ast.expr:
return ast.Call(
func=ast.Name(id="len"),
args=[
ast.Attribute(
value=df,
attr="index",
)
],
keywords=[],
)
self.assertions.define(
func_name,
ast.FunctionDef(
@@ -451,17 +466,9 @@ class FrameMethodRegistry(MethodRegistry[Call]):
body=[
ast.Return(
value=ast.Compare(
left=ast.Attribute(
value=ast.Name(id="frame1"),
attr="size",
),
left=len_of_df(ast.Name(id="frame1")),
ops=[ast.Eq()],
comparators=[
ast.Attribute(
value=ast.Name(id="frame2"),
attr="size",
)
],
comparators=[len_of_df(ast.Name(id="frame2"))],
)
)
],