examples/02_demonstration/weather/custom_types.midas

@@ -5,7 +5,7 @@ type Celsius = float
 type Kelvin = float where _ >= 0
 type Hectopascal = float
 
-type Temperature = Celsius where in_range(-30.0, 100.0)
+type Temperature = Celsius where in_range(-30.0, 100.0)(_)
 type Pressure = Hectopascal where in_range(800.0, 1100.0)(_)
 type Humidity = float where is_percentage(_)
 type HeatIndex = float
@@ -38,7 +38,7 @@ alias RawData = Frame[
 
 alias Data = Frame[
     station_id: StationID,
-    timestamp: object,
+    timestamp: Any,
     temperature: Temperature,
     pressure: Pressure,
     humidity: Humidity,
@@ -46,7 +46,7 @@ alias Data = Frame[
 
 alias DataWithHI = Frame[
     station_id: StationID,
-    timestamp: object,
+    timestamp: Any,
     temperature: Temperature,
     pressure: Pressure,
     humidity: Humidity,
@@ -54,12 +54,12 @@ alias DataWithHI = Frame[
 ]
 
 alias DailyAverages = Frame[
-    timestamp: object,
+    timestamp: Any,
     temperature: Mean[Temperature],
     pressure: Mean[Pressure],
     humidity: Mean[Humidity],
     heat_index: Mean[HeatIndex],
 ]
 
-predicate limit_amplitude(max_amp: float)(ls: list[float]) = max(ls) - min(ls) <= max_amp
-type LowAmplitudeWave = list[float where _ >= 1] where limit_amplitude(10)(_)
+// predicate limit_amplitude(max_amp: float)(ls: list[float]) = max(ls) - min(ls) <= max_amp
+// type LowAmplitudeWave = list[float where _ >= 1] where limit_amplitude(10)(_)

examples/02_demonstration/weather/gen_data.py

@@ -11,7 +11,7 @@ delta = end_ts - start_ts
 
 min_temp, max_temp = -30.0, 100.0
 min_pres, max_pres = 800.0, 1100.0
-min_hum, max_hum = 0.0, 1.0
+min_hum, max_hum = 0.0, 100.0
 
 N = 3000
 

examples/02_demonstration/weather/pipeline.py

@@ -47,7 +47,7 @@ def daily_avg(df: DataWithHI):
         DailyAverages,
         df.groupby(
             by=[
-                df["station_id"],
+                "station_id",
                 df["timestamp"].dt.day.rename("day"),
             ],
         )

midas/checker/frames/frame_groupby_methods.py

@@ -51,24 +51,25 @@ class FrameGroupByMethodRegistry(MethodRegistry[Call]):
         new_columns: list[DataFrameType.Column] = []
 
         for column in call.groupby.frame.columns:
-            column_groupby: ColumnGroupBy = ColumnGroupBy(column=column.type)
-            result_type: Type = self.typer.call_method(
-                location=call.location,
-                call_expr=call.call_expr,
-                obj=(call.groupby_expr, column_groupby),
-                method_name=method,
-                positional=call.positional,
-                keywords=call.keywords,
-            )
-            if not isinstance(result_type, ColumnType):
-                result_type = ColumnType(type=UnknownType())
-            new_columns.append(
-                DataFrameType.Column(
-                    index=column.index,
-                    name=column.name,
-                    type=result_type,
+            with self.reporter.with_context(f"in column '{column.name}'"):
+                column_groupby: ColumnGroupBy = ColumnGroupBy(column=column.type)
+                result_type: Type = self.typer.call_method(
+                    location=call.location,
+                    call_expr=call.call_expr,
+                    obj=(call.groupby_expr, column_groupby),
+                    method_name=method,
+                    positional=call.positional,
+                    keywords=call.keywords,
+                )
+                if not isinstance(result_type, ColumnType):
+                    result_type = ColumnType(type=UnknownType())
+                new_columns.append(
+                    DataFrameType.Column(
+                        index=column.index,
+                        name=column.name,
+                        type=result_type,
+                    )
                 )
-            )
 
         return DataFrameType(columns=new_columns)
 

midas/checker/frames/frame_methods.py

@@ -159,7 +159,10 @@ class FrameMethodRegistry(MethodRegistry[Call]):
                 col_type2 = ColumnType(type=operand[1])
 
             if col_type2 is not None:
-                col_type = self._get_method_result(call, col_type1, col_type2, method)
+                with self.reporter.with_context(f"in column '{column.name}'"):
+                    col_type = self._get_method_result(
+                        call, col_type1, col_type2, method
+                    )
 
             new_column = DataFrameType.Column(
                 index=column.index,
@@ -595,8 +598,62 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         )
         return result.result
 
+    def _filter_groupby_columns(
+        self, frame: DataFrameType, by: TypedExpr
+    ) -> DataFrameType:
+        """Remove columns passed as string literals in groupby's `by` argument
+
+        Args:
+            frame (DataFrameType): the original dataframe
+            by (TypedExpr): the by argument
+
+        Returns:
+            DataFrameType: the filtered dataframe
+        """
+        by_columns: list[str] = []
+
+        by_expr, _ = by
+
+        match by_expr:
+            case p.ListExpr(items=items):
+                for item in items:
+                    match item:
+                        case p.LiteralExpr(value=str() as name):
+                            by_columns.append(name)
+
+            case p.LiteralExpr(value=str() as name):
+                by_columns.append(name)
+
+        if len(by_columns) == 0:
+            return frame
+
+        new_columns: list[DataFrameType.Column] = []
+        for column in frame.columns:
+            if column.name in by_columns:
+                continue
+            new_columns.append(
+                DataFrameType.Column(
+                    index=len(new_columns),
+                    name=column.name,
+                    type=column.type,
+                )
+            )
+
+        return DataFrameType(columns=new_columns)
+
     @method()
     def groupby(self, call: Call) -> Type:
+        new_frame: DataFrameType = call.frame
+
+        by: Optional[TypedExpr] = None
+        if len(call.positional) != 0:
+            by = call.positional[0]
+        elif "by" in call.keywords:
+            by = call.keywords["by"]
+
+        if by is not None:
+            new_frame = self._filter_groupby_columns(call.frame, by)
+
         bool_: Type = self.types.get_type("bool")
         function: Function = Function(
             params=ParamSpec(
@@ -626,7 +683,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
                     )
                 ],
             ),
-            returns=FrameGroupBy(frame=call.frame),
+            returns=FrameGroupBy(frame=new_frame),
         )
 
         result: CallResult = self.dispatcher.get_result(

midas/checker/midas.py

@@ -407,8 +407,18 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
             return UnknownType()
 
     def visit_constraint_type(self, type: m.ConstraintType) -> Type:
+        base_type: Type = type.type.accept(self)
+        self._predicate_params["_"] = base_type
+        constraint_type: Type = self.type_of(type.constraint)
+        self._predicate_params = {}
+        if not self.types.is_subtype(constraint_type, self._bool):
+            self.reporter.error(
+                type.location,
+                f"Constraint must evaluate to a boolean, got {constraint_type}",
+            )
+
         return ConstraintType(
-            type=type.type.accept(self),
+            type=base_type,
             constraint=type.constraint,
         )
 

midas/checker/python.py

@@ -257,6 +257,9 @@ class PythonTyper(
         """
         unfolded: Type = unfold_type(obj[1])
         match unfolded:
+            case TopType() | UnknownType():
+                return UnknownType()
+
             case DataFrameType():
                 return self.frame_mgr.call(
                     method=method_name,

midas/checker/reporter.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from contextlib import contextmanager
 from typing import Optional
 
 from midas.ast.location import Location
@@ -54,6 +55,7 @@ class FileReporter:
     def __init__(self, base_reporter: Reporter, path: Optional[str]) -> None:
         self.base_reporter: Reporter = base_reporter
         self.path: Optional[str] = path
+        self._context: list[str] = []
 
     def for_file(self, path: Optional[str]) -> FileReporter:
         """Create a new file reporter for the given path with the same base reporter
@@ -66,6 +68,14 @@ class FileReporter:
         """
         return FileReporter(self.base_reporter, path)
 
+    @contextmanager
+    def with_context(self, ctx: str):
+        self._context.append(ctx)
+        try:
+            yield
+        finally:
+            self._context.pop()
+
     def report(self, type: DiagnosticType, location: Location, message: str):
         """Report a diagnostic to the base reporter
 
@@ -74,6 +84,8 @@ class FileReporter:
             location (Location): the location of the diagnostic in the file
             message (str): the diagnostic's message
         """
+        for ctx in self._context:
+            message = message + ", " + ctx
         self.base_reporter.report(self.path, type, location, message)
 
     def error(self, location: Location, message: str):