Merge pull request 'Column subscripts' (#30) from feat/column-subscript into main

Reviewed-on: #30

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,

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),
+        )