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 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__"),
             ),
         )
 

midas/checker/frames/frame_manager.py

@@ -47,12 +47,41 @@ class FrameManager:
                 return self.assign_column(reporter, location, frame, name, value_type)
 
             case p.ListExpr(items=indices) if is_list_of_literals(indices) and all(
-                isinstance(idx, str) for idx in indices
+                isinstance(index.value, str) for index in indices
             ):
-                raise NotImplementedError
+                names: list[str] = [cast(str, index.value) for index in indices]
+
+                if not isinstance(value_type, TupleType):
+                    reporter.error(
+                        location,
+                        f"Cannot assign {type} to dataframe columns. Must be a tuple of columns",
+                    )
+                    return UnknownType()
+
+                if len(names) != len(value_type.items):
+                    reporter.error(
+                        location,
+                        f"Wrong number of columns. Cannot assign {len(value_type.items)} to {len(names)} targets",
+                    )
+                    return UnknownType()
+
+                new_frame: Type = frame
+                for name, value in zip(names, value_type.items):
+                    new_frame = self.assign_column(
+                        reporter,
+                        location,
+                        new_frame,
+                        name,
+                        value,
+                    )
+                    if not isinstance(new_frame, DataFrameType):
+                        return new_frame
+                return new_frame
 
             case _:
-                reporter.error(location, f"Invalid index type {index} on {frame}")
+                reporter.error(
+                    location, f"Invalid index type {index} on {frame} (assignment)"
+                )
                 return UnknownType()
 
     def assign_column(
@@ -100,7 +129,9 @@ class FrameManager:
                 return TupleType(items=tuple(columns))
 
             case _:
-                reporter.error(location, f"Invalid index type {index} on {frame}")
+                reporter.error(
+                    location, f"Invalid index type {index} on {frame} (access)"
+                )
                 return UnknownType()
 
     def groupby_get(

midas/checker/frames/frame_methods.py

@@ -35,19 +35,67 @@ class Call:
 
 
 class FrameMethodRegistry(MethodRegistry[Call]):
-    @method("add", "__add__")
-    def add(self, call: Call) -> Type:
-        # TODO: support add with scalar, sequence, Series, dict
-        # TODO: check operation exists on inner column types
+    def _get_method_result(
+        self,
+        call: Call,
+        column1: ColumnType,
+        column2: ColumnType,
+        method: str,
+    ) -> ColumnType:
+        """Get the result of calling a method on a column, passing a second
 
+        This function delegates to the main typer the resolution of the method
+        member, as well as computing the result type. Because we don't have any
+        AST expression for the individual columns, the frame expressions are
+        used instead.
+
+        Args:
+            call (Call): the call that triggered this resolution
+            column1 (ColumnType): the first column, i.e. left operand
+            column2 (ColumnType): the second column, i.e. right operand
+            method (str): the method name
+
+        Returns:
+            ColumnType: the resulting column.
+                If the operation is invalid / doesn't exist,
+                `ColumnType(type=UnknownType())` is returned
+        """
+
+        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):
+            return ColumnType(type=UnknownType())
+        return result
+
+    def _element_binary_op(self, call: Call, method: str) -> DataFrameType:
+        """Compute the result of an element-wise binary operation
+
+        This function delegates to the matching columns for computing resulting
+        types. Any column only present in one of the frames is forwarded as a
+        generic `ColumnType(type=UnknownType())`. Columns only in the second
+        frame are append at the end of the schema.
+
+        Args:
+            call (Call): the call that triggered this resolution
+            method (str): the method name
+
+        Returns:
+            DataFrameType: the resulting frame type
+        """
         new_columns: list[DataFrameType.Column] = []
 
         by_name: dict[str, DataFrameType.Column] = {}
         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 = {
@@ -56,20 +104,20 @@ class FrameMethodRegistry(MethodRegistry[Call]):
 
         # Compute new schema:
         # Step 1: for all columns in frame1:
-        #  - if present in frame2 with equivalent type -> add to schema as is
+        #  - if present in frame2 -> delegate operation to columns
         #  - if not -> add to schema as unknown
         in_frame1: set[str] = set()
         for column in call.frame.columns:
             if column.name is not None:
                 in_frame1.add(column.name)
 
-            col_type1: Type = column.type
-            col_type: Type = ColumnType(type=UnknownType())
+            col_type1: ColumnType = column.type
+            col_type: ColumnType = ColumnType(type=UnknownType())
             if column.name in by_name:
                 column2 = by_name[column.name]
-                col_type2: Type = column2.type
-                if self.types.are_equivalent(col_type2, col_type1):
-                    col_type = col_type1
+                col_type2: ColumnType = column2.type
+
+                col_type = self._get_method_result(call, col_type1, col_type2, method)
 
             new_column = DataFrameType.Column(
                 index=column.index,
@@ -92,6 +140,12 @@ class FrameMethodRegistry(MethodRegistry[Call]):
                     )
                 )
 
+        return DataFrameType(columns=new_columns)
+
+    @method("add", "__add__")
+    def add(self, call: Call) -> Type:
+        # TODO: support add with scalar, sequence, Series, dict
+
         # Build signature with new schema and generic operand
         signature = Function(
             args=[
@@ -102,7 +156,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
                     required=True,
                 ),
             ],
-            returns=DataFrameType(columns=new_columns),
+            returns=self._element_binary_op(call, "__add__"),
         )
 
         # Map arguments and compute result type

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()