diff --git a/gen/python.py b/gen/python.py index 56a5f41..73f64c7 100644 --- a/gen/python.py +++ b/gen/python.py @@ -174,6 +174,10 @@ class SliceExpr: step: Optional[Expr] +class TupleExpr: + items: tuple[Expr, ...] + + class RawExpr: expr: ast.expr diff --git a/midas/ast/printer.py b/midas/ast/printer.py index 82cf4b3..95c4109 100644 --- a/midas/ast/printer.py +++ b/midas/ast/printer.py @@ -879,6 +879,17 @@ class PythonAstPrinter( self._write_optional_child("upper", expr.upper) self._write_optional_child("step", expr.step, last=True) + def visit_tuple_expr(self, expr: p.TupleExpr) -> None: + self._write_line("TupleExpr") + with self._child_level(): + self._write_line("items", last=True) + with self._child_level(): + for i, item in enumerate(expr.items): + self._idx = i + if i == len(expr.items) - 1: + self._mark_last() + item.accept(self) + def visit_raw_expr(self, expr: p.RawExpr) -> None: self._write_line("RawExpr") with self._child_level(single=True): diff --git a/midas/ast/python.py b/midas/ast/python.py index b67a043..6202fcc 100644 --- a/midas/ast/python.py +++ b/midas/ast/python.py @@ -268,6 +268,9 @@ class Expr(ABC): @abstractmethod def visit_slice_expr(self, expr: SliceExpr) -> T: ... + @abstractmethod + def visit_tuple_expr(self, expr: TupleExpr) -> T: ... + @abstractmethod def visit_raw_expr(self, expr: RawExpr) -> T: ... @@ -402,6 +405,14 @@ class SliceExpr(Expr): return visitor.visit_slice_expr(self) +@dataclass(frozen=True) +class TupleExpr(Expr): + items: tuple[Expr, ...] + + def accept(self, visitor: Expr.Visitor[T]) -> T: + return visitor.visit_tuple_expr(self) + + @dataclass(frozen=True) class RawExpr(Expr): expr: ast.expr diff --git a/midas/checker/python.py b/midas/checker/python.py index 26b6301..be3d05f 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -750,6 +750,11 @@ class PythonTyper( def visit_slice_expr(self, expr: p.SliceExpr) -> Type: return self.types.get_type("slice") + def visit_tuple_expr(self, expr: p.TupleExpr) -> Type: + return TupleType( + items=tuple(self.type_of(item) for item in expr.items), + ) + def visit_raw_expr(self, expr: p.RawExpr) -> Type: return UnknownType() diff --git a/midas/checker/resolver.py b/midas/checker/resolver.py index 753f4db..43768bd 100644 --- a/midas/checker/resolver.py +++ b/midas/checker/resolver.py @@ -236,5 +236,9 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]): if expr.step is not None: self.resolve(expr.step) + def visit_tuple_expr(self, expr: p.TupleExpr) -> None: + for item in expr.items: + self.resolve(item) + def visit_raw_expr(self, expr: p.RawExpr) -> None: pass diff --git a/midas/cli/highlighter.py b/midas/cli/highlighter.py index bced867..8d142f1 100644 --- a/midas/cli/highlighter.py +++ b/midas/cli/highlighter.py @@ -247,6 +247,10 @@ class PythonHighlighter( if expr.step is not None: expr.step.accept(self) + def visit_tuple_expr(self, expr: p.TupleExpr) -> None: + for item in expr.items: + item.accept(self) + def visit_raw_expr(self, expr: p.RawExpr) -> None: ... def visit_raw_stmt(self, stmt: p.RawStmt) -> None: ... diff --git a/midas/parser/python.py b/midas/parser/python.py index 9868229..a6c9a78 100644 --- a/midas/parser/python.py +++ b/midas/parser/python.py @@ -30,6 +30,7 @@ from midas.ast.python import ( Stmt, SubscriptExpr, TernaryExpr, + TupleExpr, TypeAssign, UnaryExpr, VariableExpr, @@ -479,6 +480,12 @@ class PythonParser: step=self.parse_expr(step) if step is not None else None, ) + case ast.Tuple(elts=items): + return TupleExpr( + location=location, + items=tuple(self.parse_expr(item) for item in items), + ) + case _: print(f"Unsupported expression: {ast.unparse(node)}") return RawExpr(location=location, expr=node) diff --git a/tests/serializer/python.py b/tests/serializer/python.py index 57e47c9..7c29616 100644 --- a/tests/serializer/python.py +++ b/tests/serializer/python.py @@ -302,6 +302,12 @@ class PythonAstJsonSerializer( "step": self._serialize_optional(expr.step), } + def visit_tuple_expr(self, expr: TupleExpr) -> dict: + return { + "_type": "TupleExpr", + "items": [item.accept(self) for item in expr.items], + } + def visit_raw_expr(self, expr: RawExpr) -> dict: return { "_type": "RawExpr",