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..dfe0698 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, 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):