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

@@ -35,11 +35,19 @@ class Call:
 
 
 class ColumnMethodRegistry(MethodRegistry[Call]):
-    @method("add", "__add__")
+    def _element_binary_op(self, call: Call, method: str) -> ColumnType:
-    def add(self, call: Call) -> Type:
+        """Compute the result of an element-wise binary operation
-        # TODO: support add with scalar
-        # TODO: check operation exists on inner column types
 
+        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__"),
             ),
         )
 

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
+        result: Type = self.typer.result_of_binary_op(
-        try:
+            location=call.location,
-            result = (
+            expr=call.call_expr,
-                self.typer.call_method(
+            left=(call.frame_expr, column1),
-                    location=call.location,
+            right=(call.positional[0][0], column2),
-                    call_expr=call.call_expr,
+            method=method,
-                    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}",
-            )
 
         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]
+            operand: TypedExpr = call.positional[0]
-            unfolded_other: Type = unfold_type(other)
+            unfolded_other: Type = unfold_type(operand[1])
             if isinstance(unfolded_other, DataFrameType):
                 frame2 = unfolded_other
                 by_name = {

midas/checker/python.py

@@ -543,8 +543,14 @@ class PythonTyper(
             )
             return UnknownType()
 
-        return self._visit_binary_expr(
+        left: Type = self.type_of(expr.left)
-            expr.location, expr, expr.left, expr.right, method
+        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(
+        left: Type = self.type_of(expr.left)
-            expr.location, expr, expr.left, expr.right, method
+        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,
+        left: TypedExpr,
-        right_expr: p.Expr,
+        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()