From be2fd4c8374ca62a2db85c8bbf01a2ba8361d40c Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 00:05:40 +0200 Subject: [PATCH] feat(checker): delegate element operation to inner type delegate element-wise binary operation on columns to their inner types --- midas/checker/frames/column_methods.py | 34 +++++++++++++++++++------ midas/checker/frames/frame_methods.py | 35 +++++++++----------------- midas/checker/python.py | 35 ++++++++++++++++---------- 3 files changed, 61 insertions(+), 43 deletions(-) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 91909ff..e8a2ce6 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -35,11 +35,19 @@ class Call: class ColumnMethodRegistry(MethodRegistry[Call]): - @method("add", "__add__") - def add(self, call: Call) -> Type: - # TODO: support add with scalar - # TODO: check operation exists on inner column types + def _element_binary_op(self, call: Call, method: str) -> ColumnType: + """Compute the result of an element-wise binary operation + This function delegates to the inner types for computing the resulting + type. + + Args: + call (Call): the call that triggered this resolution + method (str): the method name + + Returns: + ColumnType: the resulting column type + """ column2: Optional[ColumnType] = None col_type1: Type = call.column.type @@ -50,8 +58,20 @@ class ColumnMethodRegistry(MethodRegistry[Call]): if isinstance(unfolded_other, ColumnType): column2 = unfolded_other col_type2: Type = column2.type - if self.types.are_equivalent(col_type2, col_type1): - new_column = ColumnType(type=col_type1) + + new_inner_type = self.typer.result_of_binary_op( + location=call.location, + expr=call.call_expr, + left=(call.column_expr, col_type1), + right=(call.positional[0][0], col_type2), + method=method, + ) + new_column = ColumnType(type=new_inner_type) + return new_column + + @method("add", "__add__") + def add(self, call: Call) -> Type: + # TODO: support add with scalar # Build signature with new column type and generic operand param_type: TypeVar = TypeVar(name="T", bound=None) @@ -67,7 +87,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): required=True, ), ], - returns=new_column, + returns=self._element_binary_op(call, "__add__"), ), ) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index a43f1fc..64c8d06 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -21,7 +21,7 @@ from midas.checker.types import ( ) if TYPE_CHECKING: - from midas.checker.python import TypedExpr, UndefinedMethodException + from midas.checker.python import TypedExpr @dataclass(frozen=True, kw_only=True) @@ -60,28 +60,17 @@ class FrameMethodRegistry(MethodRegistry[Call]): If the operation is invalid / doesn't exist, `ColumnType(type=UnknownType())` is returned """ - unknown: ColumnType = ColumnType(type=UnknownType()) - result: Type = unknown - try: - result = ( - self.typer.call_method( - location=call.location, - call_expr=call.call_expr, - obj=(call.frame_expr, column1), - method_name=method, - positional=[(call.positional[0][0], column2)], - keywords={}, - ) - or unknown - ) - except UndefinedMethodException: - self.reporter.error( - call.location, - f"Undefined operation {method} between {column1} and {column2}", - ) + + result: Type = self.typer.result_of_binary_op( + location=call.location, + expr=call.call_expr, + left=(call.frame_expr, column1), + right=(call.positional[0][0], column2), + method=method, + ) if not isinstance(result, ColumnType): - result = unknown + return ColumnType(type=UnknownType()) return result def _element_binary_op(self, call: Call, method: str) -> DataFrameType: @@ -105,8 +94,8 @@ class FrameMethodRegistry(MethodRegistry[Call]): frame2: Optional[DataFrameType] = None # Get map of operand's columns by name, if there is at least 1 operand, which is a dataframe if len(call.positional) != 0: - other: Type = call.positional[0][1] - unfolded_other: Type = unfold_type(other) + operand: TypedExpr = call.positional[0] + unfolded_other: Type = unfold_type(operand[1]) if isinstance(unfolded_other, DataFrameType): frame2 = unfolded_other by_name = { diff --git a/midas/checker/python.py b/midas/checker/python.py index 8e4b59b..d095162 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -543,8 +543,14 @@ class PythonTyper( ) return UnknownType() - return self._visit_binary_expr( - expr.location, expr, expr.left, expr.right, method + left: Type = self.type_of(expr.left) + right: Type = self.type_of(expr.right) + return self.result_of_binary_op( + expr.location, + expr, + (expr.left, left), + (expr.right, right), + method, ) def visit_compare_expr(self, expr: p.CompareExpr) -> Type: @@ -556,35 +562,38 @@ class PythonTyper( ) return UnknownType() - return self._visit_binary_expr( - expr.location, expr, expr.left, expr.right, method + left: Type = self.type_of(expr.left) + right: Type = self.type_of(expr.right) + return self.result_of_binary_op( + expr.location, + expr, + (expr.left, left), + (expr.right, right), + method, ) - def _visit_binary_expr( + def result_of_binary_op( self, location: Location, expr: p.Expr, - left_expr: p.Expr, - right_expr: p.Expr, + left: TypedExpr, + right: TypedExpr, method: str, ) -> Type: - left: Type = self.type_of(left_expr) - right: Type = self.type_of(right_expr) - result: Optional[Type] try: result = self.call_method( location=location, call_expr=expr, - obj=(left_expr, left), + obj=left, method_name=method, - positional=[(right_expr, right)], + positional=[right], keywords={}, ) except UndefinedMethodException: self.reporter.error( location, - f"Undefined operation {method} between {left} and {right}", + f"Undefined operation {method} between {left[1]} and {right[1]}", ) return UnknownType()