tests: update with new aggregate return types

examples/02_demonstration/weather/pipeline.py

@@ -8,6 +8,7 @@ from midas.typing import Column, cast, unsafe_cast
 
 
 def load_data(path: Path) -> RawData:
+    # Check base types and dataframe structure
     return cast(RawData, pd.read_csv(path))
 
 
@@ -17,10 +18,15 @@ def convert_data(raw_df: RawData) -> Data:
         Column[object],
         pd.to_datetime(new_df["timestamp"]),
     )
+
+    # Check types and constraints at runtime, catches out-of-range values and
+    # invalid types / malformed data
     return cast(Data, new_df)
 
 
 def compute_heat_index(df: Data):
+    # The computation's result can only be typed as `Column[float]`
+    # Casting is necessary to bring back semantic
     df["heat_index"] = cast(
         Column[HeatIndex],
         (
@@ -33,6 +39,10 @@ def compute_heat_index(df: Data):
 
 
 def daily_avg(df: DataWithHI):
+    # Group-by and aggregation methods keep the structure of the dataframe but
+    # may erase the exact types
+    # The type checker is still very conservative and often the result of most
+    # aggregation methods as `Column[Any]`
     return cast(
         DailyAverages,
         df.groupby(
@@ -47,21 +57,23 @@ def daily_avg(df: DataWithHI):
 
 
 def plot(df: DailyAverages):
+    # Some operations are not implemented in Midas but the user can still use
+    # them, they will just not be fully type-checked
+    # `unsafe_cast` can also be used to avoid trivial, redundant or costly checks
     stations = unsafe_cast(list[str], list(df.index.get_level_values(0).unique()))
     for station in stations:
         sub_df = unsafe_cast(DailyAverages, df.loc[station])
-        # plt.plot(sub_df["timestamp"], sub_df["temperature"])
         plt.plot(sub_df["timestamp"], sub_df["heat_index"])
     plt.show()
 
 
 def main():
+    # Assigning to annotated variables help catch errors
     raw_df: RawData = load_data(Path("data.csv"))
     df: Data = convert_data(raw_df)
 
     with_hi = compute_heat_index(df)
     dailies = daily_avg(with_hi)
-    print(dailies)
     plot(dailies)
 
 

midas/checker/frames/column_groupby_methods.py

@@ -12,8 +12,8 @@ from midas.checker.types import (
     ColumnType,
     Function,
     ParamSpec,
-    TopType,
     Type,
+    UnknownType,
 )
 
 if TYPE_CHECKING:
@@ -49,21 +49,18 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def _aggregate(
         self,
         call: Call,
+        method: str,
         params: list[str | tuple[str, str, bool]] = [],
-        *,
-        preserve_inner_type: bool = False,
     ) -> Type:
         """Compute the result type of an aggregate method call
 
         Args:
             call (Call): the call object
+            method (str): the method name to delegate on :class:`Column`
             params (list[str | tuple[str, str, bool], optional): a list of extra
                 mixed parameters. The list can contain strings to include
                 parameters predefined in `NAMED_ARGS`, or tuples containing the
                 parameter's name, type and required flag. Defaults to [].
-            preserve_inner_type (bool, optional): If `True`, the result type
-                will preserve the column's inner type (e.g. for `min`/`max`),
-                otherwise the inner type is widened to `TopType`. Defaults to False.
 
         Returns:
             Type: the result type
@@ -87,13 +84,21 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
                     )
             real_params.append(param)
 
+        # TODO: maybe better to filter arguments and pass some, in case the
+        # return type depends on them
+        returns: Type = self.typer.call_method(
+            location=call.location,
+            call_expr=call.call_expr,
+            obj=(call.groupby_expr, call.groupby.column),
+            method_name=method,
+            positional=[],
+            keywords={},
+        )
+        if not isinstance(returns, ColumnType):
+            returns = ColumnType(type=UnknownType())
         signature = Function(
             params=ParamSpec(mixed=real_params),
-            returns=(
-                call.groupby.column
-                if preserve_inner_type
-                else ColumnType(type=TopType())
-            ),
+            returns=returns,
         )
 
         result: CallResult = self.dispatcher.get_result(
@@ -108,6 +113,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def kurt(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "kurt",
             ["skipna", "numeric_only"],
         )
 
@@ -115,6 +121,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def max(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "max",
             [
                 "numeric_only",
                 (
@@ -126,13 +133,13 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
                 "engine",
                 "engine_kwargs",
             ],
-            preserve_inner_type=True,
         )
 
     @method()
     def mean(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "mean",
             ["numeric_only", "skipna", "engine", "engine_kwargs"],
         )
 
@@ -140,14 +147,15 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def median(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "median",
             ["numeric_only", "skipna"],
-            preserve_inner_type=True,
         )
 
     @method()
     def min(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "min",
             [
                 "numeric_only",
                 (
@@ -159,13 +167,13 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
                 "engine",
                 "engine_kwargs",
             ],
-            preserve_inner_type=True,
         )
 
     @method()
     def prod(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "prod",
             [
                 "numeric_only",
                 (
@@ -181,6 +189,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def std(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "std",
             [
                 (
                     "ddof",
@@ -198,6 +207,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def sum(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "sum",
             [
                 "numeric_only",
                 (
@@ -215,6 +225,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def var(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "var",
             [
                 (
                     "var",

midas/checker/frames/column_methods.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import ast
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable, Optional, TypeAlias, Union
 
 import midas.ast.python as p
 from midas.ast.location import Location
@@ -24,6 +24,31 @@ from midas.checker.types import (
 if TYPE_CHECKING:
     from midas.checker.python import TypedExpr
 
+FormulaOperand: TypeAlias = Union["Formula", str, Type]
+"""
+A operand type in a :data:`Formula`
+
+Must be one of the following:
+- a nested formula
+- a type name (a string)
+- a type instance
+"""
+
+Formula: TypeAlias = Union[Type, tuple[FormulaOperand, str, FormulaOperand]]
+"""
+A formula to compute the output type of a function
+
+Must be either a type, or a tuple containing:
+- a left operand
+- an operation / method name (e.g. `"__add__"`)
+- a right operand
+
+For example, to compute the result of a `mean` function, given the input type `T`:
+```python
+mean_formula = ((T, "__add__", T), "__truediv__", "int")
+```
+"""
+
 
 @dataclass(frozen=True, kw_only=True)
 class Call:
@@ -44,6 +69,52 @@ class Call:
 class ColumnMethodRegistry(MethodRegistry[Call]):
     """The method registry for column types"""
 
+    def _resolve_formula_operand(self, call: Call, operand: FormulaOperand) -> Type:
+        """Resolve the type of a formula operand
+
+        See :data:`FormulaOperand` for more information on the accepted format
+
+        Args:
+            call (Call): the call that triggered this resolution
+            operand (FormulaOperand): the formula operand
+
+        Returns:
+            Type: the type of the operand
+        """
+        match operand:
+            case str():
+                return self.types.get_type(operand)
+            case (_, _, _):
+                return self._resolve_formula_type(call, operand)
+            case _:
+                return operand
+
+    def _resolve_formula_type(self, call: Call, formula: Formula) -> Type:
+        """Resolve the return type of a formula
+
+        See :data:`Formula` for more information on the accepted format
+
+        Args:
+            call (Call): the call that triggered this resolution
+            formula (Formula): the formula to evaluate
+
+        Returns:
+            Type: the return type of the formula
+        """
+        if not isinstance(formula, tuple):
+            return formula
+
+        op1, operator, op2 = formula
+        op1_type: Type = self._resolve_formula_operand(call, op1)
+        op2_type: Type = self._resolve_formula_operand(call, op2)
+        return self.typer.result_of_binary_op(
+            location=call.location,
+            expr=call.call_expr,
+            left=(call.column_expr, op1_type),
+            right=(call.column_expr, op2_type),
+            method=operator,
+        )
+
     def _simple_call(self, call: Call, function: Type) -> Type:
         """Get the result of calling a simple method
 
@@ -274,7 +345,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         call: Call,
         kwargs: list[Function.Parameter] = [],
         *,
-        preserve_inner_type: bool = False,
+        formula: Optional[Callable[[Type], Formula]] = None,
     ) -> Type:
         """Compute the result type of an aggregate method call
 
@@ -282,13 +353,23 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
             call (Call): the call object
             kwargs (list[Function.Parameter], optional): a list of extra
                 keyword-only parameters. Defaults to [].
-            preserve_inner_type (bool, optional): If `True`, the result type
-                will preserve the column's inner type (e.g. for `min`/`max`),
-                otherwise the inner type is widened to `TopType`. Defaults to False.
+            formula (Callable[[Type], Formula], optional): optional formula
+                builder function to compute the return type. If set, the function
+                should accept the inner column type and return a formula.
+                If `None`, the result is typed as `Column[Any]`. Defaults to None.
 
         Returns:
             Type: the result type
         """
+
+        returns: Type = ColumnType(type=TopType())
+        if formula:
+            returns = ColumnType(
+                type=self._resolve_formula_type(
+                    call,
+                    formula(call.column.type),
+                )
+            )
         signature = Function(
             params=ParamSpec(
                 kw=[
@@ -301,7 +382,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
                     *kwargs,
                 ],
             ),
-            returns=call.column if preserve_inner_type else ColumnType(type=TopType()),
+            returns=returns,
         )
 
         result: CallResult = self.dispatcher.get_result(
@@ -318,27 +399,29 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
 
     @method()
     def max(self, call: Call) -> Type:
-        return self._aggregate(call, preserve_inner_type=True)
+        return self._aggregate(call, formula=lambda t: t)
 
     @method()
     def mean(self, call: Call) -> Type:
-        return self._aggregate(call)
+        return self._aggregate(
+            call, formula=lambda t: ((t, "__add__", t), "__truediv__", "int")
+        )
 
     @method()
     def median(self, call: Call) -> Type:
-        return self._aggregate(call, preserve_inner_type=True)
+        return self._aggregate(call, formula=lambda t: t)
 
     @method()
     def min(self, call: Call) -> Type:
-        return self._aggregate(call, preserve_inner_type=True)
+        return self._aggregate(call, formula=lambda t: t)
 
     @method()
     def mode(self, call: Call) -> Type:
-        return self._aggregate(call, preserve_inner_type=True)
+        return self._aggregate(call, formula=lambda t: t)
 
     @method("product", "prod")
     def product(self, call: Call) -> Type:
-        return self._aggregate(call)
+        return self._aggregate(call, formula=lambda t: (t, "__mul__", t))
 
     @method()
     def std(self, call: Call) -> Type:
@@ -356,7 +439,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
 
     @method()
     def sum(self, call: Call) -> Type:
-        return self._aggregate(call)
+        return self._aggregate(call, formula=lambda t: (t, "__add__", t))
 
     @method()
     def var(self, call: Call) -> Type:

midas/cli/hl_diagnostic.css

@@ -4,9 +4,11 @@ span {
     &.error {
         --col: 255, 0, 0;
     }
+
     &.warning {
         --col: 250, 160, 0;
     }
+
     &.info {
         --col: 150, 190, 250;
     }
@@ -19,12 +21,12 @@ span {
         }
 
         &:hover:not(:has(.with-msg:hover)) {
-            .message {
+            &>.message {
                 display: inline-block;
             }
         }
 
-        .message {
+        &>.message {
             position: absolute;
             top: calc(100% + 0.2em);
             left: -.2em;
@@ -33,7 +35,8 @@ span {
             padding: 0.2em 0.4em;
             border-radius: .2em;
             z-index: 10;
-            width: 300%;
+            width: max-content;
+            max-width: 60vw;
         }
     }
 }

tests/cases/checker/09_frame_ops.py.ref.json

@@ -2296,14 +2296,18 @@
             "index": 0,
             "name": "a",
             "type": {
-              "type": {}
+              "type": {
+                "name": "float"
+              }
             }
           },
           {
             "index": 1,
             "name": "b",
             "type": {
-              "type": {}
+              "type": {
+                "name": "float"
+              }
             }
           }
         ]
@@ -2517,14 +2521,18 @@
             "index": 0,
             "name": "a",
             "type": {
-              "type": {}
+              "type": {
+                "name": "int"
+              }
             }
           },
           {
             "index": 1,
             "name": "b",
             "type": {
-              "type": {}
+              "type": {
+                "name": "float"
+              }
             }
           }
         ]
@@ -2659,14 +2667,18 @@
             "index": 0,
             "name": "a",
             "type": {
-              "type": {}
+              "type": {
+                "name": "int"
+              }
             }
           },
           {
             "index": 1,
             "name": "b",
             "type": {
-              "type": {}
+              "type": {
+                "name": "float"
+              }
             }
           }
         ]
@@ -3687,7 +3699,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "float"
+        }
       }
     },
     {
@@ -3841,7 +3855,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "int"
+        }
       }
     },
     {
@@ -3878,7 +3894,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "int"
+        }
       }
     },
     {
@@ -3952,7 +3970,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "int"
+        }
       }
     },
     {
@@ -4167,7 +4187,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "float"
+        }
       }
     },
     {
@@ -4288,7 +4310,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "int"
+        }
       }
     },
     {
@@ -4366,7 +4390,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "int"
+        }
       }
     },
     {