feat(checker): type check subscripts

examples/01_simple_type_checking/04_complex_types.py

@@ -28,3 +28,8 @@ bar: list[list[Meter]]
 bar.append([p2.x])
 
 foo2 = foo + foo
+
+a = foo[0]
+b = bar[0][1]
+c = bar[0][1][2]  # invalid, not method __getitem__ on Meter
+c = bar[""]  # invalid, wrong index type

gen/python.py

@@ -143,4 +143,9 @@ class ListExpr:
     items: list[Expr]
 
 
+class SubscriptExpr:
+    object: Expr
+    index: Expr
+
+
 ###<

midas/ast/printer.py

@@ -664,7 +664,7 @@ class PythonAstPrinter(
     def visit_literal_expr(self, expr: p.LiteralExpr) -> None:
         self._write_line("LiteralExpr")
         with self._child_level(single=True):
-            self._write_line(f"value: {expr.value}")
+            self._write_line(f"value: {expr.value!r}")
 
     def visit_variable_expr(self, expr: p.VariableExpr) -> None:
         self._write_line("VariableExpr")
@@ -719,3 +719,13 @@ class PythonAstPrinter(
                     if i == len(expr.items) - 1:
                         self._mark_last()
                     item.accept(self)
+
+    def visit_subscript_expr(self, expr: p.SubscriptExpr) -> None:
+        self._write_line("SubscriptExpr")
+        with self._child_level():
+            self._write_line("object")
+            with self._child_level(single=True):
+                expr.object.accept(self)
+            self._write_line("index", last=True)
+            with self._child_level(single=True):
+                expr.index.accept(self)

midas/ast/python.py

@@ -224,6 +224,9 @@ class Expr(ABC):
         @abstractmethod
         def visit_list_expr(self, expr: ListExpr) -> T: ...
 
+        @abstractmethod
+        def visit_subscript_expr(self, expr: SubscriptExpr) -> T: ...
+
 
 @dataclass(frozen=True)
 class BinaryExpr(Expr):
@@ -324,3 +327,12 @@ class ListExpr(Expr):
 
     def accept(self, visitor: Expr.Visitor[T]) -> T:
         return visitor.visit_list_expr(self)
+
+
+@dataclass(frozen=True)
+class SubscriptExpr(Expr):
+    object: Expr
+    index: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_subscript_expr(self)

midas/checker/python.py

@@ -356,7 +356,7 @@ class PythonTyper(
 
         match operation:
             case Function() as function:
-                if not self._is_binary_function(function):
+                if not self._check_arity(function, 1, 0, 0):
                     self.reporter.error(
                         location,
                         f"Wrong definition of binary operation. Expected function with 1 positional-only parameters, got {function}",
@@ -395,7 +395,7 @@ class PythonTyper(
 
         match operation:
             case Function() as function:
-                if not self._is_unary_function(function):
+                if not self._check_arity(function, 0, 0, 0):
                     self.reporter.error(
                         expr.location,
                         f"Wrong definition of unary operation. Expected function with 0 parameters, got {function}",
@@ -512,6 +512,41 @@ class PythonTyper(
         )
         return self.types.apply_generic(list_type, [UnknownType()])
 
+    def visit_subscript_expr(self, expr: p.SubscriptExpr) -> Type:
+        object: Type = self.type_of(expr.object)
+        operation: Optional[Type] = self.types.lookup_member(object, "__getitem__")
+        if operation is None:
+            self.reporter.error(
+                expr.location,
+                f"Undefined method __getitem__ on {object}",
+            )
+            return UnknownType()
+
+        index: Type = self.type_of(expr.index)
+
+        match operation:
+            case Function() as function:
+                if not self._check_arity(function, 1, 0, 0):
+                    self.reporter.error(
+                        expr.location,
+                        f"Wrong definition of __getitem__. Expected function with 1 positional-only parameters, got {function}",
+                    )
+                    return UnknownType()
+
+                index_arg: Function.Argument = function.pos_args[0]
+                if not self.is_subtype(index, index_arg.type):
+                    self.reporter.error(
+                        expr.location,
+                        f"Wrong index type, expected {index_arg.type}, got {index}",
+                    )
+                    return UnknownType()
+                return function.returns
+            case _:
+                self.reporter.warning(
+                    expr.location, f"Unsupported operation {operation}"
+                )
+                return UnknownType()
+
     def visit_base_type(self, node: p.BaseType) -> Type:
         base: Type
         try:
@@ -654,20 +689,17 @@ class PythonTyper(
 
         return mapped
 
-    def _is_binary_function(self, function: Function) -> bool:
-        if len(function.pos_args) != 1:
+    def _check_arity(
+        self,
+        function: Function,
+        n_pos: Optional[int] = None,
+        n_mixed: Optional[int] = None,
+        n_keyword: Optional[int] = None,
+    ) -> bool:
+        if n_pos is not None and len(function.pos_args) != n_pos:
             return False
-        if len(function.args) != 0:
+        if n_mixed is not None and len(function.args) != n_mixed:
             return False
-        if len(function.kw_args) != 0:
-            return False
-        return True
-
-    def _is_unary_function(self, function: Function) -> bool:
-        if len(function.pos_args) != 0:
-            return False
-        if len(function.args) != 0:
-            return False
-        if len(function.kw_args) != 0:
+        if n_keyword is not None and len(function.kw_args) != n_keyword:
             return False
         return True

midas/checker/resolver.py

@@ -196,3 +196,7 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
     def visit_list_expr(self, expr: p.ListExpr) -> None:
         for item in expr.items:
             self.resolve(item)
+
+    def visit_subscript_expr(self, expr: p.SubscriptExpr) -> None:
+        self.resolve(expr.object)
+        self.resolve(expr.index)

midas/cli/highlighter.py

@@ -218,6 +218,10 @@ class PythonHighlighter(
         for item in expr.items:
             item.accept(self)
 
+    def visit_subscript_expr(self, expr: p.SubscriptExpr) -> None:
+        expr.object.accept(self)
+        expr.index.accept(self)
+
 
 class MidasHighlighter(
     Highlighter, m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[None]

midas/parser/python.py

@@ -23,6 +23,7 @@ from midas.ast.python import (
     MidasType,
     ReturnStmt,
     Stmt,
+    SubscriptExpr,
     TernaryExpr,
     TypeAssign,
     UnaryExpr,
@@ -423,6 +424,13 @@ class PythonParser:
                     items=[self.parse_expr(item) for item in items],
                 )
 
+            case ast.Subscript(value=value, slice=index):
+                return SubscriptExpr(
+                    location=location,
+                    object=self.parse_expr(value),
+                    index=self.parse_expr(index),
+                )
+
             case _:
                 raise UnsupportedSyntaxError(node)
 

tests/serializer/python.py

@@ -22,6 +22,7 @@ from midas.ast.python import (
     MidasType,
     ReturnStmt,
     Stmt,
+    SubscriptExpr,
     TernaryExpr,
     TypeAssign,
     UnaryExpr,
@@ -252,3 +253,10 @@ class PythonAstJsonSerializer(
             "_type": "ListExpr",
             "items": [item.accept(self) for item in expr.items],
         }
+
+    def visit_subscript_expr(self, expr: SubscriptExpr) -> dict:
+        return {
+            "_type": "SubscriptExpr",
+            "object": expr.object.accept(self),
+            "index": expr.index.accept(self),
+        }