feat(checker): process frame type definitions

Reviewed-on: #22

examples/02_demonstration/demo.py

@@ -1,6 +1,8 @@
-from typing import TypeVar, cast
+from typing import TypeVar
 
-from demo_stubs import CHF, EUR, USD, Currency, Price, Discount
+from demo_stubs import CHF, EUR, USD, Currency, Discount, Price
+
+from midas.typing import cast, unsafe_cast
 
 T = TypeVar("T", bound=Currency)
 
@@ -28,3 +30,6 @@ discounted = apply_discount(
 )
 
 print(f"Discounted: CHF {discounted}")
+
+large_data = [i * 10 for i in range(100)]
+prices = unsafe_cast(list[Price[EUR]], large_data)

gen/midas.py

@@ -152,4 +152,14 @@ class FunctionType:
         required: bool
 
 
+class FrameType:
+    columns: list[Column]
+
+    @dataclass(frozen=True, kw_only=True)
+    class Column:
+        location: Optional[Location] = None
+        name: Token
+        type: Type
+
+
 ###<

gen/python.py

@@ -145,6 +145,7 @@ class LogicalExpr:
 class CastExpr:
     type: MidasType
     expr: Expr
+    unsafe: bool
 
 
 class TernaryExpr:

midas/ast/midas.py

@@ -253,6 +253,9 @@ class Type(ABC):
         @abstractmethod
         def visit_function_type(self, type: FunctionType) -> T: ...
 
+        @abstractmethod
+        def visit_frame_type(self, type: FrameType) -> T: ...
+
 
 @dataclass(frozen=True)
 class NamedType(Type):
@@ -311,3 +314,17 @@ class FunctionType(Type):
 
     def accept(self, visitor: Type.Visitor[T]) -> T:
         return visitor.visit_function_type(self)
+
+
+@dataclass(frozen=True)
+class FrameType(Type):
+    columns: list[Column]
+
+    @dataclass(frozen=True, kw_only=True)
+    class Column:
+        location: Optional[Location] = None
+        name: Token
+        type: Type
+
+    def accept(self, visitor: Type.Visitor[T]) -> T:
+        return visitor.visit_frame_type(self)

midas/ast/printer.py

@@ -350,6 +350,25 @@ class MidasAstPrinter(
                 arg.type.accept(self)
             self._write_line(f"required: {arg.required}", last=True)
 
+    def visit_frame_type(self, type: m.FrameType) -> None:
+        self._write_line("FrameType")
+        with self._child_level(single=True):
+            self._write_line("columns")
+            with self._child_level():
+                for i, column in enumerate(type.columns):
+                    self._idx = i
+                    if i == len(type.columns) - 1:
+                        self._mark_last()
+                    self._print_frame_column(column)
+
+    def _print_frame_column(self, column: m.FrameType.Column) -> None:
+        self._write_line("Column")
+        with self._child_level():
+            self._write_line(f'name: "{column.name.lexeme}"')
+            self._write_line("type")
+            with self._child_level(single=True):
+                column.type.accept(self)
+
 
 class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str]):
     def __init__(self, indent: int = 4):
@@ -502,6 +521,23 @@ class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str]
             res += "?"
         return res
 
+    def visit_frame_type(self, type: m.FrameType) -> str:
+        res: str = self.indented("Frame[")
+        if len(type.columns) != 0:
+            res += "\n"
+            self.level += 1
+            columns: list[str] = []
+            for column in type.columns:
+                columns.append(self.indented(self._print_frame_column(column)))
+            res += ",\n".join(columns)
+            self.level -= 1
+            res += "\n"
+        res += "]"
+        return res
+
+    def _print_frame_column(self, column: m.FrameType.Column) -> str:
+        return f"{column.name.lexeme}: {column.type.accept(self)}"
+
 
 class PythonAstPrinter(
     AstPrinter,
@@ -757,9 +793,10 @@ class PythonAstPrinter(
             self._write_line("type")
             with self._child_level(single=True):
                 expr.type.accept(self)
-            self._write_line("expr", last=True)
+            self._write_line("expr")
             with self._child_level(single=True):
                 expr.expr.accept(self)
+            self._write_line(f"unsafe: {expr.unsafe}", last=True)
 
     def visit_ternary_expr(self, expr: p.TernaryExpr) -> None:
         self._write_line("TernaryExpr")

midas/ast/python.py

@@ -350,6 +350,7 @@ class LogicalExpr(Expr):
 class CastExpr(Expr):
     type: MidasType
     expr: Expr
+    unsafe: bool
 
     def accept(self, visitor: Expr.Visitor[T]) -> T:
         return visitor.visit_cast_expr(self)

midas/checker/builtins.py

@@ -15,7 +15,7 @@ if TYPE_CHECKING:
 
 
 BUILTIN_SUBTYPES: dict[str, set[str]] = {
-    "object": {"float", "list", "dict"},
+    "object": {"float", "list", "dict", "str"},
     "float": {"int"},
     "int": {"bool"},
 }

midas/checker/diagnostic.py

@@ -9,6 +9,7 @@ class DiagnosticType(StrEnum):
     ERROR = "Error"
     WARNING = "Warning"
     INFO = "Info"
+    DEBUG = "Debug"
 
 
 @dataclass(frozen=True)

midas/checker/evaluator.py

@@ -0,0 +1,172 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Optional
+
+import midas.ast.midas as m
+from midas.checker.preamble import Preamble
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter
+from midas.checker.types import Function, Predicate
+from midas.lexer.token import TokenType
+
+
+@dataclass(frozen=True, kw_only=True)
+class PartialPredicate(Predicate):
+    scope: dict[str, Any]
+
+
+class Evaluator(m.Expr.Visitor[Any]):
+    def __init__(self, types: TypesRegistry, reporter: Optional[FileReporter] = None):
+        self.types: TypesRegistry = types
+        self.reporter: Optional[FileReporter] = reporter
+        self.preamble: Preamble = Preamble(self.types)
+        self.scopes: list[dict[str, Any]] = [{}]
+
+    def evaluate(self, expr: m.Expr) -> Any:
+        value: Any = expr.accept(self)
+        if self.reporter is not None:
+            self.reporter.debug(expr.location, f"Value: {value}")
+        return value
+
+    def get_value(self, name: str) -> Any:
+        scope: dict[str, Any] = self.scopes[-1]
+        return scope[name]
+
+    def set_value(self, name: str, value: Any, force_declare: bool = False):
+        if not force_declare:
+            for scope in reversed(self.scopes):
+                if name in scope:
+                    scope[name] = value
+                    return
+        self.scopes[-1][name] = value
+
+    def visit_logical_expr(self, expr: m.LogicalExpr) -> Any:
+        def left():
+            return self.evaluate(expr.left)
+
+        def right():
+            return self.evaluate(expr.right)
+
+        match expr.operator.type:
+            case TokenType.AND:
+                return left() and right()
+            case _:
+                raise NotImplementedError
+
+    def visit_binary_expr(self, expr: m.BinaryExpr) -> Any:
+        left: Any = self.evaluate(expr.left)
+        right: Any = self.evaluate(expr.right)
+        match expr.operator.type:
+            case TokenType.MINUS:
+                return left - right
+            case TokenType.STAR:
+                return left * right
+            case TokenType.SLASH:
+                return left / right
+            case TokenType.GREATER:
+                return left > right
+            case TokenType.GREATER_EQUAL:
+                return left >= right
+            case TokenType.LESS:
+                return left < right
+            case TokenType.LESS_EQUAL:
+                return left <= right
+            case TokenType.EQUAL_EQUAL:
+                return left == right
+            case TokenType.BANG_EQUAL:
+                return left != right
+            case _:
+                raise NotImplementedError
+
+    def visit_unary_expr(self, expr: m.UnaryExpr) -> Any:
+        right: Any = self.evaluate(expr.right)
+        match expr.operator.type:
+            case TokenType.MINUS:
+                return -right
+            case _:
+                raise NotImplementedError
+
+    def visit_call_expr(self, expr: m.CallExpr) -> Any:
+        callee: Any = self.evaluate(expr.callee)
+        args: list[Any] = [self.evaluate(arg) for arg in expr.arguments]
+        kwargs: dict[str, Any] = {
+            name: self.evaluate(arg) for name, arg in expr.keywords.items()
+        }
+
+        match callee:
+            case Predicate():
+                return self._evaluate_predicate(callee, args, kwargs)
+            case _ if callable(callee):
+                return callee(*args, **kwargs)
+            case _:
+                return NotImplementedError
+
+    def visit_get_expr(self, expr: m.GetExpr) -> Any:
+        obj: Any = self.evaluate(expr.expr)
+        return getattr(obj, expr.name.lexeme)
+
+    def visit_variable_expr(self, expr: m.VariableExpr) -> Any:
+        name: str = expr.name.lexeme
+        for scope in reversed(self.scopes):
+            if name in scope:
+                return scope[name]
+
+        predicate: Optional[Predicate] = self.types.lookup_predicate(name)
+        if predicate is not None:
+            if predicate.alias:
+                return self.evaluate(predicate.body)
+            return predicate
+
+        glob: Optional[Callable] = self.preamble.get_py_func(name)
+        if glob is not None:
+            return glob
+        raise NameError(f"Unknown variable '{name}'")
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr) -> Any:
+        return self.evaluate(expr.expr)
+
+    def visit_literal_expr(self, expr: m.LiteralExpr) -> Any:
+        return expr.value
+
+    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> Any:
+        return self.get_value("_")
+
+    def _evaluate_predicate(
+        self, predicate: Predicate, args: list[Any], kwargs: dict[str, Any]
+    ) -> Any:
+        res: Any = None
+        if isinstance(predicate, PartialPredicate):
+            self.scopes.append(predicate.scope)
+        else:
+            self.scopes.append({})
+        match predicate.type:
+            case Function(returns=Function() as inner):
+                self._map_args(predicate.type, args, kwargs)
+                res = PartialPredicate(
+                    type=inner,
+                    body=predicate.body,
+                    alias=False,
+                    scope=self.scopes[-1],
+                )
+
+            case Function():
+                self._map_args(predicate.type, args, kwargs)
+                res = self.evaluate(predicate.body)
+
+            case _:
+                raise NotImplementedError
+        self.scopes.pop()
+        return res
+
+    def _map_args(self, function: Function, args: list[Any], kwargs: dict[str, Any]):
+        positional: list[Function.Argument] = function.pos_args + function.args
+        keywords: dict[str, Function.Argument] = {
+            arg.name: arg for arg in function.args + function.kw_args
+        }
+
+        for i, arg in enumerate(args):
+            param: Function.Argument = positional[i]
+            self.set_value(param.name, arg)
+
+        for name, arg in kwargs.items():
+            param: Function.Argument = keywords[name]
+            self.set_value(param.name, arg)

midas/checker/frames.py

@@ -0,0 +1,133 @@
+from typing import Optional, TypeGuard, cast
+
+from midas.ast.location import Location
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter
+from midas.checker.types import ColumnType, DataFrameType, TupleType, Type, UnknownType
+
+import midas.ast.python as p
+
+
+def is_list_of_literals(exprs: list[p.Expr]) -> TypeGuard[list[p.LiteralExpr]]:
+    return all(isinstance(expr, p.LiteralExpr) for expr in exprs)
+
+
+class FrameManager:
+    def __init__(self, types: TypesRegistry) -> None:
+        self.types: TypesRegistry = types
+
+    def assign(
+        self,
+        reporter: FileReporter,
+        location: Location,
+        frame: DataFrameType,
+        index: p.Expr,
+        value_type: Type,
+    ) -> Type:
+        match index:
+            case p.LiteralExpr(value=str() as name):
+                return self.assign_column(reporter, location, frame, name, value_type)
+
+            case p.ListExpr(items=indices) if is_list_of_literals(indices) and all(
+                isinstance(idx, str) for idx in indices
+            ):
+                raise NotImplementedError
+
+            case _:
+                reporter.error(location, f"Invalid index type {index} on {frame}")
+                return UnknownType()
+
+    def assign_column(
+        self,
+        reporter: FileReporter,
+        location: Location,
+        frame: DataFrameType,
+        name: str,
+        type: Type,
+    ) -> Type:
+        if not isinstance(type, ColumnType):
+            reporter.error(
+                location,
+                f"Cannot assign {type} to dataframe column. Must be a ColumnType",
+            )
+            return frame
+        return self._set_column(frame, name, type)
+
+    def get(
+        self,
+        reporter: FileReporter,
+        location: Location,
+        frame: DataFrameType,
+        index: p.Expr,
+    ) -> Type:
+        match index:
+            case p.LiteralExpr(value=str() as name):
+                column: Optional[ColumnType] = FrameManager._get_column(frame, name)
+                if column is None:
+                    reporter.error(location, f"Unknown column '{name}' on {frame}")
+                    return UnknownType()
+                return column
+
+            case p.ListExpr(items=indices) if is_list_of_literals(indices) and all(
+                isinstance(index.value, str) for index in indices
+            ):
+                names: list[str] = [cast(str, index.value) for index in indices]
+                columns: list[ColumnType] = []
+                for name in names:
+                    column: Optional[ColumnType] = FrameManager._get_column(frame, name)
+                    if column is None:
+                        reporter.error(location, f"Unknown column '{name}' on {frame}")
+                        return UnknownType()
+                    columns.append(column)
+                return TupleType(items=tuple(columns))
+
+            case _:
+                reporter.error(location, f"Invalid index type {index} on {frame}")
+                return UnknownType()
+
+    @classmethod
+    def _set_column(
+        cls, frame: DataFrameType, name: str, column: ColumnType
+    ) -> DataFrameType:
+        new_columns: list[DataFrameType.Column] = []
+        index: int = len(frame.columns)
+        replace: bool = False
+        for i, col in enumerate(frame.columns):
+            if col.name == name:
+                index = i
+                replace = True
+                # TODO: check column type here to prevent changing it
+            new_columns.append(col)
+
+        new_col: DataFrameType.Column = DataFrameType.Column(
+            index=index,
+            name=name,
+            type=column,
+        )
+        if replace:
+            new_columns[index] = new_col
+        else:
+            new_columns.append(new_col)
+
+        return DataFrameType(columns=new_columns)
+
+    @classmethod
+    def _set_columns(
+        cls, frame: DataFrameType, names: list[str], columns: list[ColumnType]
+    ) -> DataFrameType:
+        for name, col in zip(names, columns):
+            frame = cls._set_column(frame, name, col)
+        return frame
+
+    @classmethod
+    def _get_column(cls, frame: DataFrameType, name: str) -> Optional[ColumnType]:
+        for col in frame.columns:
+            if col.name == name:
+                return col.type
+        return None
+
+    @classmethod
+    def _get_columns(
+        cls, frame: DataFrameType, names: list[str]
+    ) -> list[Optional[ColumnType]]:
+        return [cls._get_column(frame, name) for name in names]

midas/checker/midas.py

@@ -14,8 +14,10 @@ from midas.checker.reporter import FileReporter, Reporter
 from midas.checker.types import (
     AliasType,
     AppliedType,
+    ColumnType,
     ComplexType,
     ConstraintType,
+    DataFrameType,
     ExtensionType,
     Function,
     GenericType,
@@ -401,6 +403,18 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
             kw=[process_arg(arg, i + n_pos + n_mixed) for i, arg in enumerate(spec.kw)],
         )
 
+    def visit_frame_type(self, type: m.FrameType) -> Type:
+        def process_column(i: int, col: m.FrameType.Column) -> DataFrameType.Column:
+            return DataFrameType.Column(
+                index=i,
+                name=col.name.lexeme,
+                type=ColumnType(type=col.type.accept(self)),
+            )
+
+        return DataFrameType(
+            columns=[process_column(i, col) for i, col in enumerate(type.columns)]
+        )
+
     def _resolve_type_params(self, params: list[m.TypeParam]):
         vars: list[TypeVar] = []
         for param in params:

midas/checker/preamble.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from typing import Callable, Optional
 
 from midas.checker.environment import Environment
 from midas.checker.registry import TypesRegistry
@@ -16,16 +17,18 @@ class Preamble(Environment):
     def __init__(self, types: TypesRegistry) -> None:
         super().__init__()
         self._types: TypesRegistry = types
+        self._python_funcs: dict[str, Callable] = {}
 
-        self._def_type_constructor("object")
+        self._def_type_constructor("object", object)
-        self._def_type_constructor("float")
+        self._def_type_constructor("float", float)
-        self._def_type_constructor("int")
+        self._def_type_constructor("int", int)
-        self._def_type_constructor("bool")
+        self._def_type_constructor("bool", bool)
-        self._def_type_constructor("str")
+        self._def_type_constructor("str", str)
         self._def_function(
             name="list",
             pos=[Param("object", TopType())],
             returns=self._list_of(TopType()),
+            py_function=list,
         )
 
         # TODO: use sink
@@ -33,6 +36,7 @@ class Preamble(Environment):
             name="print",
             pos=[Param("object", TopType())],
             returns=UnitType(),
+            py_function=print,
         )
 
         map_in = TypeVar(name="T", bound=None)
@@ -53,6 +57,7 @@ class Preamble(Environment):
             ],
             returns=self._list_of(map_out),  # TODO: replace with Iterable[U]
             type_vars=[map_in, map_out],
+            py_function=map,
         )
         self._def_function(
             name="input",
@@ -63,12 +68,13 @@ class Preamble(Environment):
     def _list_of(self, item_type: Type) -> Type:
         return self._types.apply_generic(self._types.get_type("list"), [item_type])
 
-    def _def_type_constructor(self, name: str):
+    def _def_type_constructor(self, name: str, py_function: Optional[Callable] = None):
         # TODO: more specific arg types
         self._def_function(
             name=name,
             pos=[Param("object", TopType(), required=False)],
             returns=self._types.get_type(name),
+            py_function=py_function,
         )
 
     def _make_function(
@@ -115,6 +121,7 @@ class Preamble(Environment):
         kw: list[Param] = [],
         returns: Type = UnitType(),
         type_vars: list[TypeVar] = [],
+        py_function: Optional[Callable] = None,
     ):
         function: Type = self._make_function(
             name=name,
@@ -125,3 +132,8 @@ class Preamble(Environment):
             type_vars=type_vars,
         )
         self.define(name, function)
+        if py_function is not None:
+            self._python_funcs[name] = py_function
+
+    def get_py_func(self, name: str) -> Optional[Callable]:
+        return self._python_funcs.get(name)

midas/checker/python.py

@@ -1,11 +1,14 @@
 import ast
 import logging
 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional
 
 import midas.ast.python as p
 from midas.ast.location import Location
+from midas.ast.printer import MidasPrinter
 from midas.checker.environment import Environment
+from midas.checker.evaluator import Evaluator
+from midas.checker.frames import FrameManager
 from midas.checker.operators import (
     PY_COMPARATOR_METHODS,
     PY_OPERATOR_METHODS,
@@ -18,9 +21,14 @@ from midas.checker.resolver import Resolver
 from midas.checker.types import (
     AliasType,
     AppliedType,
+    BaseType,
+    ColumnType,
+    ConstraintType,
+    DataFrameType,
     Function,
     GenericType,
     OverloadedFunction,
+    TupleType,
     Type,
     TypeVar,
     UnitType,
@@ -67,10 +75,12 @@ class PythonTyper(
         self.logger: logging.Logger = logging.getLogger("PythonTyper")
         self.reporter: FileReporter = reporter.for_file(None)
         self.types: TypesRegistry = types
+        self.frame_mgr: FrameManager = FrameManager(self.types)
         self.global_env: Environment = Preamble(self.types)
         self.env: Environment = self.global_env
         self.locals: dict[p.Expr, int] = {}
         self.judgements: list[tuple[p.Expr, Type]] = []
+        self.evaluated_casts: list[p.CastExpr] = []
 
     def process(self, source: str, path: Optional[str]) -> TypedAST:
         self.reporter = self.reporter.for_file(path)
@@ -84,10 +94,15 @@ class PythonTyper(
         self.env = self.global_env
         self.locals = resolver.locals
         self.judgements = []
+        self.evaluated_casts = []
 
         self.check(stmts)
 
-        return TypedAST(stmts=stmts, judgements=self.judgements)
+        return TypedAST(
+            stmts=stmts,
+            judgements=self.judgements,
+            evaluated_casts=self.evaluated_casts,
+        )
 
     def judge(self, expr: p.Expr, type: Type):
         """Record a typing judgement
@@ -309,9 +324,15 @@ class PythonTyper(
             case p.VariableExpr():
                 self._assign_var(location, target, value_type)
 
+            # Allow any kind of object because we disallow creating new attributes
             case p.GetExpr(object=object, name=name):
                 self._assign_attr(location, object, name, value_type)
 
+            # Only support variable expressions because modifying
+            # the underlying value would require reference types
+            case p.SubscriptExpr(object=p.VariableExpr() as var, index=index):
+                self._assign_sub(location, var, index, value_type)
+
             case _:
                 if not isinstance(target, p.VariableExpr):
                     self.logger.warning(f"Unsupported assignment to {target}")
@@ -350,6 +371,27 @@ class PythonTyper(
                 f"Cannot assign {value_type} to member '{object_type}.{name}' of type {member}",
             )
 
+    def _assign_sub(
+        self,
+        location: Location,
+        var: p.VariableExpr,
+        index: p.Expr,
+        value_type: Type,
+    ):
+        var_type: Type = self.type_of(var)
+        # TODO: what happens if type is an alias of a dataframe type
+        match var_type:
+            case DataFrameType() as frame:
+                new_type: Type = self.frame_mgr.assign(
+                    self.reporter, location, frame, index, value_type
+                )
+                self.env.assign(var.name, new_type)
+            case _:
+                self.reporter.error(
+                    location,
+                    f"Cannot assign {value_type} to index {index} of {var_type}",
+                )
+
     def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
         type: Type = self.type_of(stmt.value) if stmt.value is not None else UnitType()
         self.env.return_types.append(type)
@@ -538,8 +580,16 @@ class PythonTyper(
         return UnknownType()
 
     def visit_cast_expr(self, expr: p.CastExpr) -> Type:
-        _ = self.type_of(expr.expr)
+        subject_type: Type = self.type_of(expr.expr)
-        return self.resolve_type_expr(expr.type)
+        target_type: Type = self.resolve_type_expr(expr.type)
+        is_lit, lit_value = self._get_literal(expr.expr)
+        if is_lit:
+            evaluated: bool = self._evaluate_cast_statically(
+                expr, subject_type, target_type, lit_value
+            )
+            if evaluated:
+                self.evaluated_casts.append(expr)
+        return target_type
 
     def visit_ternary_expr(self, expr: p.TernaryExpr) -> Type:
         test_type: Type = self.type_of(expr.test)
@@ -622,6 +672,13 @@ class PythonTyper(
 
     def visit_subscript_expr(self, expr: p.SubscriptExpr) -> Type:
         object: Type = self.type_of(expr.object)
+        unfolded: Type = unfold_type(object)
+        match unfolded:
+            case TupleType():
+                return self._visit_tuple_subscript(unfolded, expr)
+            case DataFrameType():
+                return self._visit_frame_subscript(unfolded, expr)
+
         operation: Optional[Type] = self.types.lookup_member(object, "__getitem__")
         if operation is None:
             self.reporter.error(
@@ -659,13 +716,26 @@ class PythonTyper(
         self.reporter.warning(node.location, "ConstraintType not yet supported")
         return UnknownType()
 
-    def visit_frame_column(self, node: p.FrameColumn) -> Type:
+    def visit_frame_column(self, node: p.FrameColumn) -> ColumnType:
-        self.reporter.warning(node.location, "FrameColumn not yet supported")
+        return ColumnType(
-        return UnknownType()
+            type=(
+                self.resolve_type_expr(node.type)
+                if node.type is not None
+                else UnknownType()
+            )
+        )
 
     def visit_frame_type(self, node: p.FrameType) -> Type:
-        self.reporter.warning(node.location, "FrameType not yet supported")
+        return DataFrameType(
-        return UnknownType()
+            columns=[
+                DataFrameType.Column(
+                    index=i,
+                    name=column.name,
+                    type=self.visit_frame_column(column),
+                )
+                for i, column in enumerate(node.columns)
+            ]
+        )
 
     def _get_call_result(
         self,
@@ -1108,3 +1178,113 @@ class PythonTyper(
                 return p.BaseType(location=location, base=name, param=None)
             case _:
                 raise NotImplementedError
+
+    def _get_literal(self, expr: p.Expr) -> tuple[bool, Any]:
+        match expr:
+            case p.LiteralExpr(value=value):
+                return True, value
+
+            case p.ListExpr(items=items):
+                values: list[Any] = []
+                for item in items:
+                    is_lit, value = self._get_literal(item)
+                    if not is_lit:
+                        return False, None
+                    values.append(value)
+                return True, values
+
+            case p.DictExpr(keys=keys, values=values):
+                pairs: list[tuple[Any, Any]] = []
+                for key, value in zip(keys, values):
+                    key_val = None
+                    if key is not None:
+                        is_lit, key_val = self._get_literal(key)
+                        if not is_lit:
+                            return False, None
+
+                    is_lit, value_val = self._get_literal(value)
+                    if not is_lit:
+                        return False, None
+
+                    if key is None:
+                        # TODO: check that value is always a dict
+                        assert isinstance(value_val, dict)
+                        pairs.extend(value_val.items())
+                    else:
+                        pairs.append((key_val, value_val))
+                return True, dict(pairs)
+
+            case _:
+                return False, None
+
+    def _evaluate_cast_statically(
+        self, expr: p.CastExpr, subject_type: Type, target_type: Type, lit_value: Any
+    ) -> bool:
+        match target_type:
+            case AliasType(type=base):
+                return self._evaluate_cast_statically(
+                    expr, subject_type, base, lit_value
+                )
+
+            case AppliedType(body=body):
+                return self._evaluate_cast_statically(
+                    expr, subject_type, body, lit_value
+                )
+
+            case ConstraintType(type=base, constraint=constraint):
+                evaluated: bool = True
+                if not self._evaluate_cast_statically(
+                    expr, subject_type, base, lit_value
+                ):
+                    evaluated = False
+
+                evaluator = Evaluator(self.types)
+                evaluator.set_value("_", lit_value)
+                res = evaluator.evaluate(constraint)
+                if not res:
+                    printer = MidasPrinter()
+                    constraint_str: str = printer.print(constraint)
+                    self.reporter.error(
+                        expr.location,
+                        f"Value {lit_value!r} does not fit constraint '{constraint_str}'",
+                    )
+                    evaluated = False
+                return evaluated
+
+            case BaseType():
+                # TODO: do we want to allow cast(float, int)? would require runtime conversion
+                if not self.types.is_subtype(
+                    subject_type, target_type
+                ) or not self.types.is_subtype(target_type, subject_type):
+                    self.reporter.error(
+                        expr.location,
+                        f"Value {lit_value!r} of type {subject_type} cannot be cast as {target_type}",
+                    )
+                    return False
+                return True
+
+            case _:
+                self.reporter.info(
+                    expr.location, f"Cannot evaluate cast to {target_type} statically"
+                )
+                return False
+
+    def _visit_tuple_subscript(self, tup: TupleType, expr: p.SubscriptExpr) -> Type:
+        match expr.index:
+            case p.LiteralExpr(value=int() as index):
+                if index < 0 or index >= len(tup.items):
+                    self.reporter.error(
+                        expr.location, f"Index {index} out of range for tuple {tup}"
+                    )
+                    return UnknownType()
+                return tup.items[index]
+            case _:
+                self.reporter.error(
+                    expr.location, f"Invalid index type {expr.index} on {tup}"
+                )
+                return UnknownType()
+
+    def _visit_frame_subscript(
+        self, frame: DataFrameType, expr: p.SubscriptExpr
+    ) -> Type:
+        return self.frame_mgr.get(self.reporter, expr.location, frame, expr.index)

midas/checker/reporter.py

@@ -61,3 +61,10 @@ class FileReporter:
             location=location,
             message=message,
         )
+
+    def debug(self, location: Location, message: str):
+        self.report(
+            type=DiagnosticType.DEBUG,
+            location=location,
+            message=message,
+        )

midas/checker/resolver.py

@@ -128,6 +128,10 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
 
             case p.GetExpr():
                 target.accept(self)
+
+            case p.SubscriptExpr():
+                target.accept(self)
+
             case _:
                 raise Exception(f"Unsupported assignment to {target}")
 

midas/checker/types.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass, field
 from enum import StrEnum
-from typing import Optional, assert_never
+from typing import Optional, assert_never, cast
 
 import midas.ast.midas as m
 from midas.ast.printer import MidasPrinter
@@ -156,6 +156,37 @@ class ConstraintType:
         return f"{self.type} where {printer.print(self.constraint)}"
 
 
+@dataclass(frozen=True, kw_only=True)
+class TupleType:
+    items: tuple[Type, ...]
+
+    def __str__(self) -> str:
+        return f"({', '.join(map(str, self.items))})"
+
+
+@dataclass(frozen=True, kw_only=True)
+class ColumnType:
+    type: Type
+
+    def __str__(self) -> str:
+        return f"Column[{self.type}]"
+
+
+@dataclass(frozen=True, kw_only=True)
+class DataFrameType:
+    columns: list[Column]
+
+    def __str__(self) -> str:
+        schema: list[str] = [f"{col.name}: {col.type}" for col in self.columns]
+        return f"Frame[{', '.join(schema)}]"
+
+    @dataclass(frozen=True, kw_only=True)
+    class Column:
+        index: int
+        name: Optional[str]
+        type: ColumnType
+
+
 def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
     def sub_argument(arg: Function.Argument):
         return Function.Argument(
@@ -165,6 +196,13 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
             required=arg.required,
         )
 
+    def sub_column(col: DataFrameType.Column):
+        return DataFrameType.Column(
+            index=col.index,
+            name=col.name,
+            type=cast(ColumnType, substitute_typevars(col.type, substitutions)),
+        )
+
     match type:
         case TopType():
             return type
@@ -250,10 +288,26 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
                 body=substitute_typevars(body, substitutions),
             )
 
+        case TupleType(items=items):
+            return TupleType(
+                items=tuple(substitute_typevars(item, substitutions) for item in items),
+            )
+
+        case ColumnType(type=items_type):
+            return ColumnType(
+                type=substitute_typevars(items_type, substitutions),
+            )
+
+        case DataFrameType(columns=columns):
+            return DataFrameType(
+                columns=list(map(sub_column, columns)),
+            )
+
         case UnknownType() | UnitType():
             return type
 
         case TopType() | GenericType():
+
             raise NotImplementedError(f"Unsupported type {type}")
 
         # Ensure exhaustiveness
@@ -317,6 +371,15 @@ def to_annotation(type: Type) -> str:
         case ConstraintType():
             return str(type)
 
+        case TupleType(items=items):
+            return f"Tuple[{', '.join(map(to_annotation, items))}]"
+
+        case ColumnType():
+            return "pd.Series"
+
+        case DataFrameType():
+            return "pd.DataFrame"
+
         case _:
             assert_never(type)
 
@@ -342,4 +405,7 @@ Type = (
     | GenericType
     | AppliedType
     | ConstraintType
+    | TupleType
+    | ColumnType
+    | DataFrameType
 )

midas/cli/commands/compile.py

@@ -19,9 +19,11 @@ from midas.utils import TypedAST
 @click.command(help="Compile source")
 @click.argument("file", type=click.File("r"))
 @click.option("-t", "--types", type=click.File("r"), multiple=True)
+@click.option("--ignore-errors", is_flag=True)
 def compile(
     file: TextIO,
     types: tuple[TextIO],
+    ignore_errors: bool,
 ):
     source: str = file.read()
     source_path: Path = Path(file.name).resolve()
@@ -35,7 +37,9 @@ def compile(
     printer = DiagnosticPrinter()
     printer.print_all(diagnostics)
 
-    if any(map(lambda d: d.type == DiagnosticType.ERROR, diagnostics)):
+    if not ignore_errors and any(
+        map(lambda d: d.type == DiagnosticType.ERROR, diagnostics)
+    ):
         sys.exit(1)
 
     generator = Generator(workdir=source_path.parent, types=checker.types)

midas/cli/commands/types.py

@@ -41,6 +41,7 @@ def types(
                 message=f"Type: {type}",
             )
         )
+    diagnostics.extend(checker.diagnostics)
     printer = DiagnosticPrinter()
     printer.print_all(diagnostics)
 

midas/cli/highlighter.py

@@ -350,6 +350,14 @@ class MidasHighlighter(
         for param in spec.pos + spec.mixed + spec.kw:
             param.type.accept(self)
 
+    def visit_frame_type(self, type: m.FrameType) -> None:
+        self.wrap(type, "frame")
+        for column in type.columns:
+            self._visit_frame_column(column)
+
+    def _visit_frame_column(self, column: m.FrameType.Column) -> None:
+        self.wrap(column, "column")
+
 
 class DiagnosticsHighlighter(Highlighter):
     EXTRA_CSS_PATH: Optional[Path] = Path(__file__).parent / "hl_diagnostic.css"

midas/cli/utils.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from pathlib import Path
 from typing import Optional
 
@@ -7,6 +8,13 @@ from midas.cli.ansi import Ansi
 
 
 class DiagnosticPrinter:
+    COLORS: dict[DiagnosticType, int] = {
+        DiagnosticType.ERROR: Ansi.RED,
+        DiagnosticType.WARNING: Ansi.YELLOW,
+        DiagnosticType.INFO: Ansi.CYAN,
+        DiagnosticType.DEBUG: Ansi.MAGENTA,
+    }
+
     def __init__(self) -> None:
         self.files: dict[Optional[str], list[str]] = {}
 
@@ -22,10 +30,25 @@ class DiagnosticPrinter:
         return self.files[filename]
 
     def print_all(self, diagnostics: list[Diagnostic], indent: int = 4):
+        by_type: dict[DiagnosticType, int] = defaultdict(int)
         for diagnostic in diagnostics:
             filename: Optional[str] = diagnostic.file_path
             lines = self.get_lines(filename)
             self.print(lines, diagnostic, indent=indent)
+            by_type[diagnostic.type] += 1
+
+        if len(diagnostics) == 0:
+            return
+
+        counts: list[str] = []
+        for type in DiagnosticType:
+            if type not in by_type:
+                continue
+            count: int = by_type[type]
+            color: int = self.COLORS.get(type, Ansi.WHITE)
+            counts.append(f"{Ansi.FG(color)}{type.value}s{Ansi.RESET}: {count}")
+
+        print("    ".join(counts))
 
     def print(self, lines: list[str], diagnostic: Diagnostic, indent: int = 4):
         """Pretty-print a diagnostic, showing some context if possible
@@ -55,11 +78,7 @@ class DiagnosticPrinter:
         before: str = line[:start_offset]
         after: str = line[end_offset:]
 
-        color: int = {
+        color: int = self.COLORS.get(diagnostic.type, Ansi.WHITE)
-            DiagnosticType.ERROR: Ansi.RED,
-            DiagnosticType.WARNING: Ansi.YELLOW,
-            DiagnosticType.INFO: Ansi.CYAN,
-        }.get(diagnostic.type, Ansi.WHITE)
 
         subject: str = Ansi.FG(color) + line[start_offset:end_offset] + Ansi.RESET
         cursor: str = (

midas/generator/generator.py

@@ -1,4 +1,5 @@
 import ast
+import logging
 import shutil
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -13,13 +14,16 @@ from midas.checker.types import (
     AliasType,
     AppliedType,
     BaseType,
+    ColumnType,
     ComplexType,
     ConstraintType,
+    DataFrameType,
     ExtensionType,
     Function,
     GenericType,
     OverloadedFunction,
     TopType,
+    TupleType,
     Type,
     TypeVar,
     UnitType,
@@ -40,10 +44,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         self.workdir: Path = workdir.resolve()
         self.build_dir: Path = self.workdir / "build" / "midas"
         self.rel_src_path: Path = Path()
+        self.logger: logging.Logger = logging.getLogger("Generator")
 
         self._typed_ast: TypedAST = TypedAST(
             stmts=[],
             judgements=[],
+            evaluated_casts=[],
         )
         self._alias_count: int = 0
         self._predicate_count: int = 0
@@ -131,6 +137,10 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
 
     def visit_cast_expr(self, expr: p.CastExpr) -> ast.expr:
         expr2: ast.expr = expr.expr.accept(self)
+
+        if expr in self._typed_ast.evaluated_casts or expr.unsafe:
+            return expr2
+
         alias: ast.expr = self._make_alias(expr2)
 
         type: Type = self._get_expr_type(expr)
@@ -327,6 +337,19 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 if bound is not None:
                     self._make_cast_asserts(src_location, expr, bound)
 
+            case TupleType(items=items):
+                self._add_assert(
+                    ast.Call(
+                        func=ast.Name(id="isinstance"),
+                        args=[expr, ast.Name(id="tuple")],
+                        keywords=[],
+                    ),
+                    self._make_cast_assert_message(src_location, expr, type),
+                )
+                assert isinstance(expr, ast.Tuple)
+                for item, item_type in zip(expr.elts, items):
+                    self._make_cast_asserts(src_location, item, item_type)
+
             case (
                 TopType()
                 | Function()
@@ -334,8 +357,10 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 | ComplexType()
                 | ExtensionType()
                 | GenericType()
+                | ColumnType()
+                | DataFrameType()
             ):
-                raise NotImplementedError(f"Can't make assertion for type {type}")
+                self.logger.warning(f"Can't make assertion for type {type}")
 
             # Ensure exhaustiveness
             case _:

midas/generator/stubs.py

@@ -7,13 +7,16 @@ from midas.checker.types import (
     AliasType,
     AppliedType,
     BaseType,
+    ColumnType,
     ComplexType,
     ConstraintType,
+    DataFrameType,
     ExtensionType,
     Function,
     GenericType,
     OverloadedFunction,
     TopType,
+    TupleType,
     Type,
     TypeVar,
     UnitType,
@@ -30,6 +33,7 @@ class StubsGenerator:
         self.types: TypesRegistry = types
         self.stubs: list[ast.stmt] = []
         self.typing_imports: set[str] = set()
+        self.import_pandas: bool = False
         self.protocol_idx: int = 0
         self.stub_idx: int = 0
         self.type_var_idx: int = 0
@@ -38,6 +42,7 @@ class StubsGenerator:
     def generate_stubs(self) -> ast.Module:
         self.stubs = []
         self.typing_imports = set()
+        self.import_pandas = False
         for name, type in self.types._types.items():
             # Skip builtin types, not just based on name so the user can override
             # TODO: check if added members on builtin type
@@ -53,7 +58,7 @@ class StubsGenerator:
                     continue
             self.generate_stub(name, type)
 
-        imports = [
+        imports: list[ast.stmt] = [
             ast.ImportFrom(
                 module="__future__",
                 names=[ast.alias(name="annotations")],
@@ -70,6 +75,17 @@ class StubsGenerator:
                     level=0,
                 )
             )
+        if self.import_pandas:
+            imports.append(
+                ast.Import(
+                    names=[
+                        ast.alias(
+                            name="pandas",
+                            asname="pd",
+                        )
+                    ],
+                )
+            )
         return ast.Module(body=imports + self.stubs, type_ignores=[])
 
     def generate_stub(self, name: str, type: Type):
@@ -231,6 +247,31 @@ class StubsGenerator:
             case ConstraintType():
                 return self.dump_type(type.type)
 
+            case TupleType(items=items):
+                return ast.Subscript(
+                    value=ast.Name(id="tuple"),
+                    slice=ast.Tuple(
+                        elts=[self.dump_type(item) for item in items],
+                    ),
+                )
+
+            case ColumnType(type=inner):
+                self.import_pandas = True
+                return ast.Subscript(
+                    value=ast.Attribute(
+                        value=ast.Name(id="pd"),
+                        attr="Series",
+                    ),
+                    slice=self.dump_type(inner),
+                )
+
+            case DataFrameType():
+                self.import_pandas = True
+                return ast.Attribute(
+                    value=ast.Name(id="pd"),
+                    attr="DataFrame",
+                )
+
             case _:
                 assert_never(type)
 

midas/parser/midas.py

@@ -9,6 +9,7 @@ from midas.ast.midas import (
     Expr,
     ExtendStmt,
     ExtensionType,
+    FrameType,
     FunctionType,
     GenericType,
     GetExpr,
@@ -204,8 +205,10 @@ class MidasParser(Parser):
         return self.generic_type()
 
     def generic_type(self) -> Type:
-        type: Type = self.named_type()
+        type: NamedType = self.named_type()
         if self.check(TokenType.LEFT_BRACKET):
+            if type.name.lexeme == "Frame":
+                return self.frame_type()
             args: list[Type] = self.type_args()
             return GenericType(
                 location=Location.span(type.location, self.previous().get_location()),
@@ -224,7 +227,7 @@ class MidasParser(Parser):
         self.consume(TokenType.RIGHT_BRACKET, "Missing ']' after generic arguments")
         return args
 
-    def named_type(self) -> Type:
+    def named_type(self) -> NamedType:
         name: Token = self.consume_identifier("Expected type name")
         return NamedType(
             location=name.get_location(),
@@ -259,6 +262,32 @@ class MidasParser(Parser):
             members=members,
         )
 
+    def frame_type(self) -> FrameType:
+        keyword: Token = self.previous()
+        self.consume(TokenType.LEFT_BRACKET, "Expected '[' to start frame schema")
+
+        columns: list[FrameType.Column] = []
+        while not self.check(TokenType.RIGHT_BRACKET) and not self.is_at_end():
+            name: Token = self.advance()
+            self.consume(TokenType.COLON, "Expected ':' between column name and type")
+            type: Type = self.type_expr()
+            columns.append(
+                FrameType.Column(
+                    location=name.location_to(self.previous()),
+                    name=name,
+                    type=type,
+                )
+            )
+            if not self.match(TokenType.COMMA):
+                break
+
+        self.consume(TokenType.RIGHT_BRACKET, "Unclosed frame schema")
+
+        return FrameType(
+            location=keyword.location_to(self.previous()),
+            columns=columns,
+        )
+
     def constraint(self) -> Expr:
         """Parse a constraint
 

midas/parser/python.py

@@ -49,6 +49,7 @@ class UnsupportedSyntaxError(Exception):
 
 class PythonParser:
     CAST_FUNCTION = "cast"
+    UNSAFE_CAST_FUNCTION = "unsafe_cast"
 
     def parse_module(self, node: ast.Module) -> list[Stmt]:
         statements: list[Stmt] = []
@@ -423,6 +424,9 @@ class PythonParser:
             case ast.Call(func=ast.Name(id=self.CAST_FUNCTION)):
                 return self.parse_cast(node)
 
+            case ast.Call(func=ast.Name(id=self.UNSAFE_CAST_FUNCTION)):
+                return self.parse_cast(node)
+
             case ast.Call():
                 return self.parse_call(node)
 
@@ -527,16 +531,19 @@ class PythonParser:
         return expr
 
     def parse_cast(self, node: ast.Call) -> CastExpr:
+        assert isinstance(node.func, ast.Name)
+        func: str = node.func.id
         match node:
             case ast.Call(args=[type, expr], keywords=[]):
                 return CastExpr(
                     location=Location.from_ast(node),
                     type=self._parse_type(type),
                     expr=self.parse_expr(expr),
+                    unsafe=func == self.UNSAFE_CAST_FUNCTION,
                 )
             case _:
                 raise InvalidSyntaxError(
-                    f"Invalid call to {self.CAST_FUNCTION}, expected type and expression"
+                    f"Invalid call to {func}, expected type and expression"
                 )
 
     def parse_call(self, node: ast.Call) -> CallExpr:

midas/typing.py

@@ -0,0 +1,34 @@
+from typing import cast as typing_cast
+
+cast = typing_cast
+"""### Midas documentation
+Cast a value to a type.
+
+- **Compile-time**: tells the type checker that the return value has the designated type.
+- **Run-time**: generates assertions to ensure the value can be interpreted as the given type.
+
+---
+<br>
+<br>
+<br>
+
+_**Internal Python documentation**_
+"""
+
+
+unsafe_cast = typing_cast
+"""### Midas documentation
+Cast a value to a type.
+
+- **Compile-time**: tells the type checker that the return value has the designated type.
+- **Run-time**: -
+
+This operation is unsound, use at your own risk!
+
+---
+<br>
+<br>
+<br>
+
+_**Internal Python documentation**_
+"""

midas/utils.py

@@ -62,3 +62,4 @@ class UniversalJSONDumper:
 class TypedAST:
     stmts: list[p.Stmt]
     judgements: list[tuple[p.Expr, Type]]
+    evaluated_casts: list[p.CastExpr]

tests/cases/checker/04_custom_types.py.ref.json

@@ -1,6 +1,19 @@
 {
   "diagnostics": [],
   "judgments": [
+    {
+      "location": {
+        "from": "L4:30",
+        "to": "L4:36"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 123.45
+      },
+      "type": {
+        "name": "float"
+      }
+    },
     {
       "location": {
         "from": "L4:18",
@@ -16,7 +29,8 @@
         "expr": {
           "_type": "LiteralExpr",
           "value": 123.45
-        }
+        },
+        "unsafe": false
       },
       "type": {
         "name": "Meter",
@@ -25,6 +39,19 @@
         }
       }
     },
+    {
+      "location": {
+        "from": "L5:28",
+        "to": "L5:31"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 6.7
+      },
+      "type": {
+        "name": "float"
+      }
+    },
     {
       "location": {
         "from": "L5:15",
@@ -40,7 +67,8 @@
         "expr": {
           "_type": "LiteralExpr",
           "value": 6.7
-        }
+        },
+        "unsafe": false
       },
       "type": {
         "name": "Second",

tests/cases/generator/01_simple_types.py.ref.txt

@@ -7,68 +7,14 @@ Module(
         alias(name='Meter'),
         alias(name='Second')],
       level=0),
-    Assign(
-      targets=[
-        Name(id='__midas_a0__')],
-      value=Constant(value=123.45)),
-    Assert(
-      test=Call(
-        func=Name(id='isinstance'),
-        args=[
-          Name(id='__midas_a0__'),
-          Name(id='float')],
-        keywords=[]),
-      msg=JoinedStr(
-        values=[
-          Constant(value='01_simple_types.py:L3:19: CastError: Cannot cast '),
-          FormattedValue(
-            value=Attribute(
-              value=Call(
-                func=Name(id='type'),
-                args=[
-                  Name(id='__midas_a0__')],
-                keywords=[]),
-              attr='__name__'),
-            conversion=-1),
-          Constant(value=' to float')])),
     Assign(
       targets=[
         Name(id='distance')],
-      value=Name(id='__midas_a0__')),
+      value=Constant(value=123.45)),
-    Delete(
-      targets=[
-        Name(id='__midas_a0__')]),
-    Assign(
-      targets=[
-        Name(id='__midas_a1__')],
-      value=Constant(value=6.7)),
-    Assert(
-      test=Call(
-        func=Name(id='isinstance'),
-        args=[
-          Name(id='__midas_a1__'),
-          Name(id='float')],
-        keywords=[]),
-      msg=JoinedStr(
-        values=[
-          Constant(value='01_simple_types.py:L4:16: CastError: Cannot cast '),
-          FormattedValue(
-            value=Attribute(
-              value=Call(
-                func=Name(id='type'),
-                args=[
-                  Name(id='__midas_a1__')],
-                keywords=[]),
-              attr='__name__'),
-            conversion=-1),
-          Constant(value=' to float')])),
     Assign(
       targets=[
         Name(id='time')],
-      value=Name(id='__midas_a1__')),
+      value=Constant(value=6.7)),
-    Delete(
-      targets=[
-        Name(id='__midas_a1__')]),
     Assign(
       targets=[
         Name(id='speed')],

tests/serializer/python.py

@@ -263,6 +263,7 @@ class PythonAstJsonSerializer(
             "_type": "CastExpr",
             "type": expr.type.accept(self),
             "expr": expr.expr.accept(self),
+            "unsafe": expr.unsafe,
         }
 
     def visit_ternary_expr(self, expr: TernaryExpr) -> dict: