From 5051e155c0b1c9692459a20366818a784d9048e9 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 7 Jul 2026 11:22:07 +0200 Subject: [PATCH] fix(gen): only generate length assertion for non-scalar ops --- midas/checker/frames/column_methods.py | 18 +++++++++++------- midas/checker/frames/frame_methods.py | 18 +++++++++++------- 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 80792d0..a52c2ed 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -63,7 +63,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): ) return result.result - def _element_binary_op(self, call: Call, method: str) -> Type: + def _element_binary_op(self, call: Call, method: str) -> tuple[Type, bool]: """Compute the result of an element-wise binary operation This function delegates to the inner types for computing the resulting @@ -74,18 +74,21 @@ class ColumnMethodRegistry(MethodRegistry[Call]): method (str): the method name Returns: - Type: the resulting type + tuple[Type, bool]: the resulting type and a boolean indicating + whether the operand is a column """ if len(call.positional) == 0: - return UnknownType() + return UnknownType(), False col_type1: Type = call.column.type operand: TypedExpr = call.positional[0] unfolded_operand: Type = unfold_type(operand[1]) col_type2: Type + column_operand: bool = isinstance(unfolded_operand, ColumnType) + # Operand is a column -> get the inner type - if isinstance(unfolded_operand, ColumnType): + if column_operand: col_type2 = unfolded_operand.type # Otherwise use the operand type itself else: @@ -98,7 +101,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): right=(operand[0], col_type2), method=method, ) - return ColumnType(type=new_inner_type) + return ColumnType(type=new_inner_type), column_operand def _element_wise(self, call: Call, method: str) -> Type: """Compute the result of an element-wise method call @@ -115,6 +118,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): """ # Build signature with new column type and generic operand + returns, column_operand = self._element_binary_op(call, method) signature = Function( params=ParamSpec( mixed=[ @@ -126,7 +130,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): ), ], ), - returns=self._element_binary_op(call, method), + returns=returns, ) # Map arguments and compute result type @@ -136,7 +140,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): positional=call.positional, keywords=call.keywords, ) - if result.is_valid: + if result.is_valid and column_operand: self._assert_same_length( call.call_expr, call.column_expr, call.positional[0][0] ) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 7d3d991..01f9a4c 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -102,7 +102,7 @@ class FrameMethodRegistry(MethodRegistry[Call]): return ColumnType(type=UnknownType()) return result - def _element_binary_op(self, call: Call, method: str) -> Type: + def _element_binary_op(self, call: Call, method: str) -> tuple[Type, bool]: """Compute the result of an element-wise binary operation This function delegates to the matching columns for computing resulting @@ -115,11 +115,12 @@ class FrameMethodRegistry(MethodRegistry[Call]): method (str): the method name Returns: - Type: the resulting type + tuple[Type, bool]: the resulting type and a boolean indicating + whether the operand is a frame """ if len(call.positional) == 0: - return UnknownType() + return UnknownType(), False operand: TypedExpr = call.positional[0] new_columns: list[DataFrameType.Column] = [] @@ -128,7 +129,8 @@ class FrameMethodRegistry(MethodRegistry[Call]): frame2: Optional[DataFrameType] = None # Get map of operand's columns by name, if the operand is a dataframe unfolded_other: Type = unfold_type(operand[1]) - if isinstance(unfolded_other, DataFrameType): + frame_operand: bool = isinstance(unfolded_other, DataFrameType) + if frame_operand: frame2 = unfolded_other by_name = {col.name: col for col in frame2.columns if col.name is not None} @@ -180,7 +182,7 @@ class FrameMethodRegistry(MethodRegistry[Call]): ) ) - return DataFrameType(columns=new_columns) + return DataFrameType(columns=new_columns), frame_operand def _element_wise(self, call: Call, method: str) -> Type: """Compute the result of an element-wise method call @@ -196,6 +198,8 @@ class FrameMethodRegistry(MethodRegistry[Call]): Type: the result type """ # TODO: support sequence, Series, dict operand + returns, frame_operand = self._element_binary_op(call, method) + # Build signature with new schema and generic operand signature = Function( params=ParamSpec( @@ -208,7 +212,7 @@ class FrameMethodRegistry(MethodRegistry[Call]): ), ], ), - returns=self._element_binary_op(call, method), + returns=returns, ) # Map arguments and compute result type @@ -218,7 +222,7 @@ class FrameMethodRegistry(MethodRegistry[Call]): positional=call.positional, keywords=call.keywords, ) - if result.is_valid: + if result.is_valid and frame_operand: self._assert_same_length( call.call_expr, call.frame_expr, call.positional[0][0] )