diff --git a/midas/checker/frames/column_manager.py b/midas/checker/frames/column_manager.py index 8adeb1b..4fa9e81 100644 --- a/midas/checker/frames/column_manager.py +++ b/midas/checker/frames/column_manager.py @@ -4,11 +4,20 @@ from typing import TYPE_CHECKING, Optional import midas.ast.python as p from midas.ast.location import Location +from midas.checker.dispatcher import CallResult from midas.checker.frames.column_groupby_methods import Call as GroupByCall from midas.checker.frames.column_groupby_methods import ColumnGroupByMethodRegistry from midas.checker.frames.column_methods import Call, ColumnMethodRegistry from midas.checker.registry import TypesRegistry -from midas.checker.types import ColumnGroupBy, ColumnType, Type +from midas.checker.reporter import FileReporter +from midas.checker.types import ( + ColumnGroupBy, + ColumnType, + Function, + OverloadedFunction, + ParamSpec, + Type, +) if TYPE_CHECKING: from midas.checker.python import PythonTyper, TypedExpr @@ -24,6 +33,62 @@ class ColumnManager: ColumnGroupByMethodRegistry(self.typer) ) + def get( + self, + reporter: FileReporter, + location: Location, + column: ColumnType, + index: TypedExpr, + ) -> Type: + """Compute the type of a subscript access + + Args: + reporter (FileReporter): the file reporter to use for diagnostics + location (Location): the subscript's location + column (DataFrameType): the column type + index (TypedExpr): the index + + Returns: + Type: the resulting type + """ + + single = Function( + params=ParamSpec( + pos=[ + Function.Parameter( + pos=0, + name="index", + type=self.typer.types.get_type("int"), + required=True, + ) + ] + ), + returns=column.type, + ) + slice = Function( + params=ParamSpec( + pos=[ + Function.Parameter( + pos=0, + name="slice", + type=self.typer.types.get_type("slice"), + required=True, + ) + ] + ), + returns=column, + ) + + overload = OverloadedFunction(overloads=[single, slice]) + + result: CallResult = self.typer.dispatcher.get_result( + location=location, + callee=overload, + positional=[index], + keywords={}, + ) + return result.result + def call( self, method: str, diff --git a/midas/checker/python.py b/midas/checker/python.py index 779a480..674727d 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -889,6 +889,8 @@ class PythonTyper( return self._visit_frame_subscript(unfolded, expr) case FrameGroupBy(): return self._visit_frame_groupby_subscript(unfolded, expr) + case ColumnType(): + return self._visit_column_subscript(unfolded, expr) operation: Optional[Type] = self.types.lookup_member(object, "__getitem__") if operation is None: @@ -1274,3 +1276,14 @@ class PythonTyper( return self.frame_mgr.groupby_get( self.reporter, expr.location, groupby, expr.index ) + + def _visit_column_subscript( + self, column: ColumnType, expr: p.SubscriptExpr + ) -> Type: + index_type: Type = self.type_of(expr.index) + return self.column_mgr.get( + self.reporter, + expr.location, + column, + (expr.index, index_type), + )