diff --git a/examples/02_demonstration/weather/custom_types.midas b/examples/02_demonstration/weather/custom_types.midas index 289b358..d945cab 100644 --- a/examples/02_demonstration/weather/custom_types.midas +++ b/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)(_) diff --git a/examples/02_demonstration/weather/gen_data.py b/examples/02_demonstration/weather/gen_data.py index bdd3b5f..f86467e 100644 --- a/examples/02_demonstration/weather/gen_data.py +++ b/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 diff --git a/examples/02_demonstration/weather/pipeline.py b/examples/02_demonstration/weather/pipeline.py index 3da4a49..3593f12 100644 --- a/examples/02_demonstration/weather/pipeline.py +++ b/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"), ], ) diff --git a/midas/checker/frames/frame_groupby_methods.py b/midas/checker/frames/frame_groupby_methods.py index cabf88a..a2efba6 100644 --- a/midas/checker/frames/frame_groupby_methods.py +++ b/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) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 99910e1..24a41f7 100644 --- a/midas/checker/frames/frame_methods.py +++ b/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( diff --git a/midas/checker/midas.py b/midas/checker/midas.py index 5213785..e8cf791 100644 --- a/midas/checker/midas.py +++ b/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, ) diff --git a/midas/checker/python.py b/midas/checker/python.py index 7569f1b..705498b 100644 --- a/midas/checker/python.py +++ b/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, diff --git a/midas/checker/reporter.py b/midas/checker/reporter.py index 970ef34..8604dcd 100644 --- a/midas/checker/reporter.py +++ b/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):