Merge pull request 'Dictionaries' (#13) from feat/dictionaries into main

Reviewed-on: #13

.gitignore

@@ -4,3 +4,6 @@ __pycache__
 venv
 .venv
 *.pyc
+uv.lock
+.python-version
+/out

README.md

@@ -5,3 +5,82 @@
 *Midas* aims at providing Python developers with a simple annotation system to enable compile-time integrity and data type checks, as well as generating runtime assertions.
 
 This framework is being developed as part of a Bachelor's Thesis by Louis Heredero at HEI Sion.
+
+## Requirements
+
+- Python 3.11+
+- [uv](https://docs.astral.sh/uv/getting-started/installation/)
+
+## Installation
+
+1. Clone the repository
+   ```shell
+   git clone https://git.kb28.ch/HEL/midas.git
+   ```
+2. Go in the project directory
+   ```shell
+   cd midas
+   ```
+3. Install the CLI as a user-wide tool
+   ```shell
+   uv tool install .
+   ```
+4. You can now run the `midas` command from anywhere
+   ```shell
+   midas --help
+   ```
+
+## Commands
+
+### Compiling
+
+> [!NOTE]
+> In the current state of the project, the `compile` command doesn't generate any runnable code, it only runs the parsers and type checker on the provided files
+
+```shell
+midas compile -t types.midas source.py
+```
+
+With the `compile` command, you can process a source Python file, with any number of custom type definition files (`-t FILE` option), and the type checker will verify the coherence of your program and generate the runnable code with valid syntax and runtime assertions.
+
+The optional `-l FILE` option lets you produce a highlighted version of the source code showing diagnostics from the type checker (see [Highlighting](#highlighting))
+
+### Highlighting
+
+```shell
+midas utils highlight source.py
+# or
+midas utils highlight types.midas
+```
+
+The `highlight` command takes in a source file (Python or Midas), runs the appropriate parser and outputs an HTML file containing the source code with added highlighting. This highlighting takes the form of hoverable annotations showing some of the parsed structures (e.g. a function definition, an assignment, a generic type, etc.)
+
+The optional `-o FILE` option can be used to specify an output path. By default, the file is printed in stdout (equivalent to `-o -`).
+
+### Dumping the AST
+
+```shell
+midas utils dump-ast source.py
+# or
+midas utils dump-ast types.midas
+```
+
+For debugging purposes, you can output the AST parsed from a Python or Midas file. For Python files, the `-p` flags lets you toggle the custom AST parsing. Without `-p`, the raw AST is returned, as produced by the builtin `ast` module. This flag has no effect on Midas files.
+
+The optional `-o FILE` option can be used to specify an output path. By default, the file is printed in stdout (equivalent to `-o -`).
+
+## Tests
+
+Several snapshot tests are available to assert the good behaviour of the parsers and type checker. They can be run as follows:
+
+```shell
+uv run -m tests.midas run -a
+uv run -m tests.python run -a
+uv run -m tests.checker run -a
+```
+
+**Available subcommands:**
+- Run all tests: `run -a`
+- Run specific tests: `run tests/cases/test1.py tests/cases/test2.py ...`
+- Update all tests: `update -a`
+- Update specific tests: `update tests/cases/test1.py tests/cases/test2.py ...`

core/ast/annotations.py

@@ -1,107 +0,0 @@
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from typing import Any, Generic, Optional, TypeVar
-
-from lexer.token import Token
-
-T = TypeVar("T")
-
-
-@dataclass(frozen=True)
-class Stmt(ABC):
-    @abstractmethod
-    def accept(self, visitor: Visitor[T]) -> T: ...
-
-    class Visitor(ABC, Generic[T]):
-        @abstractmethod
-        def visit_annotation_stmt(self, stmt: AnnotationStmt) -> T: ...
-
-
-@dataclass(frozen=True)
-class AnnotationStmt(Stmt):
-    name: Token
-    schema: Optional[SchemaExpr]
-
-    def accept(self, visitor: Stmt.Visitor[T]) -> T:
-        return visitor.visit_annotation_stmt(self)
-
-
-@dataclass(frozen=True)
-class Expr(ABC):
-    @abstractmethod
-    def accept(self, visitor: Visitor[T]) -> T: ...
-
-    class Visitor(ABC, Generic[T]):
-        @abstractmethod
-        def visit_wildcard_expr(self, expr: WildcardExpr) -> T: ...
-
-        @abstractmethod
-        def visit_literal_expr(self, expr: LiteralExpr) -> T: ...
-
-        @abstractmethod
-        def visit_type_expr(self, expr: TypeExpr) -> T: ...
-
-        @abstractmethod
-        def visit_constraint_expr(self, expr: ConstraintExpr) -> T: ...
-
-        @abstractmethod
-        def visit_schema_expr(self, expr: SchemaExpr) -> T: ...
-
-        @abstractmethod
-        def visit_schema_element_expr(self, expr: SchemaElementExpr) -> T: ...
-
-
-@dataclass(frozen=True)
-class WildcardExpr(Expr):
-    token: Token
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_wildcard_expr(self)
-
-
-@dataclass(frozen=True)
-class LiteralExpr(Expr):
-    value: Any
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_literal_expr(self)
-
-
-@dataclass(frozen=True)
-class TypeExpr(Expr):
-    name: Token
-    constraints: list[ConstraintExpr]
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_type_expr(self)
-
-
-@dataclass(frozen=True)
-class ConstraintExpr(Expr):
-    left: Expr
-    op: Token
-    right: Expr
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_constraint_expr(self)
-
-
-@dataclass(frozen=True)
-class SchemaExpr(Expr):
-    left: Token
-    elements: list[Expr]
-    right: Token
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_schema_expr(self)
-
-
-@dataclass(frozen=True)
-class SchemaElementExpr(Expr):
-    name: Optional[Token]
-    type: Optional[Expr]
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_schema_element_expr(self)

core/ast/midas.py

@@ -1,138 +0,0 @@
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from typing import Any, Generic, Optional, TypeVar
-
-from lexer.token import Token
-
-T = TypeVar("T")
-
-
-# Statements
-
-
-@dataclass(frozen=True)
-class Stmt(ABC):
-    @abstractmethod
-    def accept(self, visitor: Visitor[T]) -> T: ...
-
-    class Visitor(ABC, Generic[T]):
-        @abstractmethod
-        def visit_type_stmt(self, stmt: TypeStmt) -> T: ...
-
-        @abstractmethod
-        def visit_property_stmt(self, stmt: PropertyStmt) -> T: ...
-
-        @abstractmethod
-        def visit_op_stmt(self, stmt: OpStmt) -> T: ...
-
-        @abstractmethod
-        def visit_constraint_stmt(self, stmt: ConstraintStmt) -> T: ...
-
-
-@dataclass(frozen=True)
-class TypeStmt(Stmt):
-    name: Token
-    bases: list[TypeExpr]
-    body: Optional[TypeBodyExpr]
-
-    def accept(self, visitor: Stmt.Visitor[T]) -> T:
-        return visitor.visit_type_stmt(self)
-
-
-@dataclass(frozen=True)
-class PropertyStmt(Stmt):
-    name: Token
-    type: TypeExpr
-
-    def accept(self, visitor: Stmt.Visitor[T]) -> T:
-        return visitor.visit_property_stmt(self)
-
-
-@dataclass(frozen=True)
-class OpStmt(Stmt):
-    left: TypeExpr
-    op: Token
-    right: TypeExpr
-    result: TypeExpr
-
-    def accept(self, visitor: Stmt.Visitor[T]) -> T:
-        return visitor.visit_op_stmt(self)
-
-
-@dataclass(frozen=True)
-class ConstraintStmt(Stmt):
-    name: Token
-    constraint: ConstraintExpr
-
-    def accept(self, visitor: Stmt.Visitor[T]) -> T:
-        return visitor.visit_constraint_stmt(self)
-
-
-# Expressions
-
-
-@dataclass(frozen=True)
-class Expr(ABC):
-    @abstractmethod
-    def accept(self, visitor: Visitor[T]) -> T: ...
-
-    class Visitor(ABC, Generic[T]):
-        @abstractmethod
-        def visit_wildcard_expr(self, expr: WildcardExpr) -> T: ...
-
-        @abstractmethod
-        def visit_literal_expr(self, expr: LiteralExpr) -> T: ...
-
-        @abstractmethod
-        def visit_type_expr(self, expr: TypeExpr) -> T: ...
-
-        @abstractmethod
-        def visit_constraint_expr(self, expr: ConstraintExpr) -> T: ...
-
-        @abstractmethod
-        def visit_type_body_expr(self, expr: TypeBodyExpr) -> T: ...
-
-
-@dataclass(frozen=True)
-class WildcardExpr(Expr):
-    token: Token
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_wildcard_expr(self)
-
-
-@dataclass(frozen=True)
-class LiteralExpr(Expr):
-    value: Any
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_literal_expr(self)
-
-
-@dataclass(frozen=True)
-class TypeExpr(Expr):
-    name: Token
-    constraints: list[ConstraintExpr]
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_type_expr(self)
-
-
-@dataclass(frozen=True)
-class ConstraintExpr(Expr):
-    left: Expr
-    op: Token
-    right: Expr
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_constraint_expr(self)
-
-
-@dataclass(frozen=True)
-class TypeBodyExpr(Expr):
-    properties: list[PropertyStmt]
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_type_body_expr(self)

core/ast/printer.py

@@ -1,360 +0,0 @@
-from __future__ import annotations
-
-from contextlib import contextmanager
-from enum import Enum, auto
-import io
-from typing import Generator, Generic, Optional, Protocol, TypeVar
-
-import core.ast.annotations as a
-import core.ast.midas as m
-
-
-class _Level(Enum):
-    EMPTY = auto()
-    ACTIVE = auto()
-    LAST = auto()
-
-
-class Expr(Protocol):
-    def accept(self, printer: AstPrinter) -> None: ...
-
-
-T = TypeVar("T", bound=Expr)
-
-
-class AstPrinter(Generic[T]):
-    LAST_CHILD = "└── "
-    CHILD = "├── "
-    VERTICAL = "│   "
-    EMPTY = "    "
-
-    def __init__(self):
-        self._levels: list[_Level] = []
-        self._idx: Optional[int] = None
-        self._buf: io.StringIO = io.StringIO()
-
-    def print(self, expr: T):
-        self._buf = io.StringIO()
-        expr.accept(self)
-        return self._buf.getvalue()
-
-    @contextmanager
-    def _child_level(self, last: bool = False) -> Generator[None, None, None]:
-        self._levels.append(_Level.LAST if last else _Level.ACTIVE)
-        try:
-            yield
-        finally:
-            self._levels.pop()
-
-    def _mark_last(self):
-        if self._levels:
-            self._levels[-1] = _Level.LAST
-
-    def _write_line(self, text: str, *, last: bool = False):
-        if last:
-            self._mark_last()
-        indent: str = self._build_indent()
-        if self._idx is not None:
-            text = f"[{self._idx}] {text}"
-            self._idx = None
-        self._buf.write(indent + text + "\n")
-
-    def _build_indent(self) -> str:
-        parts: list[str] = []
-        for level in self._levels[:-1]:
-            parts.append(self.EMPTY if level == _Level.EMPTY else self.VERTICAL)
-        if self._levels:
-            if self._levels[-1] == _Level.LAST:
-                parts.append(self.LAST_CHILD)
-                self._levels[-1] = _Level.EMPTY
-            else:
-                parts.append(self.CHILD)
-        return "".join(parts)
-
-    def _write_optional_child(
-        self, label: str, child: Optional[T], *, last: bool = False
-    ):
-        if last:
-            self._mark_last()
-        if child is None:
-            self._write_line(f"{label}: None")
-        else:
-            self._write_line(label)
-            with self._child_level(last=True):
-                child.accept(self)
-
-
-class AnnotationAstPrinter(AstPrinter, a.Expr.Visitor[None], a.Stmt.Visitor[None]):
-    def visit_annotation_stmt(self, stmt: a.AnnotationStmt) -> None:
-        self._write_line("AnnotationStmt")
-        with self._child_level():
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_optional_child("schema", stmt.schema, last=True)
-
-    def visit_type_expr(self, expr: a.TypeExpr):
-        self._write_line("TypeExpr")
-        with self._child_level():
-            self._write_line(f'name: "{expr.name.lexeme}"')
-            self._write_line("constraints", last=True)
-            with self._child_level():
-                for i, constraint in enumerate(expr.constraints):
-                    self._idx = i
-                    if i == len(expr.constraints) - 1:
-                        self._mark_last()
-                    constraint.accept(self)
-
-    def visit_constraint_expr(self, expr: a.ConstraintExpr) -> None:
-        self._write_line("ConstraintExpr")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level():
-                self._mark_last()
-                expr.left.accept(self)
-
-            self._write_line(f"operator: {expr.op.lexeme}")
-
-            self._write_line("right", last=True)
-            with self._child_level():
-                self._mark_last()
-                expr.right.accept(self)
-
-    def visit_schema_expr(self, expr: a.SchemaExpr):
-        self._write_line("SchemaExpr")
-        with self._child_level():
-            for i, elmt in enumerate(expr.elements):
-                self._idx = i
-                if i == len(expr.elements) - 1:
-                    self._mark_last()
-                elmt.accept(self)
-
-    def visit_schema_element_expr(self, expr: a.SchemaElementExpr):
-        self._write_line("SchemaElementExpr")
-        with self._child_level():
-            name_text: str = "None" if expr.name is None else f'"{expr.name.lexeme}"'
-            self._write_line(f"name: {name_text}")
-            self._write_optional_child("type", expr.type, last=True)
-
-    def visit_wildcard_expr(self, expr: a.WildcardExpr) -> None:
-        self._write_line("WildcardExpr")
-
-    def visit_literal_expr(self, expr: a.LiteralExpr) -> None:
-        self._write_line("LiteralExpr")
-        with self._child_level():
-            self._write_line(f"value: {expr.value}", last=True)
-
-
-class AnnotationPrinter(a.Expr.Visitor[str], a.Stmt.Visitor[str]):
-    def print(self, expr: a.Expr | a.Stmt):
-        return expr.accept(self)
-
-    def visit_annotation_stmt(self, stmt: a.AnnotationStmt) -> str:
-        schema: str = ""
-        if stmt.schema is not None:
-            schema = stmt.schema.accept(self)
-        return f"{stmt.name.lexeme}{schema}"
-
-    def visit_type_expr(self, expr: a.TypeExpr) -> str:
-        parts: list[str] = [expr.name.lexeme]
-        for constraint in expr.constraints:
-            parts.append("(" + constraint.accept(self) + ")")
-        return " + ".join(parts)
-
-    def visit_constraint_expr(self, expr: a.ConstraintExpr) -> str:
-        parts: list[str] = [
-            expr.left.accept(self),
-            expr.op.lexeme,
-            expr.right.accept(self),
-        ]
-        return " ".join(parts)
-
-    def visit_schema_expr(self, expr: a.SchemaExpr) -> str:
-        res: str = expr.left.lexeme
-        res += ", ".join(elmt.accept(self) for elmt in expr.elements)
-        res += expr.right.lexeme
-        return res
-
-    def visit_schema_element_expr(self, expr: a.SchemaElementExpr) -> str:
-        parts: list[str] = []
-        if expr.name is not None:
-            parts.append(expr.name.lexeme)
-
-        if expr.type is None:
-            parts.append("_")
-        else:
-            parts.append(expr.type.accept(self))
-        return ": ".join(parts)
-
-    def visit_wildcard_expr(self, expr: a.WildcardExpr) -> str:
-        return "_"
-
-    def visit_literal_expr(self, expr: a.LiteralExpr) -> str:
-        return str(expr.value)
-
-
-class MidasAstPrinter(AstPrinter, m.Expr.Visitor[None], m.Stmt.Visitor[None]):
-    def visit_type_stmt(self, stmt: m.TypeStmt):
-        self._write_line("TypeStmt")
-        with self._child_level():
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line("bases")
-            with self._child_level():
-                for i, base in enumerate(stmt.bases):
-                    self._idx = i
-                    if i == len(stmt.bases) - 1:
-                        self._mark_last()
-                    base.accept(self)
-            self._write_optional_child("body", stmt.body, last=True)
-
-    def visit_property_stmt(self, stmt: m.PropertyStmt):
-        self._write_line("PropertyStmt")
-        with self._child_level():
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line("type", last=True)
-            with self._child_level():
-                self._mark_last()
-                stmt.type.accept(self)
-
-    def visit_op_stmt(self, stmt: m.OpStmt) -> None:
-        self._write_line("OpStmt")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level():
-                self._mark_last()
-                stmt.left.accept(self)
-
-            self._write_line(f'op: "{stmt.op.lexeme}"')
-
-            self._write_line("right")
-            with self._child_level():
-                self._mark_last()
-                stmt.right.accept(self)
-
-            self._write_line("result", last=True)
-            with self._child_level():
-                self._mark_last()
-                stmt.result.accept(self)
-
-    def visit_constraint_stmt(self, stmt: m.ConstraintStmt):
-        self._write_line("ConstraintStmt")
-        with self._child_level():
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line("constraint", last=True)
-            with self._child_level():
-                self._mark_last()
-                stmt.constraint.accept(self)
-
-    def visit_type_expr(self, expr: m.TypeExpr):
-        self._write_line("TypeExpr")
-        with self._child_level():
-            self._write_line(f'name: "{expr.name.lexeme}"')
-            self._write_line("constraints", last=True)
-            with self._child_level():
-                for i, constraint in enumerate(expr.constraints):
-                    self._idx = i
-                    if i == len(expr.constraints) - 1:
-                        self._mark_last()
-                    constraint.accept(self)
-
-    def visit_constraint_expr(self, expr: m.ConstraintExpr):
-        self._write_line("ConstraintExpr")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level():
-                self._mark_last()
-                expr.left.accept(self)
-
-            self._write_line(f"operator: {expr.op.lexeme}")
-
-            self._write_line("right", last=True)
-            with self._child_level():
-                self._mark_last()
-                expr.right.accept(self)
-
-    def visit_type_body_expr(self, expr: m.TypeBodyExpr):
-        self._write_line("TypeBodyExpr")
-        with self._child_level():
-            self._write_line("properties", last=True)
-            with self._child_level():
-                for i, property in enumerate(expr.properties):
-                    self._idx = i
-                    if i == len(expr.properties) - 1:
-                        self._mark_last()
-                    property.accept(self)
-
-    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> None:
-        self._write_line("WildcardExpr")
-
-    def visit_literal_expr(self, expr: m.LiteralExpr) -> None:
-        self._write_line("LiteralExpr")
-        with self._child_level():
-            self._write_line(f"value: {expr.value}", last=True)
-
-class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str]):
-    def __init__(self, indent: int = 4):
-        self.indent: int = indent
-        self.level: int = 0
-
-    def indented(self, text: str) -> str:
-        return " " * (self.level * self.indent) + text
-
-    def print(self, expr: m.Expr | m.Stmt):
-        self.level = 0
-        return expr.accept(self)
-
-    def visit_type_stmt(self, stmt: m.TypeStmt):
-        bases: list[str] = [
-            b.accept(self)
-            for b in stmt.bases
-        ]
-        
-        res: str = self.indented(f"type {stmt.name.lexeme}<{', '.join(bases)}>")
-        if stmt.body is not None:
-            res += " {\n"
-            self.level += 1
-            res += stmt.body.accept(self)
-            self.level -= 1
-            res += "\n" + self.indented("}")
-
-        return res
-
-    def visit_property_stmt(self, stmt: m.PropertyStmt):
-        return f"{stmt.name.lexeme}: {stmt.type.accept(self)}"
-
-    def visit_op_stmt(self, stmt: m.OpStmt):
-        left: str = stmt.left.accept(self)
-        op: str = stmt.op.lexeme
-        right: str = stmt.right.accept(self)
-        result: str = stmt.result.accept(self)
-        return self.indented(f"op <{left}> {op} <{right}> = <{result}>")
-
-    def visit_constraint_stmt(self, stmt: m.ConstraintStmt):
-        name: str = stmt.name.lexeme
-        constraint: str = stmt.constraint.accept(self)
-        return self.indented(f"constraint {name} = {constraint}")
-
-    def visit_type_expr(self, expr: m.TypeExpr):
-        parts: list[str] = [expr.name.lexeme]
-        for constraint in expr.constraints:
-            parts.append("(" + constraint.accept(self) + ")")
-        return " + ".join(parts)
-
-    def visit_constraint_expr(self, expr: m.ConstraintExpr):
-        parts: list[str] = [
-            expr.left.accept(self),
-            expr.op.lexeme,
-            expr.right.accept(self),
-        ]
-        return " ".join(parts)
-
-    def visit_type_body_expr(self, expr: m.TypeBodyExpr):
-        properties: list[str] = [
-            self.indented(prop.accept(self))
-            for prop in expr.properties
-        ]
-        return "\n".join(properties)
-
-    def visit_wildcard_expr(self, expr: m.WildcardExpr):
-        return "_"
-
-    def visit_literal_expr(self, expr: m.LiteralExpr):
-        return str(expr.value)

docs/architecture.typ

@@ -0,0 +1,150 @@
+#import "@preview/cetz:0.5.2": canvas, draw
+
+#let diagram-only = false
+
+#set document(
+  title: [Midas Architecture],
+  //author: "Louis Heredero",
+)
+
+#set text(
+  font: "Source Sans 3",
+)
+
+#let diagram = canvas({
+  let framed = draw.content.with(
+    padding: (x: .8em, y: 1em),
+    frame: "rect",
+    stroke: black,
+  )
+  let arrow = draw.line.with(mark: (end: ">", fill: black))
+  framed(
+    (0, 0),
+    name: "python-parser",
+  )[Python parser]
+
+  draw.content(
+    (rel: (0, 1), to: "python-parser.north"),
+    padding: 5pt,
+    anchor: "south",
+    name: "source-py",
+  )[_`source.py`_]
+  arrow("source-py", "python-parser")
+
+  framed(
+    (rel: (3, 0), to: "python-parser.east"),
+    anchor: "west",
+    name: "custom-parser",
+    align(center)[Custom python\ parser],
+  )
+
+  arrow("python-parser", "custom-parser", name: "arrow-python-ast")
+  draw.content(
+    "arrow-python-ast",
+    anchor: "south",
+    padding: 5pt,
+  )[`ast.Module`]
+
+  framed(
+    (rel: (-3, -2), to: "custom-parser.south"),
+    anchor: "east",
+    name: "python-resolver",
+  )[Python Resolver]
+  arrow(
+    "custom-parser",
+    ((), "|-", "python-resolver.east"),
+    "python-resolver",
+    name: "arrow-python-custom-ast",
+  )
+  draw.content(
+    (rel: (1.5, 0), to: "arrow-python-custom-ast.end"),
+    padding: 5pt,
+    anchor: "south",
+  )[P-AST#footnote[#strong[P]ython *AST*]<fn-past>]
+  draw.content(
+    "python-resolver.west",
+    padding: 5pt,
+    anchor: "south-east",
+  )[Resolved P-AST@fn-past]
+
+  draw.circle(
+    (rel: (1, -2), to: "custom-parser.south-east"),
+    radius: .4,
+    name: "midas-loader",
+  )
+  arrow(
+    "custom-parser",
+    "midas-loader",
+    name: "arrow-load-midas",
+    mark: (end: (symbol: ">", fill: black), start: "o"),
+  )
+  draw.content(
+    "arrow-load-midas",
+    anchor: "west",
+    padding: 5pt,
+  )[```python midas.using("types.midas")```]
+
+  framed(
+    (rel: (0, -2), to: "midas-loader.south"),
+    name: "midas-parser",
+  )[Midas lexer/parser]
+  arrow("midas-loader", "midas-parser", name: "arrow-midas-source")
+  draw.content(
+    "arrow-midas-source",
+    anchor: "west",
+    padding: 5pt,
+  )[_`types.midas`_]
+
+
+  framed(
+    (rel: (-2, 0), to: "midas-parser.west"),
+    anchor: "east",
+    name: "midas-resolver",
+  )[Midas Resolver]
+  arrow("midas-parser", "midas-resolver", name: "arrow-midas-ast")
+  draw.content(
+    "arrow-midas-ast",
+    anchor: "south",
+    padding: 5pt,
+  )[M-AST#footnote[#strong[M]idas *AST*]<fn-mast>]
+
+  framed(
+    (rel: (-3, 0), to: "midas-resolver.west"),
+    anchor: "east",
+    name: "checker",
+  )[Checker]
+  arrow("midas-resolver", "checker", name: "arrow-type-ctx")
+  arrow(
+    "python-resolver",
+    ((), "-|", "checker.north"),
+    "checker",
+  )
+  draw.content(
+    "arrow-type-ctx",
+    anchor: "south",
+    padding: 5pt,
+  )[Types context]
+})
+
+#show: doc => if diagram-only {
+  set page(width: auto, height: auto, margin: .5cm)
+  diagram
+} else { doc }
+
+#align(center, title())
+
+#v(1cm)
+
+#figure(
+  diagram,
+  caption: [Midas type-checker architecture],
+)
+
+== Components
+
+- *Python parser*: builtin Python AST parser, extracts abstract syntax from the raw Python source (```python ast.parse(...)```)
+- *Custom python parser*: converts the raw Python AST into custom, more suitable constructs, especially for type annotations
+- *Python resolver*: resolves bindings and references, tracks binding scopes
+- *Midas lexer/parser*: parses a Midas type definition file and extracts its AST
+- *Midas resolver*: walks the AST and fills the environment with the defined types and operations
+- *Checker*: evaluates expressions and checks type coherence

examples/00_syntax_prototype/02_custom_types.py

@@ -2,10 +2,6 @@
 # ruff: disable[F821]
 from __future__ import annotations
 
-# Prototype of custom type import to use valid Python syntax
-import midas
-midas.using("02_custom_types.midas")
-
 # A data-frame using a custom type
 df: Frame[
     location: GeoLocation
@@ -21,7 +17,7 @@ lat + lon  # Invalid operation
 # Registered operations are permitted
 lat1: Latitude = lat[0]
 lat2: Latitude = lat[1]
-lat_diff: LatitudeDiff = lat2 - lat1  # Valid operation
+lat_diff: Difference[Latitude] = lat2 - lat1  # Valid operation
 
 # In addition to the type, a column can have one or more constraints, either defined inline or in a separate file
 df2: Frame[

examples/00_syntax_prototype/03_custom_types_v2.midas

@@ -0,0 +1,73 @@
+// Simple custom type derived from float
+type Custom(float)
+
+// Simple custom types with constraints
+type Latitude(float) where (-90 <= _ <= 90)
+type Longitude(float) where (-180 <= _ <= 180)
+
+// Generic custom type (a Difference of T is derived from T, e.g. a difference of floats is a float
+type Difference[T](T)
+
+// Complex custom type, containing two values accessible through properties
+type GeoLocation {
+    lat: Latitude
+    lon: Longitude
+}
+
+// Define operations on our custom type
+extend GeoLocation {
+    // This type is compatible with the `-` operation with another GeoLocation
+    // i.e. you can subtract a GeoLocation from another GeoLocation, resulting
+    // in a Difference of GeoLocations
+    op __sub__(GeoLocation) -> Difference[GeoLocation]
+}
+
+// For complex generics, you need to specify how the genericity the properties
+// are handled
+type Difference[GeoLocation] {
+    lat: Difference[Latitude]
+    lon: Difference[Longitude]
+}
+
+// Simple operation defined on our custom types
+extend Latitude {
+    op __sub__(Latitude) -> Difference[Latitude]
+}
+
+extend Longitude {
+    op __sub__(Longitude) -> Difference[Longitude]
+}
+
+// Predefined custom predicates that can be referenced in other definitions
+predicate Positive(v: float) = v >= 0
+predicate StrictlyPositive(v: float) = v > 0
+predicate Equatorial(loc: GeoLocation) = (-10 <= loc.lat <= 10)
+predicate Arctic(loc: GeoLocation) = (loc.lat >= 66)
+
+type Person {
+    name: str
+
+    // Property with an inline constraint
+    age: int? where (0 <= _ < 150)
+
+    // Property referencing a predicate
+    height: float where StrictlyPositive
+
+    home: GeoLocation
+}
+
+// Custom complex type derived from another complex type, with a constraint
+// on a property
+// Multiple proposed syntaxes, not yet defined
+
+// Explicit, but new keyword
+type EquatorialPerson refines Person where Equatorial(_.home)
+
+// Explicit with existing keyword, might be confusing if expectations regarding 'is'
+type EquatorialPerson is Person where Equatorial(_.home)
+
+// Consistent and Python-friendly but can be confused with structural extension
+type EquatorialPerson(Person) where Equatorial(_.home)
+
+// Allow new properties, probably not useful
+type EquatorialPerson extends Person where Equatorial(_.home)

examples/00_syntax_prototype/04_functions.py

@@ -0,0 +1,15 @@
+# type: ignore
+# ruff: disable[F821]
+from __future__ import annotations
+
+
+def func(
+    col1: Column[float + (0 <= _ <= 1)],
+    col2: Column[float + (0 <= _ <= 1)],
+) -> Column[float + (0 <= _ <= 2)]:
+    result: Column[float + (0 <= _ <= 2)] = col1 + col2
+    return result
+
+
+def func2(a: int, /, b: float, *, c: str):
+    pass

examples/00_syntax_prototype/05_custom_types_v3.midas

@@ -0,0 +1,33 @@
+type Foo1 = float
+type Foo2 = float where (_ > 3)
+type Foo3 = int | float
+type Foo4 = int where (_ > 3) | float where (_ > 3)
+type Foo5 = (int | float) where (_ > 3)
+type Foo6 = {
+    foo: float
+    bar: float where (_ > 3)
+}
+
+type Foo7[T] = T where (_ > 3)
+type Foo8[A, B<:int] = {
+    a: A
+    b: B
+}
+
+type Complex = {
+    a: int
+    b: int
+}
+type Complex2 = Complex where (_.a > 3 & _.b < 5)
+
+predicate Positive(n: int) = n >= 0
+
+extend Foo1 {
+    op __add__(Foo1) -> Foo1
+}
+
+extend Foo7[T] {
+    op __add__(Foo7[T]) -> Foo7[T]
+}
+
+type Optional[T] = None | T

examples/01_simple_type_checking/01_simple_operations.py

@@ -0,0 +1,13 @@
+a: int = 3
+b: int = 4
+
+c = a + b  # -> int
+
+c = "invalid"  # -> can't assign str to int variable
+
+d = True
+e = d + d
+
+f: float = a
+
+f = -f

examples/01_simple_type_checking/02_simple_types.midas

@@ -0,0 +1,14 @@
+type Meter = float
+type Second = float
+type MeterPerSecond = float
+
+extend Meter {
+    def __add__: fn(Meter, /) -> Meter
+    def __sub__: fn(Meter, /) -> Meter
+    def __truediv__: fn(Second, /) -> MeterPerSecond
+}
+
+extend Second {
+    def __add__: fn(Second, /) -> Second
+    def __sub__: fn(Second, /) -> Second
+}

examples/01_simple_type_checking/02_simple_types.py

@@ -0,0 +1,6 @@
+# type: ignore
+# ruff: disable [F821]
+
+distance: Meter = cast(Meter, 123.45)
+time: Second = cast(Second, 6.7)
+speed = distance / time

examples/01_simple_type_checking/03_control_flow.py

@@ -0,0 +1,23 @@
+def minimum(x: int, y: int):
+    if x < y:
+        return x
+    else:
+        return y
+
+
+a = 15
+b = 72
+c = minimum(a, b)
+
+
+def factorial(n: int) -> int:
+    if n <= 1:
+        return 1
+    return n * factorial(n - 1)
+
+
+category = "Category 1" if a < 10 else "Category 2"
+
+
+def foo() -> None:
+    pass

examples/01_simple_type_checking/04_complex_types.midas

@@ -0,0 +1,21 @@
+type Meter = float
+
+extend Meter {
+    def __add__: fn(Meter, /) -> Meter
+    def __sub__: fn(Meter, /) -> Meter
+}
+
+type Coordinate = object
+
+extend Coordinate {
+    prop x: Meter
+    prop y: Meter
+}
+
+type Difference[T <: float] = T
+type MeterDifference = Difference[Meter]
+
+type CompDiff[T <: float] = {
+    prop d1: Difference[T]
+    prop d2: Difference[T]
+}

examples/01_simple_type_checking/04_complex_types.py

@@ -0,0 +1,37 @@
+# type: ignore
+# ruff: disable [F821]
+
+p1: Coordinate
+p2: Coordinate
+
+diff_x = p2.x - p1.x
+diff_y = p2.y - p1.y
+
+dist = diff_x + diff_y
+
+p2.x += cast(Meter, 1)
+p2.y = True  # invalid, wrong type
+p2.z = 3  # invalid, no property 'z' on Coordinate
+p2.x.a = 3  # invalid, no properties on Meter
+
+foo: list[float] = []
+
+append = foo.append
+
+foo.append("")  # invalid, must be float
+foo.append(2)
+append(True)  # invalid, must be float
+append(2)
+
+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
+
+d = foo[1:2]

examples/01_simple_type_checking/05_functions.py

@@ -0,0 +1,28 @@
+def incr(value: int):
+    return value + 1
+
+
+def decr(value: int):
+    return value - 1
+
+
+def foo(a: int, /, b: float, *, c: str):
+    return True
+
+
+r1 = foo()  # foo() missing 2 required positional arguments: 'a' and 'b'
+r2 = foo(1)  # foo() missing 1 required positional argument: 'b'
+r3 = foo(1, 2.0)  # foo() missing 1 required keyword-only argument: 'c'
+r4 = foo(1, b=2.0)  # foo() missing 1 required keyword-only argument: 'c'
+r5 = foo(1, 2.0, "test")  # foo() takes 2 positional arguments but 3 were given
+r6 = foo(1, 2.0, b=3.0)  # foo() got multiple values for argument 'b'
+r7 = foo(
+    a=1
+)  # foo() got some positional-only arguments passed as keyword arguments: 'a'
+r8 = foo(g="test")  # foo() got an unexpected keyword argument 'g'
+
+r9a = foo(1, 2.0, c="test")
+r9b = foo(1, b=2.0, c="test")
+r9c = foo(1, c="test", b=2.0)
+
+r10 = foo("a", 3, c=False)  # wrong argument types

examples/01_simple_type_checking/06_overloads.midas

@@ -0,0 +1,10 @@
+type T1 = object
+type T2 = object
+type Foo = object
+type T2b = T2
+
+extend Foo {
+    def bar: fn(T1, /) -> int
+    def bar: fn(T2, /) -> float
+    def bar: fn(T2b, /) -> int
+}

examples/01_simple_type_checking/06_overloads.py

@@ -0,0 +1,18 @@
+# type: ignore
+# ruff: disable [F821]
+
+foo: Foo
+t1: T1
+t2: T2
+
+a = foo.bar(t1)
+b = foo.bar(t2)
+
+func = foo.bar
+
+c = func(t1)
+d = func(t2)
+
+t2b: T2b
+
+e = foo.bar(t2b)

gen/gen.py

@@ -0,0 +1,159 @@
+import re
+from pathlib import Path
+
+HEADER = '''"""
+This file was generated by a script. Any manual changes might be overwritten.
+Please modify {defs_path} instead and run {gen_path}
+"""'''
+
+SECTION_TEMPLATE = """{banner}
+
+
+@dataclass(frozen=True, kw_only=True)
+class {base}(ABC):
+    location: Location
+
+    @abstractmethod
+    def accept(self, visitor: Visitor[T]) -> T: ...
+
+    class Visitor(ABC, Generic[T]):
+{visitor_methods}
+
+
+{classes}"""
+
+TEMPLATE = """{header}
+
+from __future__ import annotations
+
+{imports}
+
+T = TypeVar("T")
+
+{preamble}
+{sections}
+"""
+
+VISITOR_METHOD_TEMPLATE = """
+        @abstractmethod
+        def visit_{func_name}(self, {param}: {cls}) -> T: ...
+"""
+
+CLASS_TEMPLATE = """
+@dataclass(frozen=True)
+class {cls}({base}):
+{body}
+
+    def accept(self, visitor: {base}.Visitor[T]) -> T:
+        return visitor.visit_{func_name}(self)
+"""
+
+SECTION_REGEX = re.compile(
+    r"^###>\s*(?P<base>[^\n]*?)\s*\|\s*(?P<name>[^\n]*?)(\s*\|\s*(?P<param>[^\n]*?))?\s*?\n(?P<body>.*?)\n###<$",
+    re.MULTILINE | re.DOTALL,
+)
+
+IMPORTS_REGEX = re.compile(
+    r"^###>\s*Imports\s*?\n(?P<body>.*?)\n###<$",
+    re.MULTILINE | re.DOTALL,
+)
+
+PREAMBLE_REGEX = re.compile(
+    r"^###>\s*Preamble\s*?\n(?P<body>.*?)\n###<$",
+    re.MULTILINE | re.DOTALL,
+)
+
+
+def snake_case(text: str) -> str:
+    return re.sub(r"[A-Z]", lambda c: "_" + c.group().lower(), text).lower().strip("_")
+
+
+def make_visitor_method(cls: str, param: str):
+    method: str = VISITOR_METHOD_TEMPLATE.format(
+        func_name=snake_case(cls), param=param, cls=cls
+    )
+    return method.strip("\n")
+
+
+def make_class(name: str, cls: str, base: str):
+    body: str = cls.split("\n", 1)[1]
+    func_name: str = snake_case(name)
+    cls_def: str = CLASS_TEMPLATE.format(
+        cls=name,
+        base=base,
+        body=body,
+        func_name=func_name,
+    )
+    return cls_def.strip("\n")
+
+
+def make_banner(text: str) -> str:
+    middle: str = f"# {text} #"
+    rule: str = "#" * len(middle)
+    return "\n".join((rule, middle, rule))
+
+
+def make_section(full_name: str, base: str, param: str, body: str) -> str:
+    print(f"    Generating {full_name}")
+    visitor_methods: list[str] = []
+    classes: list[str] = []
+    definitions: list[str] = body.strip("\n").split("\n\n\n")
+    for cls in definitions:
+        cls = cls.strip("\n")
+        name: str = re.match("class (.*?):", cls).group(1)  # type: ignore
+        print(f"        Processing {name}")
+        visitor_methods.append(make_visitor_method(name, param))
+        classes.append(make_class(name, cls, base))
+
+    return SECTION_TEMPLATE.format(
+        banner=make_banner(full_name),
+        base=base,
+        visitor_methods="\n\n".join(visitor_methods),
+        classes="\n\n\n".join(classes),
+    )
+
+
+def generate(definitions_path: Path, out_path: Path):
+    print(f"Processing generating {out_path} from {definitions_path}")
+    root_dir: Path = Path(__file__).parent.parent
+    rel_path: Path = definitions_path.relative_to(root_dir)
+    src: str = definitions_path.read_text()
+    sections: list[str] = []
+
+    imports: str = ""
+    if m := IMPORTS_REGEX.search(src):
+        imports = m.group("body").strip("\n")
+
+    preamble: str = ""
+    if m := PREAMBLE_REGEX.search(src):
+        preamble = m.group("body")
+
+    for section_m in SECTION_REGEX.finditer(src):
+        full_name: str = section_m.group("name")
+        base: str = section_m.group("base")
+        param: str = section_m.group("param") or base.lower()
+        body: str = section_m.group("body")
+        sections.append(make_section(full_name, base, param, body))
+
+    result: str = TEMPLATE.format(
+        header=HEADER.format(
+            defs_path=rel_path,
+            gen_path=Path(__file__).relative_to(root_dir),
+        ),
+        imports=imports,
+        preamble=preamble,
+        sections="\n\n\n".join(sections),
+    )
+    out_path.write_text(result)
+
+
+def main():
+    root: Path = Path(__file__).parent.parent
+    defs_dir: Path = root / "gen"
+    ast_dir: Path = root / "midas" / "ast"
+    generate(defs_dir / "midas.py", ast_dir / "midas.py")
+    generate(defs_dir / "python.py", ast_dir / "python.py")
+
+
+if __name__ == "__main__":
+    main()

gen/midas.py

@@ -0,0 +1,144 @@
+# type: ignore
+# ruff: disable[F821, F401]
+
+###> Imports
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import Any, Generic, Optional, TypeVar
+
+from midas.ast.location import Location
+from midas.lexer.token import Token
+
+###<
+
+
+###> Preamble
+@dataclass(frozen=True, kw_only=True)
+class TypeParam:
+    location: Location
+    name: Token
+    bound: Optional[Type]
+
+
+class MemberKind(Enum):
+    PROPERTY = auto()
+    METHOD = auto()
+
+
+###<
+
+
+###> Stmt | Statements
+class TypeStmt:
+    name: Token
+    params: list[TypeParam]
+    type: Type
+
+
+class MemberStmt:
+    name: Token
+    type: Type
+    kind: MemberKind
+
+
+class ExtendStmt:
+    name: Token
+    params: list[TypeParam]
+    members: list[MemberStmt]
+
+
+class PredicateStmt:
+    name: Token
+    subject: Token
+    type: Type
+    condition: Expr
+
+
+###<
+
+
+###> Expr | Expressions
+
+
+class LogicalExpr:
+    left: Expr
+    operator: Token
+    right: Expr
+
+
+class BinaryExpr:
+    left: Expr
+    operator: Token
+    right: Expr
+
+
+class UnaryExpr:
+    operator: Token
+    right: Expr
+
+
+class GetExpr:
+    expr: Expr
+    name: Token
+
+
+class VariableExpr:
+    name: Token
+
+
+class GroupingExpr:
+    expr: Expr
+
+
+class LiteralExpr:
+    value: Any
+
+
+class WildcardExpr:
+    token: Token
+
+
+###<
+
+###> Type | Types
+
+
+class NamedType:
+    name: Token
+
+
+class GenericType:
+    type: Type
+    args: list[Type]
+
+
+class ConstraintType:
+    type: Type
+    constraint: Expr
+
+
+class ComplexType:
+    members: list[MemberStmt]
+
+
+class ExtensionType:
+    base: Type
+    extension: ComplexType
+
+
+class FunctionType:
+    pos_args: list[Argument]
+    args: list[Argument]
+    kw_args: list[Argument]
+    returns: Type
+
+    @dataclass(frozen=True, kw_only=True)
+    class Argument:
+        location: Optional[Location] = None
+        name: Optional[Token]
+        type: Type
+        required: bool
+
+
+###<

gen/python.py

@@ -0,0 +1,180 @@
+# type: ignore
+# ruff: disable[F821, F401]
+
+###> Imports
+import ast
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Generic, Optional, TypeVar
+
+from midas.ast.location import Location
+
+###<
+
+
+###> MidasType | Type annotations | node
+class BaseType:
+    base: str
+    param: Optional[MidasType]
+
+
+class ConstraintType:
+    type: MidasType
+    constraint: ast.expr
+
+
+class FrameColumn:
+    name: Optional[str]
+    type: Optional[MidasType]
+
+
+class FrameType:
+    columns: list[FrameColumn]
+
+
+###<
+
+
+###> Stmt | Statements
+class ExpressionStmt:
+    expr: Expr
+
+
+class Function:
+    name: str
+    posonlyargs: list[Argument]
+    args: list[Argument]
+    sink: Optional[Argument]
+    kwonlyargs: list[Argument]
+    kw_sink: Optional[Argument]
+    returns: Optional[MidasType]
+    body: list[Stmt]
+
+    @dataclass(frozen=True, kw_only=True)
+    class Argument:
+        location: Optional[Location] = None
+        name: str
+        type: Optional[MidasType]
+        default: Optional[Expr]
+
+    @property
+    def all_args(self) -> list[Argument]:
+        return self.posonlyargs + self.args + self.kwonlyargs
+
+
+class TypeAssign:
+    name: str
+    type: MidasType
+
+
+class AssignStmt:
+    targets: list[Expr]
+    value: Expr
+
+
+class ReturnStmt:
+    value: Optional[Expr]
+
+
+class IfStmt:
+    test: Expr
+    body: list[Stmt]
+    orelse: list[Stmt]
+
+
+class Pass:
+    pass
+
+
+class ForStmt:
+    target: Expr
+    iterator: Expr
+    body: list[Stmt]
+
+
+class RawStmt:
+    stmt: ast.stmt
+
+
+###<
+
+
+###> Expr | Expressions
+class BinaryExpr:
+    left: Expr
+    operator: ast.operator
+    right: Expr
+
+
+class CompareExpr:
+    left: Expr
+    operator: ast.cmpop
+    right: Expr
+
+
+class UnaryExpr:
+    operator: ast.unaryop
+    right: Expr
+
+
+class CallExpr:
+    callee: Expr
+    arguments: list[Expr]
+    keywords: dict[str, Expr]
+
+
+class GetExpr:
+    object: Expr
+    name: str
+
+
+class LiteralExpr:
+    value: Any
+
+
+class VariableExpr:
+    name: str
+
+
+class LogicalExpr:
+    left: Expr
+    operator: ast.boolop
+    right: Expr
+
+
+class CastExpr:
+    type: MidasType
+    expr: Expr
+
+
+class TernaryExpr:
+    test: Expr
+    if_true: Expr
+    if_false: Expr
+
+
+class ListExpr:
+    items: list[Expr]
+
+
+class DictExpr:
+    keys: list[Optional[Expr]]
+    values: list[Expr]
+
+
+class SubscriptExpr:
+    object: Expr
+    index: Expr
+
+
+class SliceExpr:
+    lower: Optional[Expr]
+    upper: Optional[Expr]
+    step: Optional[Expr]
+
+
+class RawExpr:
+    expr: ast.expr
+
+
+###<

lexer/annotations.py

@@ -1,102 +0,0 @@
-from lexer.base import Lexer
-from lexer.keyword import ANNOTATION_KEYWORDS
-from lexer.token import TokenType
-
-
-class AnnotationLexer(Lexer):
-    def scan_token(self) -> None:
-        char: str = self.advance()
-        match char:
-            case "(":
-                self.add_token(TokenType.LEFT_PAREN)
-            case ")":
-                self.add_token(TokenType.RIGHT_PAREN)
-            case "[":
-                self.add_token(TokenType.LEFT_BRACKET)
-            case "]":
-                self.add_token(TokenType.RIGHT_BRACKET)
-            case "<":
-                self.add_token(
-                    TokenType.LESS_EQUAL if self.match("=") else TokenType.LESS
-                )
-            case ">":
-                self.add_token(
-                    TokenType.GREATER_EQUAL if self.match("=") else TokenType.GREATER
-                )
-            case "=":
-                self.add_token(
-                    TokenType.EQUAL_EQUAL if self.match("=") else TokenType.EQUAL
-                )
-            case "!":
-                if self.match("="):
-                    self.add_token(TokenType.BANG_EQUAL)
-                else:
-                    self.error("Unexpected single bang. Did you mean '!=' ?")
-            case ":":
-                self.add_token(TokenType.COLON)
-            case ",":
-                self.add_token(TokenType.COMMA)
-            case "_":
-                self.add_token(TokenType.UNDERSCORE)
-            case "+":
-                self.add_token(TokenType.PLUS)
-            case "#":
-                self.scan_comment()
-            case "\n":
-                self.add_token(TokenType.NEWLINE)
-            case " " | "\r" | "\t":
-                # Consume all whitespace characters until EOL or EOF
-                while (
-                    self.peek().isspace()
-                    and self.peek() != "\n"
-                    and not self.is_at_end()
-                ):
-                    self.advance()
-                self.add_token(TokenType.WHITESPACE)
-            case _:
-                if char.isdigit():
-                    self.scan_number()
-                elif char.isalpha():
-                    self.scan_identifier()
-                else:
-                    self.error("Unexpected character")
-        return None
-
-    def scan_number(self):
-        """Scan the rest of number and add it as a token
-
-        This method handles both simple integers and floats. Scientific notation
-        and base prefixes (0x, 0b, 0o) are not supported
-        """
-        while self.peek().isdigit():
-            self.advance()
-
-        if self.peek() == "." and self.peek_next().isdigit():
-            self.advance()
-            while self.peek().isdigit():
-                self.advance()
-
-        value: float = float(self.source[self.start : self.idx])
-        self.add_token(TokenType.NUMBER, value)
-
-    def scan_identifier(self):
-        """Scan the rest of an identifier and add it as a token
-
-        An identifier starts with a letter, followed by any number of
-        alphanumerical characters or underscores
-        """
-        while self.peek().isalnum() or self.peek() == "_":
-            self.advance()
-        
-        lexeme: str = self.source[self.start : self.idx]
-        token_type: TokenType = ANNOTATION_KEYWORDS.get(lexeme, TokenType.IDENTIFIER)
-        self.add_token(token_type)
-
-    def scan_comment(self):
-        """Scan the rest of a comment and add it as a token
-
-        A comment starts with a `#` character and ends at the EOL/EOF
-        """
-        while self.peek() != "\n" and not self.is_at_end():
-            self.advance()
-        self.add_token(TokenType.COMMENT)

lexer/keyword.py

@@ -1,16 +0,0 @@
-from lexer.token import TokenType
-
-ANNOTATION_KEYWORDS: dict[str, TokenType] = {
-    "True": TokenType.TRUE,
-    "False": TokenType.FALSE,
-    "None": TokenType.NONE,
-}
-
-MIDAS_KEYWORDS: dict[str, TokenType] = {
-    "type": TokenType.TYPE,
-    "op": TokenType.OP,
-    "constraint": TokenType.CONSTRAINT,
-    "true": TokenType.TRUE,
-    "false": TokenType.FALSE,
-    "none": TokenType.NONE,
-}

lexer/token.py

@@ -1,59 +0,0 @@
-from dataclasses import dataclass
-from enum import Enum, auto
-from typing import Any
-
-from lexer.position import Position
-
-
-class TokenType(Enum):
-    # Punctuation
-    LEFT_PAREN = auto()
-    RIGHT_PAREN = auto()
-    LEFT_BRACKET = auto()
-    RIGHT_BRACKET = auto()
-    LEFT_BRACE = auto()
-    RIGHT_BRACE = auto()
-    COLON = auto()
-    COMMA = auto()
-    UNDERSCORE = auto()
-
-    # Operators
-    PLUS = auto()
-    MINUS = auto()
-    STAR = auto()
-    SLASH = auto()
-    GREATER = auto()
-    GREATER_EQUAL = auto()
-    LESS = auto()
-    LESS_EQUAL = auto()
-    EQUAL = auto()
-    EQUAL_EQUAL = auto()
-    BANG_EQUAL = auto()
-
-    # Literals
-    IDENTIFIER = auto()
-    NUMBER = auto()
-    TRUE = auto()
-    FALSE = auto()
-    NONE = auto()
-
-    # Keywords
-    TYPE = auto()
-    OP = auto()
-    CONSTRAINT = auto()
-
-    # Misc
-    COMMENT = auto()
-    WHITESPACE = auto()
-    EOF = auto()
-    NEWLINE = auto()
-
-
-@dataclass(frozen=True)
-class Token:
-    """A scanned token"""
-
-    type: TokenType
-    lexeme: str
-    value: Any
-    position: Position

midas/ast/location.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Protocol
+
+
+class HasLocation(Protocol):
+    lineno: int
+    col_offset: int
+    end_lineno: Optional[int]
+    end_col_offset: Optional[int]
+
+
+@dataclass(frozen=True, kw_only=True)
+class Location:
+    lineno: int
+    col_offset: int
+    end_lineno: Optional[int]
+    end_col_offset: Optional[int]
+
+    @staticmethod
+    def from_ast(obj: HasLocation) -> Location:
+        return Location(
+            lineno=obj.lineno,
+            col_offset=obj.col_offset,
+            end_lineno=obj.end_lineno,
+            end_col_offset=obj.end_col_offset,
+        )
+
+    @staticmethod
+    def span(start: Location, end: Location) -> Location:
+        return Location(
+            lineno=start.lineno,
+            col_offset=start.col_offset,
+            end_lineno=end.lineno,
+            end_col_offset=end.end_col_offset,
+        )

midas/ast/midas.py

@@ -0,0 +1,295 @@
+"""
+This file was generated by a script. Any manual changes might be overwritten.
+Please modify gen/midas.py instead and run gen/gen.py
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import Any, Generic, Optional, TypeVar
+
+from midas.ast.location import Location
+from midas.lexer.token import Token
+
+T = TypeVar("T")
+
+@dataclass(frozen=True, kw_only=True)
+class TypeParam:
+    location: Location
+    name: Token
+    bound: Optional[Type]
+
+
+class MemberKind(Enum):
+    PROPERTY = auto()
+    METHOD = auto()
+
+
+##############
+# Statements #
+##############
+
+
+@dataclass(frozen=True, kw_only=True)
+class Stmt(ABC):
+    location: Location
+
+    @abstractmethod
+    def accept(self, visitor: Visitor[T]) -> T: ...
+
+    class Visitor(ABC, Generic[T]):
+        @abstractmethod
+        def visit_type_stmt(self, stmt: TypeStmt) -> T: ...
+
+        @abstractmethod
+        def visit_member_stmt(self, stmt: MemberStmt) -> T: ...
+
+        @abstractmethod
+        def visit_extend_stmt(self, stmt: ExtendStmt) -> T: ...
+
+        @abstractmethod
+        def visit_predicate_stmt(self, stmt: PredicateStmt) -> T: ...
+
+
+@dataclass(frozen=True)
+class TypeStmt(Stmt):
+    name: Token
+    params: list[TypeParam]
+    type: Type
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_type_stmt(self)
+
+
+@dataclass(frozen=True)
+class MemberStmt(Stmt):
+    name: Token
+    type: Type
+    kind: MemberKind
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_member_stmt(self)
+
+
+@dataclass(frozen=True)
+class ExtendStmt(Stmt):
+    name: Token
+    params: list[TypeParam]
+    members: list[MemberStmt]
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_extend_stmt(self)
+
+
+@dataclass(frozen=True)
+class PredicateStmt(Stmt):
+    name: Token
+    subject: Token
+    type: Type
+    condition: Expr
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_predicate_stmt(self)
+
+
+###############
+# Expressions #
+###############
+
+
+@dataclass(frozen=True, kw_only=True)
+class Expr(ABC):
+    location: Location
+
+    @abstractmethod
+    def accept(self, visitor: Visitor[T]) -> T: ...
+
+    class Visitor(ABC, Generic[T]):
+        @abstractmethod
+        def visit_logical_expr(self, expr: LogicalExpr) -> T: ...
+
+        @abstractmethod
+        def visit_binary_expr(self, expr: BinaryExpr) -> T: ...
+
+        @abstractmethod
+        def visit_unary_expr(self, expr: UnaryExpr) -> T: ...
+
+        @abstractmethod
+        def visit_get_expr(self, expr: GetExpr) -> T: ...
+
+        @abstractmethod
+        def visit_variable_expr(self, expr: VariableExpr) -> T: ...
+
+        @abstractmethod
+        def visit_grouping_expr(self, expr: GroupingExpr) -> T: ...
+
+        @abstractmethod
+        def visit_literal_expr(self, expr: LiteralExpr) -> T: ...
+
+        @abstractmethod
+        def visit_wildcard_expr(self, expr: WildcardExpr) -> T: ...
+
+
+@dataclass(frozen=True)
+class LogicalExpr(Expr):
+    left: Expr
+    operator: Token
+    right: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_logical_expr(self)
+
+
+@dataclass(frozen=True)
+class BinaryExpr(Expr):
+    left: Expr
+    operator: Token
+    right: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_binary_expr(self)
+
+
+@dataclass(frozen=True)
+class UnaryExpr(Expr):
+    operator: Token
+    right: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_unary_expr(self)
+
+
+@dataclass(frozen=True)
+class GetExpr(Expr):
+    expr: Expr
+    name: Token
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_get_expr(self)
+
+
+@dataclass(frozen=True)
+class VariableExpr(Expr):
+    name: Token
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_variable_expr(self)
+
+
+@dataclass(frozen=True)
+class GroupingExpr(Expr):
+    expr: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_grouping_expr(self)
+
+
+@dataclass(frozen=True)
+class LiteralExpr(Expr):
+    value: Any
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_literal_expr(self)
+
+
+@dataclass(frozen=True)
+class WildcardExpr(Expr):
+    token: Token
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_wildcard_expr(self)
+
+
+#########
+# Types #
+#########
+
+
+@dataclass(frozen=True, kw_only=True)
+class Type(ABC):
+    location: Location
+
+    @abstractmethod
+    def accept(self, visitor: Visitor[T]) -> T: ...
+
+    class Visitor(ABC, Generic[T]):
+        @abstractmethod
+        def visit_named_type(self, type: NamedType) -> T: ...
+
+        @abstractmethod
+        def visit_generic_type(self, type: GenericType) -> T: ...
+
+        @abstractmethod
+        def visit_constraint_type(self, type: ConstraintType) -> T: ...
+
+        @abstractmethod
+        def visit_complex_type(self, type: ComplexType) -> T: ...
+
+        @abstractmethod
+        def visit_extension_type(self, type: ExtensionType) -> T: ...
+
+        @abstractmethod
+        def visit_function_type(self, type: FunctionType) -> T: ...
+
+
+@dataclass(frozen=True)
+class NamedType(Type):
+    name: Token
+
+    def accept(self, visitor: Type.Visitor[T]) -> T:
+        return visitor.visit_named_type(self)
+
+
+@dataclass(frozen=True)
+class GenericType(Type):
+    type: Type
+    args: list[Type]
+
+    def accept(self, visitor: Type.Visitor[T]) -> T:
+        return visitor.visit_generic_type(self)
+
+
+@dataclass(frozen=True)
+class ConstraintType(Type):
+    type: Type
+    constraint: Expr
+
+    def accept(self, visitor: Type.Visitor[T]) -> T:
+        return visitor.visit_constraint_type(self)
+
+
+@dataclass(frozen=True)
+class ComplexType(Type):
+    members: list[MemberStmt]
+
+    def accept(self, visitor: Type.Visitor[T]) -> T:
+        return visitor.visit_complex_type(self)
+
+
+@dataclass(frozen=True)
+class ExtensionType(Type):
+    base: Type
+    extension: ComplexType
+
+    def accept(self, visitor: Type.Visitor[T]) -> T:
+        return visitor.visit_extension_type(self)
+
+
+@dataclass(frozen=True)
+class FunctionType(Type):
+    pos_args: list[Argument]
+    args: list[Argument]
+    kw_args: list[Argument]
+    returns: Type
+
+    @dataclass(frozen=True, kw_only=True)
+    class Argument:
+        location: Optional[Location] = None
+        name: Optional[Token]
+        type: Type
+        required: bool
+
+    def accept(self, visitor: Type.Visitor[T]) -> T:
+        return visitor.visit_function_type(self)

midas/ast/printer.py

@@ -0,0 +1,789 @@
+from __future__ import annotations
+
+import ast
+import io
+from contextlib import contextmanager
+from enum import Enum, auto
+from typing import Generator, Generic, Optional, Protocol, TypeVar
+
+import midas.ast.midas as m
+import midas.ast.python as p
+
+
+class _Level(Enum):
+    EMPTY = auto()
+    ACTIVE = auto()
+    LAST = auto()
+
+
+class Expr(Protocol):
+    def accept(self, printer: AstPrinter) -> None: ...
+
+
+T = TypeVar("T", bound=Expr)
+
+
+class AstPrinter(Generic[T]):
+    LAST_CHILD = "└── "
+    CHILD = "├── "
+    VERTICAL = "│   "
+    EMPTY = "    "
+
+    def __init__(self):
+        self._levels: list[_Level] = []
+        self._idx: Optional[int] = None
+        self._buf: io.StringIO = io.StringIO()
+
+    def print(self, expr: T):
+        self._buf = io.StringIO()
+        expr.accept(self)
+        return self._buf.getvalue()
+
+    @contextmanager
+    def _child_level(self, single: bool = False) -> Generator[None, None, None]:
+        self._levels.append(_Level.LAST if single else _Level.ACTIVE)
+        try:
+            yield
+        finally:
+            self._levels.pop()
+
+    def _mark_last(self):
+        if self._levels:
+            self._levels[-1] = _Level.LAST
+
+    def _write_line(self, text: str, *, last: bool = False):
+        if last:
+            self._mark_last()
+        indent: str = self._build_indent()
+        if self._idx is not None:
+            text = f"[{self._idx}] {text}"
+            self._idx = None
+        self._buf.write(indent + text + "\n")
+
+    def _build_indent(self) -> str:
+        parts: list[str] = []
+        for level in self._levels[:-1]:
+            parts.append(self.EMPTY if level == _Level.EMPTY else self.VERTICAL)
+        if self._levels:
+            if self._levels[-1] == _Level.LAST:
+                parts.append(self.LAST_CHILD)
+                self._levels[-1] = _Level.EMPTY
+            else:
+                parts.append(self.CHILD)
+        return "".join(parts)
+
+    def _write_optional_child(
+        self, label: str, child: Optional[T], *, last: bool = False
+    ):
+        if last:
+            self._mark_last()
+        if child is None:
+            self._write_line(f"{label}: None")
+        else:
+            self._write_line(label)
+            with self._child_level(single=True):
+                child.accept(self)
+
+
+class MidasAstPrinter(
+    AstPrinter, m.Expr.Visitor[None], m.Stmt.Visitor[None], m.Type.Visitor[None]
+):
+    # Statements
+
+    def visit_type_stmt(self, stmt: m.TypeStmt) -> None:
+        self._write_line("TypeStmt")
+        with self._child_level():
+            self._write_line(f'name: "{stmt.name.lexeme}"')
+            self._write_line("params")
+            with self._child_level():
+                for i, param in enumerate(stmt.params):
+                    self._idx = i
+                    if i == len(stmt.params) - 1:
+                        self._mark_last()
+                    self._print_type_param(param)
+            self._write_line("type", last=True)
+            with self._child_level(single=True):
+                stmt.type.accept(self)
+
+    def _print_type_param(self, param: m.TypeParam) -> None:
+        self._write_line("Param")
+        with self._child_level():
+            self._write_line(f'name: "{param.name.lexeme}"')
+            self._write_optional_child("bound", param.bound, last=True)
+
+    def visit_member_stmt(self, stmt: m.MemberStmt):
+        self._write_line("MemberStmt")
+        with self._child_level():
+            self._write_line(f"kind: {stmt.kind.name}")
+            self._write_line(f'name: "{stmt.name.lexeme}"')
+            self._write_line("type", last=True)
+            with self._child_level(single=True):
+                stmt.type.accept(self)
+
+    def visit_extend_stmt(self, stmt: m.ExtendStmt) -> None:
+        self._write_line("ExtendStmt")
+        with self._child_level():
+            self._write_line("params")
+            with self._child_level():
+                for i, param in enumerate(stmt.params):
+                    self._idx = i
+                    if i == len(stmt.params) - 1:
+                        self._mark_last()
+                    self._print_type_param(param)
+            self._write_line(f'name: "{stmt.name.lexeme}"')
+            self._write_line("params")
+            with self._child_level():
+                for i, param in enumerate(stmt.params):
+                    self._idx = i
+                    if i == len(stmt.params) - 1:
+                        self._mark_last()
+                    self._print_type_param(param)
+            self._write_line("members", last=True)
+            with self._child_level():
+                for i, member in enumerate(stmt.members):
+                    self._idx = i
+                    if i == len(stmt.members) - 1:
+                        self._mark_last()
+                    member.accept(self)
+
+    def visit_predicate_stmt(self, stmt: m.PredicateStmt):
+        self._write_line("PredicateStmt")
+        with self._child_level():
+            self._write_line(f'name: "{stmt.name.lexeme}"')
+            self._write_line(f'subject: "{stmt.subject.lexeme}"')
+            self._write_line("type")
+            with self._child_level(single=True):
+                stmt.type.accept(self)
+            self._write_line("condition", last=True)
+            with self._child_level(single=True):
+                stmt.condition.accept(self)
+
+    # Expressions
+
+    def visit_logical_expr(self, expr: m.LogicalExpr):
+        self._write_line("LogicalExpr")
+        with self._child_level():
+            self._write_line("left")
+            with self._child_level(single=True):
+                expr.left.accept(self)
+
+            self._write_line(f"operator: {expr.operator.lexeme}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_binary_expr(self, expr: m.BinaryExpr):
+        self._write_line("BinaryExpr")
+        with self._child_level():
+            self._write_line("left")
+            with self._child_level(single=True):
+                expr.left.accept(self)
+
+            self._write_line(f"operator: {expr.operator.lexeme}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_unary_expr(self, expr: m.UnaryExpr):
+        self._write_line("UnaryExpr")
+        with self._child_level():
+            self._write_line(f"operator: {expr.operator.lexeme}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_get_expr(self, expr: m.GetExpr):
+        self._write_line("GetExpr")
+        with self._child_level():
+            self._write_line("expr")
+            with self._child_level(single=True):
+                expr.expr.accept(self)
+            self._write_line(f'name: "{expr.name.lexeme}"', last=True)
+
+    def visit_variable_expr(self, expr: m.VariableExpr):
+        self._write_line("VariableExpr")
+        with self._child_level():
+            self._write_line(f'name: "{expr.name.lexeme}"', last=True)
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr):
+        self._write_line("GroupingExpr")
+        with self._child_level():
+            self._write_line("expr", last=True)
+            with self._child_level(single=True):
+                expr.expr.accept(self)
+
+    def visit_literal_expr(self, expr: m.LiteralExpr) -> None:
+        self._write_line("LiteralExpr")
+        with self._child_level():
+            self._write_line(f"value: {expr.value}", last=True)
+
+    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> None:
+        self._write_line("WildcardExpr")
+
+    def visit_named_type(self, type: m.NamedType) -> None:
+        self._write_line("NamedType")
+        with self._child_level():
+            self._write_line(f'name: "{type.name.lexeme}"', last=True)
+
+    def visit_generic_type(self, type: m.GenericType) -> None:
+        self._write_line("GenericType")
+        with self._child_level():
+            self._write_line("type")
+            with self._child_level():
+                type.type.accept(self)
+            self._write_line("args", last=True)
+            with self._child_level():
+                for i, param in enumerate(type.args):
+                    self._idx = i
+                    if i == len(type.args) - 1:
+                        self._mark_last()
+                    param.accept(self)
+
+    def visit_constraint_type(self, type: m.ConstraintType) -> None:
+        self._write_line("ConstraintType")
+        with self._child_level():
+            self._write_line("type")
+            with self._child_level(single=True):
+                type.type.accept(self)
+            self._write_line("constraint", last=True)
+            with self._child_level(single=True):
+                type.constraint.accept(self)
+
+    def visit_complex_type(self, type: m.ComplexType) -> None:
+        self._write_line("ComplexType")
+        with self._child_level():
+            self._write_line("members", last=True)
+            with self._child_level():
+                for i, member in enumerate(type.members):
+                    self._idx = i
+                    if i == len(type.members) - 1:
+                        self._mark_last()
+                    member.accept(self)
+
+    def visit_extension_type(self, type: m.ExtensionType) -> None:
+        self._write_line("ExtensionType")
+        with self._child_level():
+            self._write_line("base")
+            with self._child_level(single=True):
+                type.base.accept(self)
+            self._write_line("extension", last=True)
+            with self._child_level(single=True):
+                type.extension.accept(self)
+
+    def visit_function_type(self, type: m.FunctionType) -> None:
+        self._write_line("FunctionType")
+        with self._child_level():
+            self._write_line("pos_args")
+            with self._child_level():
+                for i, arg in enumerate(type.pos_args):
+                    self._idx = i
+                    if i == len(type.pos_args) - 1:
+                        self._mark_last()
+                    self._print_function_arg(arg)
+
+            self._write_line("args")
+            with self._child_level():
+                for i, arg in enumerate(type.args):
+                    self._idx = i
+                    if i == len(type.args) - 1:
+                        self._mark_last()
+                    self._print_function_arg(arg)
+
+            self._write_line("kw_args")
+            with self._child_level():
+                for i, arg in enumerate(type.kw_args):
+                    self._idx = i
+                    if i == len(type.kw_args) - 1:
+                        self._mark_last()
+                    self._print_function_arg(arg)
+
+            self._write_line("returns", last=True)
+            with self._child_level(single=True):
+                type.returns.accept(self)
+
+    def _print_function_arg(self, arg: m.FunctionType.Argument) -> None:
+        self._write_line("Argument")
+        with self._child_level():
+            name: str = "None"
+            if arg.name is not None:
+                name = f'"{arg.name.lexeme}"'
+            self._write_line(f"name: {name}")
+            self._write_line("type")
+            with self._child_level(single=True):
+                arg.type.accept(self)
+            self._write_line(f"required: {arg.required}", last=True)
+
+
+class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str]):
+    def __init__(self, indent: int = 4):
+        self.indent: int = indent
+        self.level: int = 0
+
+    def indented(self, text: str) -> str:
+        return " " * (self.level * self.indent) + text
+
+    def print(self, expr: m.Expr | m.Stmt | m.Type) -> str:
+        self.level = 0
+        return expr.accept(self)
+
+    def visit_type_stmt(self, stmt: m.TypeStmt) -> str:
+        template: str = ""
+        if len(stmt.params) != 0:
+            params: list[str] = [self._print_type_param(param) for param in stmt.params]
+            template = f"[{', '.join(params)}]"
+        res: str = f"type {stmt.name.lexeme}{template} = {stmt.type.accept(self)}"
+        return self.indented(res)
+
+    def _print_type_param(self, param: m.TypeParam) -> str:
+        res: str = param.name.lexeme
+        if param.bound is not None:
+            res += "<:" + param.bound.accept(self)
+        return res
+
+    def visit_member_stmt(self, stmt: m.MemberStmt):
+        keyword: str = {
+            m.MemberKind.PROPERTY: "prop",
+            m.MemberKind.METHOD: "def",
+        }.get(stmt.kind, "")
+        res: str = f"{keyword} {stmt.name.lexeme}: {stmt.type.accept(self)}"
+        return self.indented(res)
+
+    def visit_extend_stmt(self, stmt: m.ExtendStmt):
+        template: str = ""
+        if len(stmt.params) != 0:
+            params: list[str] = [self._print_type_param(param) for param in stmt.params]
+            template = f"[{', '.join(params)}]"
+        res: str = self.indented(f"extend {stmt.name.lexeme}{template}")
+        res += " {\n"
+        self.level += 1
+        for member in stmt.members:
+            res += member.accept(self) + "\n"
+        self.level -= 1
+        res += self.indented("}")
+        return res
+
+    def visit_predicate_stmt(self, stmt: m.PredicateStmt):
+        name: str = stmt.name.lexeme
+        subject: str = stmt.subject.lexeme
+        type: str = stmt.type.accept(self)
+        condition: str = stmt.condition.accept(self)
+        return self.indented(f"predicate {name}({subject}: {type}) = {condition}")
+
+    def visit_logical_expr(self, expr: m.LogicalExpr):
+        left: str = expr.left.accept(self)
+        operator: str = expr.operator.lexeme
+        right: str = expr.right.accept(self)
+        return f"{left} {operator} {right}"
+
+    def visit_binary_expr(self, expr: m.BinaryExpr):
+        left: str = expr.left.accept(self)
+        operator: str = expr.operator.lexeme
+        right: str = expr.right.accept(self)
+        return f"{left} {operator} {right}"
+
+    def visit_unary_expr(self, expr: m.UnaryExpr):
+        operator: str = expr.operator.lexeme
+        right: str = expr.right.accept(self)
+        return f"{operator}{right}"
+
+    def visit_get_expr(self, expr: m.GetExpr):
+        expr_: str = expr.expr.accept(self)
+        name: str = expr.name.lexeme
+        return f"{expr_}.{name}"
+
+    def visit_variable_expr(self, expr: m.VariableExpr):
+        return expr.name.lexeme
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr):
+        expr_: str = expr.expr.accept(self)
+        return f"({expr_})"
+
+    def visit_literal_expr(self, expr: m.LiteralExpr):
+        return str(expr.value)
+
+    def visit_wildcard_expr(self, expr: m.WildcardExpr):
+        return "_"
+
+    def visit_named_type(self, type: m.NamedType) -> str:
+        return type.name.lexeme
+
+    def visit_generic_type(self, type: m.GenericType) -> str:
+        res: str = type.type.accept(self)
+        if len(type.args) != 0:
+            args: list[str] = [param.accept(self) for param in type.args]
+            res += f"[{', '.join(args)}]"
+        return res
+
+    def visit_constraint_type(self, type: m.ConstraintType) -> str:
+        res: str = type.type.accept(self)
+        res += " where " + type.constraint.accept(self)
+        return res
+
+    def visit_complex_type(self, type: m.ComplexType) -> str:
+        res: str = "{\n"
+        self.level += 1
+        for member in type.members:
+            res += member.accept(self)
+            res += "\n"
+        self.level -= 1
+        res += self.indented("}")
+        return res
+
+    def visit_extension_type(self, type: m.ExtensionType) -> str:
+        return f"{type.base.accept(self)} & {type.extension.accept(self)}"
+
+    def visit_function_type(self, type: m.FunctionType) -> str:
+        pos_args: list[str] = [self._print_arg(arg) for arg in type.pos_args]
+        mixed_args: list[str] = [self._print_arg(arg) for arg in type.args]
+        kw_args: list[str] = [self._print_arg(arg) for arg in type.kw_args]
+        args: list[str] = pos_args
+
+        if len(pos_args) != 0:
+            args.append("/")
+        args += mixed_args
+        if len(kw_args) != 0:
+            args.append("*")
+        args += kw_args
+
+        return f"fn ({', '.join(args)}) -> {type.returns.accept(self)}"
+
+    def _print_arg(self, arg: m.FunctionType.Argument) -> str:
+        res: str = ""
+        if arg.name is not None:
+            res += arg.name.lexeme
+            res += ": "
+        res += arg.type.accept(self)
+        if not arg.required:
+            res += "?"
+        return res
+
+
+class PythonAstPrinter(
+    AstPrinter,
+    p.MidasType.Visitor[None],
+    p.Stmt.Visitor[None],
+    p.Expr.Visitor[None],
+):
+    def visit_base_type(self, node: p.BaseType) -> None:
+        self._write_line("BaseType")
+        with self._child_level():
+            self._write_line(f"base: {node.base}")
+            self._write_optional_child("param", node.param, last=True)
+
+    def visit_constraint_type(self, node: p.ConstraintType) -> None:
+        self._write_line("ConstraintType")
+        with self._child_level():
+            self._write_line("type")
+            with self._child_level(single=True):
+                node.type.accept(self)
+            self._write_line(f"constraint: {ast.unparse(node.constraint)}", last=True)
+
+    def visit_frame_column(self, node: p.FrameColumn) -> None:
+        self._write_line("FrameColumn")
+        with self._child_level():
+            self._write_line(f"name: {node.name}")
+            self._write_optional_child("type", node.type, last=True)
+
+    def visit_frame_type(self, node: p.FrameType) -> None:
+        self._write_line("FrameType")
+        with self._child_level():
+            self._write_line("columns", last=True)
+            with self._child_level():
+                for i, col in enumerate(node.columns):
+                    self._idx = i
+                    if i == len(node.columns) - 1:
+                        self._mark_last()
+                    col.accept(self)
+
+    def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
+        stmt.expr.accept(self)
+
+    def visit_function(self, stmt: p.Function) -> None:
+        self._write_line("Function")
+        with self._child_level():
+            self._write_line(f"name: {stmt.name}")
+
+            self._write_line("posonlyargs")
+            with self._child_level():
+                for i, arg in enumerate(stmt.posonlyargs):
+                    self._idx = i
+                    if i == len(stmt.posonlyargs) - 1:
+                        self._mark_last()
+                    self._print_argument(arg)
+
+            self._write_line("args")
+            with self._child_level():
+                for i, arg in enumerate(stmt.args):
+                    self._idx = i
+                    if i == len(stmt.args) - 1:
+                        self._mark_last()
+                    self._print_argument(arg)
+
+            self._write_line("kwonlyargs")
+            with self._child_level():
+                for i, arg in enumerate(stmt.kwonlyargs):
+                    self._idx = i
+                    if i == len(stmt.kwonlyargs) - 1:
+                        self._mark_last()
+                    self._print_argument(arg)
+
+            self._write_optional_child("returns", stmt.returns)
+            self._write_line("body", last=True)
+            with self._child_level():
+                for i, body_stmt in enumerate(stmt.body):
+                    self._idx = i
+                    if i == len(stmt.body) - 1:
+                        self._mark_last()
+                    body_stmt.accept(self)
+
+    def _print_argument(self, arg: p.Function.Argument) -> None:
+        self._write_line("FunctionArgument")
+        with self._child_level():
+            self._write_line(f"name: {arg.name}")
+            self._write_optional_child("type", arg.type, last=True)
+
+    def visit_type_assign(self, stmt: p.TypeAssign) -> None:
+        self._write_line("TypeAssign")
+        with self._child_level():
+            self._write_line(f"name: {stmt.name}")
+            self._write_line("type", last=True)
+            with self._child_level(single=True):
+                stmt.type.accept(self)
+
+    def visit_assign_stmt(self, stmt: p.AssignStmt) -> None:
+        self._write_line("AssignStmt")
+        with self._child_level():
+            self._write_line("targets")
+            with self._child_level():
+                for i, target in enumerate(stmt.targets):
+                    self._idx = i
+                    if i == len(stmt.targets) - 1:
+                        self._mark_last()
+                    target.accept(self)
+            self._write_line("value", last=True)
+            with self._child_level(single=True):
+                stmt.value.accept(self)
+
+    def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
+        self._write_line("ReturnStmt")
+        with self._child_level():
+            self._write_optional_child("value", stmt.value, last=True)
+
+    def visit_if_stmt(self, stmt: p.IfStmt) -> None:
+        self._write_line("IfStmt")
+        with self._child_level():
+            self._write_line("test")
+            with self._child_level(single=True):
+                stmt.test.accept(self)
+            self._write_line("body")
+            with self._child_level():
+                for i, body_stmt in enumerate(stmt.body):
+                    self._idx = i
+                    if i == len(stmt.body) - 1:
+                        self._mark_last()
+                    body_stmt.accept(self)
+            self._write_line("orelse", last=True)
+            with self._child_level():
+                for i, else_stmt in enumerate(stmt.orelse):
+                    self._idx = i
+                    if i == len(stmt.orelse) - 1:
+                        self._mark_last()
+                    else_stmt.accept(self)
+
+    def visit_pass(self, stmt: p.Pass) -> None:
+        self._write_line("Pass")
+
+    def visit_for_stmt(self, stmt: p.ForStmt) -> None:
+        self._write_line("ForStmt")
+        with self._child_level():
+            self._write_line("target")
+            with self._child_level(single=True):
+                stmt.target.accept(self)
+            self._write_line("iterator")
+            with self._child_level(single=True):
+                stmt.iterator.accept(self)
+            self._write_line("body", last=True)
+            with self._child_level():
+                for i, body_stmt in enumerate(stmt.body):
+                    self._idx = i
+                    if i == len(stmt.body) - 1:
+                        self._mark_last()
+                    body_stmt.accept(self)
+
+    def visit_raw_stmt(self, stmt: p.RawStmt) -> None:
+        self._write_line("RawStmt")
+        with self._child_level(single=True):
+            self._write_line(f"stmt: {ast.unparse(stmt.stmt)}")
+
+    def visit_binary_expr(self, expr: p.BinaryExpr) -> None:
+        self._write_line("BinaryExpr")
+        with self._child_level():
+            self._write_line("left")
+            with self._child_level(single=True):
+                expr.left.accept(self)
+
+            self._write_line(f"operator: {expr.operator.__class__.__name__}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_compare_expr(self, expr: p.CompareExpr) -> None:
+        self._write_line("CompareExpr")
+        with self._child_level():
+            self._write_line("left")
+            with self._child_level(single=True):
+                expr.left.accept(self)
+
+            self._write_line(f"operator: {expr.operator.__class__.__name__}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_unary_expr(self, expr: p.UnaryExpr) -> None:
+        self._write_line("UnaryExpr")
+        with self._child_level():
+            self._write_line(f"operator: {expr.operator.__class__.__name__}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_call_expr(self, expr: p.CallExpr) -> None:
+        self._write_line("CallExpr")
+        with self._child_level():
+            self._write_line("callee")
+            with self._child_level(single=True):
+                expr.callee.accept(self)
+
+            self._write_line("arguments")
+            with self._child_level():
+                for i, arg in enumerate(expr.arguments):
+                    self._idx = i
+                    if i == len(expr.arguments) - 1:
+                        self._mark_last()
+                    arg.accept(self)
+
+            self._write_line("keywords", last=True)
+            with self._child_level():
+                for i, (name, arg) in enumerate(expr.keywords.items()):
+                    self._idx = i
+                    if i == len(expr.keywords) - 1:
+                        self._mark_last()
+                    self._write_line(name)
+                    with self._child_level(single=True):
+                        arg.accept(self)
+
+    def visit_get_expr(self, expr: p.GetExpr) -> None:
+        self._write_line("GetExpr")
+        with self._child_level():
+            self._write_line("object")
+            with self._child_level(single=True):
+                expr.object.accept(self)
+            self._write_line(f"name: {expr.name}", last=True)
+
+    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!r}")
+
+    def visit_variable_expr(self, expr: p.VariableExpr) -> None:
+        self._write_line("VariableExpr")
+        with self._child_level(single=True):
+            self._write_line(f"name: {expr.name}")
+
+    def visit_logical_expr(self, expr: p.LogicalExpr) -> None:
+        self._write_line("LogicalExpr")
+        with self._child_level():
+            self._write_line("left")
+            with self._child_level(single=True):
+                expr.left.accept(self)
+
+            self._write_line(f"operator: {expr.operator.__class__.__name__}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_cast_expr(self, expr: p.CastExpr) -> None:
+        self._write_line("CastExpr")
+        with self._child_level():
+            self._write_line("type")
+            with self._child_level(single=True):
+                expr.type.accept(self)
+            self._write_line("expr", last=True)
+            with self._child_level(single=True):
+                expr.expr.accept(self)
+
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> None:
+        self._write_line("TernaryExpr")
+        with self._child_level():
+            self._write_line("test")
+            with self._child_level(single=True):
+                expr.test.accept(self)
+
+            self._write_line("if_true")
+            with self._child_level(single=True):
+                expr.if_true.accept(self)
+
+            self._write_line("if_false", last=True)
+            with self._child_level(single=True):
+                expr.if_false.accept(self)
+
+    def visit_list_expr(self, expr: p.ListExpr) -> None:
+        self._write_line("ListExpr")
+        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_dict_expr(self, expr: p.DictExpr) -> None:
+        self._write_line("DictExpr")
+        with self._child_level():
+            self._write_line("keys")
+            with self._child_level():
+                for i, key in enumerate(expr.keys):
+                    self._idx = i
+                    if i == len(expr.keys) - 1:
+                        self._mark_last()
+                    if key is None:
+                        self._write_line("None")
+                    else:
+                        key.accept(self)
+            self._write_line("values", last=True)
+            with self._child_level():
+                for i, value in enumerate(expr.values):
+                    self._idx = i
+                    if i == len(expr.values) - 1:
+                        self._mark_last()
+                    value.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)
+
+    def visit_slice_expr(self, expr: p.SliceExpr) -> None:
+        self._write_line("SliceExpr")
+        with self._child_level():
+            self._write_optional_child("lower", expr.lower)
+            self._write_optional_child("upper", expr.upper)
+            self._write_optional_child("step", expr.step, last=True)
+
+    def visit_raw_expr(self, expr: p.RawExpr) -> None:
+        self._write_line("RawExpr")
+        with self._child_level(single=True):
+            self._write_line(f"expr: {ast.unparse(expr.expr)}")

midas/ast/python.py

@@ -0,0 +1,409 @@
+"""
+This file was generated by a script. Any manual changes might be overwritten.
+Please modify gen/python.py instead and run gen/gen.py
+"""
+
+from __future__ import annotations
+
+import ast
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Generic, Optional, TypeVar
+
+from midas.ast.location import Location
+
+T = TypeVar("T")
+
+
+####################
+# Type annotations #
+####################
+
+
+@dataclass(frozen=True, kw_only=True)
+class MidasType(ABC):
+    location: Location
+
+    @abstractmethod
+    def accept(self, visitor: Visitor[T]) -> T: ...
+
+    class Visitor(ABC, Generic[T]):
+        @abstractmethod
+        def visit_base_type(self, node: BaseType) -> T: ...
+
+        @abstractmethod
+        def visit_constraint_type(self, node: ConstraintType) -> T: ...
+
+        @abstractmethod
+        def visit_frame_column(self, node: FrameColumn) -> T: ...
+
+        @abstractmethod
+        def visit_frame_type(self, node: FrameType) -> T: ...
+
+
+@dataclass(frozen=True)
+class BaseType(MidasType):
+    base: str
+    param: Optional[MidasType]
+
+    def accept(self, visitor: MidasType.Visitor[T]) -> T:
+        return visitor.visit_base_type(self)
+
+
+@dataclass(frozen=True)
+class ConstraintType(MidasType):
+    type: MidasType
+    constraint: ast.expr
+
+    def accept(self, visitor: MidasType.Visitor[T]) -> T:
+        return visitor.visit_constraint_type(self)
+
+
+@dataclass(frozen=True)
+class FrameColumn(MidasType):
+    name: Optional[str]
+    type: Optional[MidasType]
+
+    def accept(self, visitor: MidasType.Visitor[T]) -> T:
+        return visitor.visit_frame_column(self)
+
+
+@dataclass(frozen=True)
+class FrameType(MidasType):
+    columns: list[FrameColumn]
+
+    def accept(self, visitor: MidasType.Visitor[T]) -> T:
+        return visitor.visit_frame_type(self)
+
+
+##############
+# Statements #
+##############
+
+
+@dataclass(frozen=True, kw_only=True)
+class Stmt(ABC):
+    location: Location
+
+    @abstractmethod
+    def accept(self, visitor: Visitor[T]) -> T: ...
+
+    class Visitor(ABC, Generic[T]):
+        @abstractmethod
+        def visit_expression_stmt(self, stmt: ExpressionStmt) -> T: ...
+
+        @abstractmethod
+        def visit_function(self, stmt: Function) -> T: ...
+
+        @abstractmethod
+        def visit_type_assign(self, stmt: TypeAssign) -> T: ...
+
+        @abstractmethod
+        def visit_assign_stmt(self, stmt: AssignStmt) -> T: ...
+
+        @abstractmethod
+        def visit_return_stmt(self, stmt: ReturnStmt) -> T: ...
+
+        @abstractmethod
+        def visit_if_stmt(self, stmt: IfStmt) -> T: ...
+
+        @abstractmethod
+        def visit_pass(self, stmt: Pass) -> T: ...
+
+        @abstractmethod
+        def visit_for_stmt(self, stmt: ForStmt) -> T: ...
+
+        @abstractmethod
+        def visit_raw_stmt(self, stmt: RawStmt) -> T: ...
+
+
+@dataclass(frozen=True)
+class ExpressionStmt(Stmt):
+    expr: Expr
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_expression_stmt(self)
+
+
+@dataclass(frozen=True)
+class Function(Stmt):
+    name: str
+    posonlyargs: list[Argument]
+    args: list[Argument]
+    sink: Optional[Argument]
+    kwonlyargs: list[Argument]
+    kw_sink: Optional[Argument]
+    returns: Optional[MidasType]
+    body: list[Stmt]
+
+    @dataclass(frozen=True, kw_only=True)
+    class Argument:
+        location: Optional[Location] = None
+        name: str
+        type: Optional[MidasType]
+        default: Optional[Expr]
+
+    @property
+    def all_args(self) -> list[Argument]:
+        return self.posonlyargs + self.args + self.kwonlyargs
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_function(self)
+
+
+@dataclass(frozen=True)
+class TypeAssign(Stmt):
+    name: str
+    type: MidasType
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_type_assign(self)
+
+
+@dataclass(frozen=True)
+class AssignStmt(Stmt):
+    targets: list[Expr]
+    value: Expr
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_assign_stmt(self)
+
+
+@dataclass(frozen=True)
+class ReturnStmt(Stmt):
+    value: Optional[Expr]
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_return_stmt(self)
+
+
+@dataclass(frozen=True)
+class IfStmt(Stmt):
+    test: Expr
+    body: list[Stmt]
+    orelse: list[Stmt]
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_if_stmt(self)
+
+
+@dataclass(frozen=True)
+class Pass(Stmt):
+    pass
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_pass(self)
+
+
+@dataclass(frozen=True)
+class ForStmt(Stmt):
+    target: Expr
+    iterator: Expr
+    body: list[Stmt]
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_for_stmt(self)
+
+
+@dataclass(frozen=True)
+class RawStmt(Stmt):
+    stmt: ast.stmt
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_raw_stmt(self)
+
+
+###############
+# Expressions #
+###############
+
+
+@dataclass(frozen=True, kw_only=True)
+class Expr(ABC):
+    location: Location
+
+    @abstractmethod
+    def accept(self, visitor: Visitor[T]) -> T: ...
+
+    class Visitor(ABC, Generic[T]):
+        @abstractmethod
+        def visit_binary_expr(self, expr: BinaryExpr) -> T: ...
+
+        @abstractmethod
+        def visit_compare_expr(self, expr: CompareExpr) -> T: ...
+
+        @abstractmethod
+        def visit_unary_expr(self, expr: UnaryExpr) -> T: ...
+
+        @abstractmethod
+        def visit_call_expr(self, expr: CallExpr) -> T: ...
+
+        @abstractmethod
+        def visit_get_expr(self, expr: GetExpr) -> T: ...
+
+        @abstractmethod
+        def visit_literal_expr(self, expr: LiteralExpr) -> T: ...
+
+        @abstractmethod
+        def visit_variable_expr(self, expr: VariableExpr) -> T: ...
+
+        @abstractmethod
+        def visit_logical_expr(self, expr: LogicalExpr) -> T: ...
+
+        @abstractmethod
+        def visit_cast_expr(self, expr: CastExpr) -> T: ...
+
+        @abstractmethod
+        def visit_ternary_expr(self, expr: TernaryExpr) -> T: ...
+
+        @abstractmethod
+        def visit_list_expr(self, expr: ListExpr) -> T: ...
+
+        @abstractmethod
+        def visit_dict_expr(self, expr: DictExpr) -> T: ...
+
+        @abstractmethod
+        def visit_subscript_expr(self, expr: SubscriptExpr) -> T: ...
+
+        @abstractmethod
+        def visit_slice_expr(self, expr: SliceExpr) -> T: ...
+
+        @abstractmethod
+        def visit_raw_expr(self, expr: RawExpr) -> T: ...
+
+
+@dataclass(frozen=True)
+class BinaryExpr(Expr):
+    left: Expr
+    operator: ast.operator
+    right: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_binary_expr(self)
+
+
+@dataclass(frozen=True)
+class CompareExpr(Expr):
+    left: Expr
+    operator: ast.cmpop
+    right: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_compare_expr(self)
+
+
+@dataclass(frozen=True)
+class UnaryExpr(Expr):
+    operator: ast.unaryop
+    right: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_unary_expr(self)
+
+
+@dataclass(frozen=True)
+class CallExpr(Expr):
+    callee: Expr
+    arguments: list[Expr]
+    keywords: dict[str, Expr]
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_call_expr(self)
+
+
+@dataclass(frozen=True)
+class GetExpr(Expr):
+    object: Expr
+    name: str
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_get_expr(self)
+
+
+@dataclass(frozen=True)
+class LiteralExpr(Expr):
+    value: Any
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_literal_expr(self)
+
+
+@dataclass(frozen=True)
+class VariableExpr(Expr):
+    name: str
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_variable_expr(self)
+
+
+@dataclass(frozen=True)
+class LogicalExpr(Expr):
+    left: Expr
+    operator: ast.boolop
+    right: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_logical_expr(self)
+
+
+@dataclass(frozen=True)
+class CastExpr(Expr):
+    type: MidasType
+    expr: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_cast_expr(self)
+
+
+@dataclass(frozen=True)
+class TernaryExpr(Expr):
+    test: Expr
+    if_true: Expr
+    if_false: Expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_ternary_expr(self)
+
+
+@dataclass(frozen=True)
+class ListExpr(Expr):
+    items: list[Expr]
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_list_expr(self)
+
+
+@dataclass(frozen=True)
+class DictExpr(Expr):
+    keys: list[Optional[Expr]]
+    values: list[Expr]
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_dict_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)
+
+
+@dataclass(frozen=True)
+class SliceExpr(Expr):
+    lower: Optional[Expr]
+    upper: Optional[Expr]
+    step: Optional[Expr]
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_slice_expr(self)
+
+
+@dataclass(frozen=True)
+class RawExpr(Expr):
+    expr: ast.expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_raw_expr(self)

midas/checker/builtins.midas

@@ -0,0 +1,181 @@
+extend float {
+    def hex: fn() -> str
+    def is_integer: fn() -> bool
+    prop real: float
+    prop imag: float
+    def conjugate: fn() -> float
+    def __add__: fn(value: float, /) -> float
+    def __sub__: fn(value: float, /) -> float
+    def __mul__: fn(value: float, /) -> float
+    def __floordiv__: fn(value: float, /) -> float
+    def __truediv__: fn(value: float, /) -> float
+    def __mod__: fn(value: float, /) -> float
+    // def __divmod__: fn(value: float, /) -> tuple[float, float]
+    
+    def __pow__: fn(value: int, /) -> float
+    // positive __value -> float; negative __value -> complex
+    // return type must be Any as `float | complex` causes too many false-positive errors
+    def __pow__: fn(value: float, /) -> Any
+    def __radd__: fn(value: float, /) -> float
+    def __rsub__: fn(value: float, /) -> float
+    def __rmul__: fn(value: float, /) -> float
+    def __rfloordiv__: fn(value: float, /) -> float
+    def __rtruediv__: fn(value: float, /) -> float
+    def __rmod__: fn(value: float, /) -> float
+    // def __rdivmod__: fn(value: float, /) -> tuple[float, float]
+    // def __rpow__: fn(value: _PositiveInteger, mod: None = None, /) -> float
+    // def __rpow__: fn(value: _NegativeInteger, mod: None = None, /) -> complex
+    // Returning `complex` for the general case gives too many false-positive errors.
+    // def __rpow__: fn(value: float, mod: None = None, /) -> Any
+    // def __getnewargs__: fn() -> tuple[float]
+    def __trunc__: fn() -> int
+    def __ceil__: fn() -> int
+    def __floor__: fn() -> int
+    def __round__: fn(ndigits: None?, /) -> int
+    def __round__: fn(ndigits: int, /) -> float
+    def __eq__: fn(value: object, /) -> bool
+    def __ne__: fn(value: object, /) -> bool
+    def __lt__: fn(value: float, /) -> bool
+    def __le__: fn(value: float, /) -> bool
+    def __gt__: fn(value: float, /) -> bool
+    def __ge__: fn(value: float, /) -> bool
+    def __neg__: fn() -> float
+    def __pos__: fn() -> float
+    def __int__: fn() -> int
+    def __float__: fn() -> float
+    def __abs__: fn() -> float
+    def __hash__: fn() -> int
+    def __bool__: fn() -> bool
+    def __format__: fn(format_spec: str, /) -> str
+}
+
+extend int {
+    prop real: int
+    prop imag: int
+    prop numerator: int
+    prop denominator: int
+    def conjugate: fn() -> int
+    def bit_length: fn() -> int
+    def bit_count: fn() -> int
+    // def to_bytes: fn(length: int?, byteorder: str?, *, signed: bool?) -> bytes
+
+    def __add__: fn(value: int, /) -> int
+    def __sub__: fn(value: int, /) -> int
+    def __mul__: fn(value: int, /) -> int
+    def __floordiv__: fn(value: int, /) -> int
+    def __truediv__: fn(value: int, /) -> float
+    def __mod__: fn(value: int, /) -> int
+    // def __divmod__: fn(value: int, /) -> tuple[int, int]
+    def __radd__: fn(value: int, /) -> int
+    def __rsub__: fn(value: int, /) -> int
+    def __rmul__: fn(value: int, /) -> int
+    def __rfloordiv__: fn(value: int, /) -> int
+    def __rtruediv__: fn(value: int, /) -> float
+    def __rmod__: fn(value: int, /) -> int
+    // def __rdivmod__: fn(value: int, /) -> tuple[int, int]
+    def __pow__: fn(value: int, /) -> int
+    // def __pow__: fn(value: _PositiveInteger, mod: None = None, /) -> int
+    // def __pow__: fn(value: _NegativeInteger, mod: None = None, /) -> float
+    // positive __value -> int; negative __value -> float
+    // return type must be Any as `int | float` causes too many false-positive errors
+    // def __pow__: fn(value: int, mod: None = None, /) -> Any
+    // def __pow__: fn(value: int, mod: int, /) -> int
+    def __rpow__: fn(value: int, /) -> Any
+    def __and__: fn(value: int, /) -> int
+    def __or__: fn(value: int, /) -> int
+    def __xor__: fn(value: int, /) -> int
+    def __lshift__: fn(value: int, /) -> int
+    def __rshift__: fn(value: int, /) -> int
+    def __rand__: fn(value: int, /) -> int
+    def __ror__: fn(value: int, /) -> int
+    def __rxor__: fn(value: int, /) -> int
+    def __rlshift__: fn(value: int, /) -> int
+    def __rrshift__: fn(value: int, /) -> int
+    def __neg__: fn() -> int
+    def __pos__: fn() -> int
+    def __invert__: fn() -> int
+    def __trunc__: fn() -> int
+    def __ceil__: fn() -> int
+    def __floor__: fn() -> int
+    def __round__: fn(ndigits: None?, /) -> int
+    def __round__: fn(ndigits: int, /) -> int
+
+    // def __getnewargs__: fn() -> tuple[int]
+    def __eq__: fn(value: object, /) -> bool
+    def __ne__: fn(value: object, /) -> bool
+    def __lt__: fn(value: int, /) -> bool
+    def __le__: fn(value: int, /) -> bool
+    def __gt__: fn(value: int, /) -> bool
+    def __ge__: fn(value: int, /) -> bool
+    def __float__: fn() -> float
+    def __int__: fn() -> int
+    def __abs__: fn() -> int
+    def __hash__: fn() -> int
+    def __bool__: fn() -> bool
+    def __index__: fn() -> int
+    def __format__: fn(format_spec: str, /) -> str
+}
+
+extend list[T] {
+    def copy: fn () -> list[T]
+    def append: fn (object: T, /) -> None
+    def extend: fn (iterable: list[T], /) -> None
+    def pop: fn (index: int?, /) -> T
+    def index: fn (value: T, start: int?, stop: int?, /) -> int
+    def count: fn (value: T, /) -> int
+    def insert: fn (index: int, object: T, /) -> None
+    def remove: fn (value: T, /) -> None
+    def sort: fn (*, reverse: bool?) -> None
+    def __len__: fn () -> int
+    // def __iter__: fn () -> Iterator[T]
+    def __getitem__: fn (i: int, /) -> T
+    def __getitem__: fn (s: slice, /) -> list[T]
+    def __setitem__: fn (key: int, value: T, /) -> None
+    def __setitem__: fn (key: slice, value: list[T], /) -> None
+    def __delitem__: fn (key: int, /) -> None
+    def __delitem__: fn (key: slice, /) -> None
+    // def __add__: fn[S <: T] (value: list[S], /) -> list[T]
+    def __add__: fn (value: list[T], /) -> list[T]
+    def __iadd__: fn (value: list[T], /) -> list[T]
+    def __mul__: fn (value: int, /) -> list[T]
+    def __rmul__: fn (value: int, /) -> list[T]
+    def __imul__: fn (value: int, /) -> list[T]
+    def __contains__: fn (key: object, /) -> bool
+    // def __reversed__: fn (self) -> Iterator[_T]
+    def __gt__: fn (value: list[T], /) -> bool
+    def __ge__: fn (value: list[T], /) -> bool
+    def __lt__: fn (value: list[T], /) -> bool
+    def __le__: fn (value: list[T], /) -> bool
+    def __eq__: fn (value: object, /) -> bool
+
+    prop __doc__: str
+}
+
+extend dict[K, V] {
+    def copy: fn() -> dict[K, V]
+    def keys: fn() -> list[K]  // TODO: use builtin types
+    def values: fn() -> list[V]  // TODO: use builtin types
+    // def items: fn() -> list[tuple[K, V]]  // TODO: use builtin types
+    
+    // def get: fn(key: K, default: None = None, /) -> V | None
+    def get: fn(key: K, default: V, /) -> V
+    // def get: fn[T](key: K, default: T, /) -> V | T
+    def pop: fn(key: K, /) -> V
+    def pop: fn(key: K, default: V, /) -> V
+    // def pop: fn[T](key: K, default: T, /) -> V | T
+    def __len__: fn() -> int
+    def __getitem__: fn(key: K, /) -> V
+    def __setitem__: fn(key: K, value: V, /) -> None
+    def __delitem__: fn(key: K, /) -> None
+    // def __iter__: fn() -> Iterator[K]
+    def __eq__: fn(value: object, /) -> bool
+    // def __reversed__: fn() -> Iterator[K]
+    
+    def __or__: fn(value: dict[K, V], /) -> dict[K, V]
+    // def __or__: fn[K2, V2](value: dict[K2, V2], /) -> dict[K | K2, V | V2]
+    def __ror__: fn(value: dict[K, V], /) -> dict[K, V]
+    // def __ror__: fn[K2, V2](value: dict[K2, V2], /) -> dict[K | K2, V | V2]
+    // def __ior__: fn(value: SupportsKeysAndGetItem[K, V], /) -> dict[K, V]
+    // def __ior__: fn(value: Iterable[tuple[K, V]], /) -> dict[K, V]
+
+}

midas/checker/builtins.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from midas.checker.types import (
+    BaseType,
+    GenericType,
+    TopType,
+    TypeVar,
+    UnitType,
+)
+
+if TYPE_CHECKING:
+    from midas.checker.registry import TypesRegistry
+
+
+BUILTIN_SUBTYPES: dict[str, set[str]] = {
+    "float": {"int"},
+    "int": {"bool"},
+}
+
+
+def define_builtins(reg: TypesRegistry):
+    """Define builtin types and operations"""
+    any = reg.define_type("Any", TopType())
+    unit = reg.define_type("None", UnitType())
+    object = reg.define_type("object", BaseType(name="object"))
+    bool = reg.define_type("bool", BaseType(name="bool"))
+    int = reg.define_type("int", BaseType(name="int"))
+    float = reg.define_type("float", BaseType(name="float"))
+    str = reg.define_type("str", BaseType(name="str"))
+    slice = reg.define_type("slice", BaseType(name="slice"))
+
+    list = reg.define_type(
+        "list",
+        GenericType(
+            name="list",
+            params=[TypeVar(name="T", bound=None)],
+            body=BaseType(name="list"),
+        ),
+    )
+    dict = reg.define_type(
+        "dict",
+        GenericType(
+            name="dict",
+            params=[
+                TypeVar(name="K", bound=None),
+                TypeVar(name="V", bound=None),
+            ],
+            body=BaseType(name="dict"),
+        ),
+    )

midas/checker/checker.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+from typing import Optional
+
+from midas.checker.diagnostic import Diagnostic
+from midas.checker.midas import MidasTyper
+from midas.checker.python import PythonTyper
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import Reporter
+from midas.utils import TypedAST
+
+
+class TypeChecker:
+    def __init__(self):
+        self.types: TypesRegistry = TypesRegistry()
+        self.reporter: Reporter = Reporter()
+
+        self.midas_typer = MidasTyper(self.types, self.reporter)
+        self.python_typer = PythonTyper(self.types, self.reporter)
+
+    def import_midas(self, path: Path):
+        source: str = path.read_text()
+        return self.import_midas_source(source, path=str(path))
+
+    def import_midas_source(self, source: str, path: Optional[str] = None):
+        self.midas_typer.process(source, path)
+
+    def type_check(self, path: Path) -> TypedAST:
+        source: str = path.read_text()
+        return self.type_check_source(source, path=str(path))
+
+    def type_check_source(self, source: str, path: Optional[str] = None) -> TypedAST:
+        return self.python_typer.process(source, path)
+
+    @property
+    def diagnostics(self) -> list[Diagnostic]:
+        return self.reporter.diagnostics

midas/checker/diagnostic.py

@@ -0,0 +1,42 @@
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import Optional
+
+from midas.ast.location import Location
+
+
+class DiagnosticType(StrEnum):
+    ERROR = "Error"
+    WARNING = "Warning"
+    INFO = "Info"
+
+
+@dataclass(frozen=True)
+class Diagnostic:
+    file_path: Optional[str]
+    location: Location
+    type: DiagnosticType
+    message: str
+
+    @property
+    def location_str(self) -> str:
+        start_loc: str = f"L{self.location.lineno}:{self.location.col_offset+1}"
+        end_loc: Optional[str] = ""
+        if (
+            self.location.end_lineno is not None
+            and self.location.end_col_offset is not None
+        ):
+            end_loc = f"L{self.location.end_lineno}:{self.location.end_col_offset+1}"
+
+        loc: str = ""
+        if self.file_path is not None:
+            loc += f" in {self.file_path}"
+        if end_loc is None:
+            loc += f" at {start_loc}"
+        else:
+            loc += f" from {start_loc} to {end_loc}"
+
+        return f"{self.type}{loc}"
+
+    def __str__(self) -> str:
+        return f"{self.location_str}: {self.message}"

midas/checker/environment.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from midas.checker.types import Type
+
+
+class Environment:
+    """
+    A scoped environment in which variables are defined
+
+    Each environment can inherit from a parent/enclosing environment.
+    """
+
+    def __init__(self, enclosing: Optional[Environment] = None) -> None:
+        self.enclosing: Optional[Environment] = enclosing
+        self.values: dict[str, Type] = {}
+        self.return_types: list[Type] = []
+
+        self._children: list[Environment] = []
+        if enclosing is not None:
+            enclosing._children.append(self)
+
+    def define(self, name: str, value: Type) -> None:
+        """Define a variable in this environment
+
+        Args:
+            name (str): the name of the variable
+            value (Type): the value
+        """
+        self.values[name] = value
+
+    def get(self, name: str) -> Optional[Type]:
+        """Get a variable in the closest environment which has a definition for it
+
+        Args:
+            name (str): the name of the variable
+
+        Returns:
+            Optional[Type]: the value of the variable, or None if it was not found
+        """
+        if name in self.values:
+            return self.values[name]
+        if self.enclosing is not None:
+            return self.enclosing.get(name)
+        # raise NameError(f"Undefined variable '{name}'")
+        return None
+
+    def assign(self, name: str, value: Type) -> bool:
+        """Assign a new value to a variable in the environment it was defined in
+
+        Args:
+            name (str): the name of the variable
+            value (Type): the new value
+
+        Returns:
+            bool: True if the variable was assigned in this environment or an ancestor, False otherwise
+        """
+        if name not in self.values:
+            if self.enclosing is None:
+                return False
+            if self.enclosing.assign(name, value):
+                return True
+        self.values[name] = value
+        return True
+
+    def clear(self):
+        """Clear all definitions in this environment"""
+        self.values = {}
+
+    def get_at(self, distance: int, name: str) -> Optional[Type]:
+        """Get the value of a variable at a given distance
+
+        A distance of 0 looks up in this environment, 1 in the parent environment, etc.
+        This methods expects `distance` to be valid. An error will be raised if
+        the stack does not extend far enough to reach `distance`
+
+        Args:
+            distance (int): the scope distance
+            name (str): the name of the variable
+
+        Returns:
+            Optional[Type]: the value at the given distance, or None if it is not defined in that environment
+
+        Raises:
+            AssertionError: if the stack does not extend far enough to reach `distance`
+        """
+        return self.ancestor(distance).values.get(name)
+
+    def assign_at(self, distance: int, name: str, value: Type) -> None:
+        """Assign a new value to a variable at a given distance
+
+        A distance of 0 assigns in this environment, 1 in the parent environment, etc.
+
+        Args:
+            distance (int): the scope distance
+            name (str): the name of the variable
+            value (Type): the new value
+
+        Raises:
+            AssertionError: if the stack does not extend far enough to reach `distance`
+        """
+        self.ancestor(distance).values[name] = value
+
+    def ancestor(self, distance: int) -> Environment:
+        """Get the ancestor at a given distance
+
+        A distance of 0 references this environment, 1 the parent environment, etc.
+
+        Args:
+            distance (int): the scope distance
+
+        Returns:
+            Environment: the environment
+
+        Raises:
+            AssertionError: if the stack does not extend far enough to reach `distance`
+        """
+        env: Environment = self
+        for _ in range(distance):
+            assert env.enclosing is not None
+            env = env.enclosing
+        return env
+
+    def flat_dict(self) -> dict[str, Type]:
+        """Get the current environment including definitions in its ancestor as a flat dictionary
+
+        This method recursively combines this environment definitions with its ancestor's
+
+        Returns:
+            dict: the combined environment
+        """
+        if self.enclosing is None:
+            return self.values
+        return self.enclosing.flat_dict() | self.values
+
+    def dump(self) -> dict:
+        return {
+            "values": self.values,
+            "return_types": self.return_types,
+            "children": [child.dump() for child in self._children],
+        }

midas/checker/midas.py

@@ -0,0 +1,206 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+import midas.ast.midas as m
+from midas.checker.builtins import define_builtins
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter, Reporter
+from midas.checker.types import (
+    AliasType,
+    ComplexType,
+    ExtensionType,
+    Function,
+    GenericType,
+    Type,
+    TypeVar,
+    UnknownType,
+)
+from midas.lexer.midas import MidasLexer
+from midas.lexer.token import Token
+from midas.parser.midas import MidasParser
+
+
+class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[Type]):
+    """A resolver which evaluates Midas type definitions and build a registry"""
+
+    def __init__(self, types: TypesRegistry, reporter: Reporter) -> None:
+        self.logger: logging.Logger = logging.getLogger("MidasTyper")
+        self.reporter: FileReporter = reporter.for_file(None)
+
+        self.types: TypesRegistry = types
+        self._local_variables: dict[str, TypeVar] = {}
+
+        self._current_name: Optional[str] = None
+
+        define_builtins(self.types)
+        builtins_path: Path = (Path(__file__).parent / "builtins.midas").resolve()
+        self.process(builtins_path.read_text(), str(builtins_path))
+
+    def process(self, source: str, path: Optional[str]):
+        self.reporter = self.reporter.for_file(path)
+        lexer: MidasLexer = MidasLexer(source)
+        tokens: list[Token] = lexer.process()
+        parser: MidasParser = MidasParser(tokens)
+        stmts: list[m.Stmt] = parser.parse()
+        for error in parser.errors:
+            self.reporter.error(error.token.get_location(), error.message)
+        self.resolve(stmts)
+
+    def get_type(self, name: str) -> Type:
+        """Get a type from its name
+
+        Args:
+            name (str): the name of the type
+
+        Raises:
+            NameError: if the type is not defined
+
+        Returns:
+            Type: the type
+        """
+        if name in self._local_variables:
+            return self._local_variables[name]
+        return self.types.get_type(name)
+
+    def resolve(self, stmts: list[m.Stmt]):
+        """Process a sequence of statements
+
+        Args:
+            stmts (list[m.Stmt]): the statements
+        """
+        for stmt in stmts:
+            stmt.accept(self)
+
+    def visit_type_stmt(self, stmt: m.TypeStmt) -> None:
+        name: str = stmt.name.lexeme
+        self._current_name = name
+        params: list[TypeVar] = self._resolve_type_params(stmt.params)
+
+        type: Type = stmt.type.accept(self)
+        if len(params) != 0:
+            type = GenericType(name=name, params=params, body=type)
+        else:
+            type = AliasType(name=name, type=type)
+        self.types.define_type(name, type)
+        self._local_variables.clear()
+        self._current_name = None
+
+    def visit_member_stmt(self, stmt: m.MemberStmt) -> None: ...
+
+    def visit_extend_stmt(self, stmt: m.ExtendStmt) -> None:
+        self._resolve_type_params(stmt.params)
+        base_name: str = stmt.name.lexeme
+        try:
+            _ = self.get_type(base_name)
+        except NameError:
+            self.reporter.error(stmt.name.get_location(), f"Unknown type '{base_name}'")
+
+        for member in stmt.members:
+            member_type: Type = member.type.accept(self)
+            self.types.define_member(
+                base_name,
+                member.name.lexeme,
+                member_type,
+                member.kind == m.MemberKind.METHOD,
+            )
+
+    def visit_predicate_stmt(self, stmt: m.PredicateStmt) -> None:
+        self.reporter.warning(stmt.location, "PredicateStmt not yet supported")
+
+    def visit_logical_expr(self, expr: m.LogicalExpr) -> None:
+        self.reporter.warning(expr.location, "LogicalExpr not yet supported")
+
+    def visit_binary_expr(self, expr: m.BinaryExpr) -> None:
+        self.reporter.warning(expr.location, "BinaryExpr not yet supported")
+
+    def visit_unary_expr(self, expr: m.UnaryExpr) -> None:
+        self.reporter.warning(expr.location, "UnaryExpr not yet supported")
+
+    def visit_get_expr(self, expr: m.GetExpr) -> None:
+        self.reporter.warning(expr.location, "GetExpr not yet supported")
+
+    def visit_variable_expr(self, expr: m.VariableExpr) -> None:
+        self.reporter.warning(expr.location, "VariableExpr not yet supported")
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr) -> None:
+        return expr.expr.accept(self)
+
+    def visit_literal_expr(self, expr: m.LiteralExpr) -> None:
+        self.reporter.warning(expr.location, "LiteralExpr not yet supported")
+
+    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> None:
+        self.reporter.warning(expr.location, "WildcardExpr not yet supported")
+
+    def visit_named_type(self, type: m.NamedType) -> Type:
+        name: str = type.name.lexeme
+        try:
+            return self.get_type(name)
+        except NameError:
+            msg: str = f"Undefined type {name}"
+            if self._current_name == name:
+                msg += ". Recursive types are not supported, use an extend block"
+            self.reporter.error(type.name.get_location(), msg)
+            return UnknownType()
+
+    def visit_generic_type(self, type: m.GenericType) -> Type:
+        type_: Type = type.type.accept(self)
+        args: list[Type] = [arg.accept(self) for arg in type.args]
+        try:
+            return self.types.apply_generic(type_, args)
+        except Exception as e:
+            self.reporter.error(type.location, f"Cannot apply generic type: {e}")
+            return UnknownType()
+
+    def visit_constraint_type(self, type: m.ConstraintType) -> Type:
+        type_: Type = type.type.accept(self)
+        type.constraint.accept(self)
+        # TODO
+        return UnknownType()
+
+    def visit_complex_type(self, type: m.ComplexType) -> ComplexType:
+        return ComplexType(
+            members={
+                member.name.lexeme: member.type.accept(self) for member in type.members
+            }
+        )
+
+    def visit_extension_type(self, type: m.ExtensionType) -> Type:
+        return ExtensionType(
+            base=type.base.accept(self),
+            extension=self.visit_complex_type(type.extension),
+        )
+
+    def visit_function_type(self, type: m.FunctionType) -> Type:
+        n_pos_args: int = len(type.pos_args)
+        n_args: int = len(type.args)
+
+        def process_arg(arg: m.FunctionType.Argument, i: int) -> Function.Argument:
+            return Function.Argument(
+                pos=i,
+                name=arg.name.lexeme if arg.name is not None else str(i),
+                type=arg.type.accept(self),
+                required=arg.required,
+            )
+
+        return Function(
+            pos_args=[process_arg(arg, i) for i, arg in enumerate(type.pos_args)],
+            args=[process_arg(arg, i + n_pos_args) for i, arg in enumerate(type.args)],
+            kw_args=[
+                process_arg(arg, i + n_pos_args + n_args)
+                for i, arg in enumerate(type.kw_args)
+            ],
+            returns=type.returns.accept(self),
+        )
+
+    def _resolve_type_params(self, params: list[m.TypeParam]):
+        vars: list[TypeVar] = []
+        for param in params:
+            name: str = param.name.lexeme
+            bound: Optional[Type] = None
+            if param.bound is not None:
+                bound = param.bound.accept(self)
+            var = TypeVar(name=name, bound=bound)
+            self._local_variables[name] = var
+            vars.append(var)
+        return vars

midas/checker/operators.py

@@ -0,0 +1,38 @@
+import ast
+from typing import Type
+
+OPERATOR_METHODS: dict[Type[ast.operator], str] = {
+    ast.Add: "__add__",
+    ast.Sub: "__sub__",
+    ast.Mult: "__mul__",
+    ast.MatMult: "__matmul__",
+    ast.Div: "__truediv__",
+    ast.Mod: "__mod__",
+    ast.Pow: "__pow__",
+    ast.LShift: "__lshift__",
+    ast.RShift: "__rshift__",
+    ast.BitOr: "__or__",
+    ast.BitXor: "__xor__",
+    ast.BitAnd: "__and__",
+    ast.FloorDiv: "__floordiv__",
+}
+
+COMPARATOR_METHODS: dict[Type[ast.cmpop], str] = {
+    ast.Eq: "__eq__",
+    # ast.NotEq: "__noteq__",
+    ast.Lt: "__lt__",
+    ast.LtE: "__le__",
+    ast.Gt: "__gt__",
+    ast.GtE: "__ge__",
+    # ast.Is: "__is__",
+    # ast.IsNot: "__isnot__",
+    # ast.In: "__in__",
+    # ast.NotIn: "__notin__",
+}
+
+UNARY_METHODS: dict[Type[ast.unaryop], str] = {
+    ast.Invert: "__invert__",
+    # ast.Not: "",
+    ast.UAdd: "__pos__",
+    ast.USub: "__neg__",
+}

midas/checker/preamble.py

@@ -0,0 +1,121 @@
+from dataclasses import dataclass
+
+from midas.checker.environment import Environment
+from midas.checker.registry import TypesRegistry
+from midas.checker.types import Function, GenericType, TopType, Type, TypeVar, UnitType
+
+
+@dataclass(frozen=True)
+class Param:
+    name: str
+    type: Type
+    required: bool = True
+
+
+class Preamble(Environment):
+    def __init__(self, types: TypesRegistry) -> None:
+        super().__init__()
+        self._types: TypesRegistry = types
+
+        self._def_type_constructor("object")
+        self._def_type_constructor("float")
+        self._def_type_constructor("int")
+        self._def_type_constructor("bool")
+        self._def_type_constructor("str")
+        self._def_function(
+            name="list",
+            pos=[Param("object", TopType())],
+            returns=self._list_of(TopType()),
+        )
+
+        # TODO: use sink
+        self._def_function(
+            name="print",
+            pos=[Param("object", TopType())],
+            returns=UnitType(),
+        )
+
+        map_in = TypeVar(name="T", bound=None)
+        map_out = TypeVar(name="U", bound=None)
+        mapper = self._make_function(
+            name="MapTransform",
+            pos=[Param("v", map_in)],
+            returns=map_out,
+        )
+        self._def_function(
+            name="map",
+            pos=[
+                Param("transform", mapper),
+                Param(
+                    "iterable",
+                    self._list_of(map_in),  # TODO: replace with Iterable[T]
+                ),
+            ],
+            returns=self._list_of(map_out),  # TODO: replace with Iterable[U]
+        )
+
+    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):
+        # TODO: more specific arg types
+        self._def_function(
+            name=name,
+            pos=[Param("object", TopType(), required=False)],
+            returns=self._types.get_type(name),
+        )
+
+    def _make_function(
+        self,
+        *,
+        name: str,
+        pos: list[Param] = [],
+        mixed: list[Param] = [],
+        kw: list[Param] = [],
+        returns: Type = UnitType(),
+        type_vars: list[TypeVar] = [],
+    ) -> Type:
+        def map_args(params: list[Param], offset: int) -> list[Function.Argument]:
+            return [
+                Function.Argument(
+                    pos=i + offset,
+                    name=param.name,
+                    type=param.type,
+                    required=param.required,
+                )
+                for i, param in enumerate(params)
+            ]
+
+        function = Function(
+            pos_args=map_args(pos, 0),
+            args=map_args(mixed, len(pos)),
+            kw_args=map_args(kw, len(pos) + len(mixed)),
+            returns=returns,
+        )
+        if len(type_vars) != 0:
+            function = GenericType(
+                name=name,
+                params=type_vars,
+                body=function,
+            )
+        return function
+
+    def _def_function(
+        self,
+        *,
+        name: str,
+        pos: list[Param] = [],
+        mixed: list[Param] = [],
+        kw: list[Param] = [],
+        returns: Type = UnitType(),
+        type_vars: list[TypeVar] = [],
+    ):
+        function: Type = self._make_function(
+            name=name,
+            pos=pos,
+            mixed=mixed,
+            kw=kw,
+            returns=returns,
+            type_vars=type_vars,
+        )
+        self.define(name, function)

midas/checker/python.py

@@ -0,0 +1,998 @@
+import ast
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.environment import Environment
+from midas.checker.operators import COMPARATOR_METHODS, OPERATOR_METHODS, UNARY_METHODS
+from midas.checker.preamble import Preamble
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter, Reporter
+from midas.checker.resolver import Resolver
+from midas.checker.types import (
+    AppliedType,
+    Function,
+    OverloadedFunction,
+    Type,
+    UnitType,
+    UnknownType,
+    unfold_type,
+)
+from midas.parser.python import PythonParser
+from midas.utils import TypedAST
+
+TypedExpr = tuple[p.Expr, Type]
+
+
+class ReturnException(Exception):
+    pass
+
+
+@dataclass(frozen=True, kw_only=True)
+class MappedArgument:
+    expr: p.Expr
+    type: Type
+    argument: Function.Argument
+
+
+@dataclass(frozen=True, kw_only=True)
+class OverloadCandidate:
+    function: Function
+    mapped: list[MappedArgument]
+
+
+class PythonTyper(
+    p.Stmt.Visitor[None],
+    p.Expr.Visitor[Type],
+    p.MidasType.Visitor[Type],
+):
+    """A type checker which can use custom type definitions"""
+
+    def __init__(
+        self,
+        types: TypesRegistry,
+        reporter: Reporter,
+    ):
+        self.logger: logging.Logger = logging.getLogger("PythonTyper")
+        self.reporter: FileReporter = reporter.for_file(None)
+        self.types: TypesRegistry = 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]] = []
+
+    def process(self, source: str, path: Optional[str]) -> TypedAST:
+        self.reporter = self.reporter.for_file(path)
+
+        tree: ast.Module = ast.parse(source, filename=path or "<unknown>")
+        parser = PythonParser()
+        stmts: list[p.Stmt] = parser.parse_module(tree)
+        resolver = Resolver()
+        resolver.resolve(*stmts)
+
+        self.env = self.global_env
+        self.locals = resolver.locals
+        self.judgements = []
+
+        self.check(stmts)
+
+        return TypedAST(stmts=stmts, judgements=self.judgements)
+
+    def judge(self, expr: p.Expr, type: Type):
+        """Record a typing judgement
+
+        Args:
+            expr (p.Expr): the judged expression
+            type (Type): the type of the expression
+        """
+        self.judgements.append((expr, type))
+
+    def compute_type(self, expr: p.Expr) -> Type:
+        """Evaluate the type of an expression
+
+        Args:
+            expr (p.Expr): the expression to type
+
+        Returns:
+            Type: the type of the given expression
+        """
+        return expr.accept(self)
+
+    def type_of(self, expr: p.Expr) -> Type:
+        """Evaluate the type of an expression and record the judgement
+
+        Args:
+            expr (p.Expr): the expression to evaluate
+
+        Returns:
+            Type: the type of the given expression
+        """
+        type: Type = self.compute_type(expr)
+        self.judge(expr, type)
+        return type
+
+    def resolve_type_expr(self, expr: p.MidasType) -> Type:
+        return expr.accept(self)
+
+    def process_stmt(self, stmt: p.Stmt) -> None:
+        stmt.accept(self)
+
+    def process_block(self, block: list[p.Stmt], env: Environment) -> bool:
+        """Evaluate a sequence of statements
+
+        Args:
+            block (list[p.Stmt]): the statements to evaluate
+            env (Environment): the environment in which to evaluate
+
+        Returns:
+            bool: whether a return statement is present in the block
+        """
+        previous_env: Environment = self.env
+        self.env = env
+        returned: bool = False
+        for i, stmt in enumerate(block):
+            try:
+                self.process_stmt(stmt)
+            except ReturnException:
+                returned = True
+                if i < len(block) - 1:
+                    self.reporter.warning(
+                        block[i + 1].location, "Unreachable statement"
+                    )
+                break
+        self.env = previous_env
+        return returned
+
+    def check(self, statements: list[p.Stmt]) -> None:
+        """Type check a sequence of statements and returns diagnostics
+
+        Args:
+            statements (list[p.Stmt]): the statements to evaluate and check
+        """
+        for stmt in statements:
+            self.process_stmt(stmt)
+
+        self.logger.debug(f"Final environment: {self.env.flat_dict()}")
+
+    def look_up_variable(self, name: str, expr: p.Expr) -> Optional[Type]:
+        """Look up a variable in the environment it was declared
+
+        Args:
+            name (str): the name of the variable
+            expr (p.Expr): the variable expression, used to lookup the scope distance
+
+        Returns:
+            Optional[Type]: the type of the variable, or None if it was not found
+        """
+        distance: Optional[int] = self.locals.get(expr)
+        if distance is not None:
+            return self.env.get_at(distance, name)
+        return self.global_env.get(name)
+
+    def is_subtype(self, type1: Type, type2: Type) -> bool:
+        return self.types.is_subtype(type1, type2)
+
+    def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
+        self.type_of(stmt.expr)
+
+    def visit_function(self, stmt: p.Function) -> None:
+        env: Environment = Environment(self.env)
+        pos_args: list[Function.Argument] = []
+        args: list[Function.Argument] = []
+        kw_args: list[Function.Argument] = []
+
+        def eval_arg_type(arg: p.Function.Argument) -> Type:
+            if arg.type is not None:
+                return self.resolve_type_expr(arg.type)
+            if arg.default is not None:
+                return self.type_of(arg.default)
+            return UnknownType()
+
+        pos: int = 0
+        for arg in stmt.posonlyargs:
+            pos_args.append(
+                Function.Argument(
+                    pos=pos,
+                    name=arg.name,
+                    type=eval_arg_type(arg),
+                    required=arg.default is None,
+                )
+            )
+            pos += 1
+        for arg in stmt.args:
+            args.append(
+                Function.Argument(
+                    pos=pos,
+                    name=arg.name,
+                    type=eval_arg_type(arg),
+                    required=arg.default is None,
+                )
+            )
+            pos += 1
+        for arg in stmt.kwonlyargs:
+            kw_args.append(
+                Function.Argument(
+                    pos=pos,  # not relevant
+                    name=arg.name,
+                    type=eval_arg_type(arg),
+                    required=arg.default is None,
+                )
+            )
+            pos += 1
+
+        for arg in pos_args + args + kw_args:
+            env.define(arg.name, arg.type)
+
+        returns_hint: Optional[Type] = None
+        if stmt.returns is not None:
+            returns_hint = self.resolve_type_expr(stmt.returns)
+            # Early define to handle simple fully-typed recursion
+            inside_function: Function = Function(
+                pos_args=pos_args,
+                args=args,
+                kw_args=kw_args,
+                returns=returns_hint,
+            )
+            self.env.define(stmt.name, inside_function)
+
+        returned: bool = self.process_block(stmt.body, env)
+        inferred_return: Type = UnknownType()
+        if not returned:
+            env.return_types.append(UnitType())
+        return_types: list[Type] = self.types.reduce_types(env.return_types)
+        if len(return_types) == 1:
+            inferred_return = return_types[0]
+        elif len(return_types) > 1:
+            self.reporter.error(
+                stmt.location,
+                f"Mixed return types: {return_types}",
+            )
+
+        returns: Type = UnknownType()
+        if returns_hint is not None:
+            assert stmt.returns is not None
+            returns = returns_hint
+            if not self.is_subtype(inferred_return, returns):
+                self.reporter.error(
+                    stmt.returns.location,
+                    f"Return type mismatch, annotated {returns} but returns {inferred_return}",
+                )
+        else:
+            returns = inferred_return
+
+        # TODO: handle *args and **kwargs sinks
+        function: Function = Function(
+            pos_args=pos_args,
+            args=args,
+            kw_args=kw_args,
+            returns=returns,
+        )
+        self.env.define(stmt.name, function)
+
+    def visit_type_assign(self, stmt: p.TypeAssign) -> None:
+        # TODO check not yet defined locally
+        type: Type = self.resolve_type_expr(stmt.type)
+        self.env.define(stmt.name, type)
+
+    def visit_assign_stmt(self, stmt: p.AssignStmt) -> None:
+        value_type: Type = self.type_of(stmt.value)
+        for target in stmt.targets:
+            self._assign(stmt.location, target, value_type)
+
+    def _assign(self, location: Location, target: p.Expr, value_type: Type):
+        match target:
+            case p.VariableExpr():
+                self._assign_var(location, target, value_type)
+
+            case p.GetExpr(object=object, name=name):
+                self._assign_attr(location, object, name, value_type)
+
+            case _:
+                if not isinstance(target, p.VariableExpr):
+                    self.logger.warning(f"Unsupported assignment to {target}")
+                    self.reporter.warning(
+                        target.location, f"Unsupported assignment to {target}"
+                    )
+
+    def _assign_var(self, location: Location, target: p.VariableExpr, value_type: Type):
+        name: str = target.name
+        var_type: Optional[Type] = self.look_up_variable(name, target)
+
+        if var_type is None:
+            self.env.define(name, value_type)
+        else:
+            # S <: T
+            # Γ, x: T   v: S
+            # x = v
+            if not self.is_subtype(value_type, var_type):
+                self.reporter.error(
+                    location,
+                    f"Cannot assign {value_type} to variable '{name}' of type {var_type}",
+                )
+
+    def _assign_attr(
+        self, location: Location, object: p.Expr, name: str, value_type: Type
+    ):
+        object_type: Type = self.type_of(object)
+        member: Optional[Type] = self.types.lookup_member(object_type, name)
+        if member is None:
+            self.reporter.error(location, f"Unknown member '{name}' of {object_type}")
+            return
+        self.logger.debug(f"Member '{name}' of {object_type} has type {member}")
+        if not self.is_subtype(value_type, member):
+            self.reporter.error(
+                location,
+                f"Cannot assign {value_type} to member '{object_type}.{name}' of type {member}",
+            )
+
+    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)
+        raise ReturnException()
+
+    def visit_if_stmt(self, stmt: p.IfStmt) -> None:
+        # Not evaluated in sub-environment because assignments in the test leak out of the if
+        # For example:
+        # if (m := 1 + 1) < 2:
+        #     ...
+        # print(m)  # <- m is still defined
+        test_type: Type = self.type_of(stmt.test)
+
+        # TODO Allow subtypes or any type
+        if test_type != self.types.get_type("bool"):
+            self.reporter.error(
+                stmt.test.location, f"If test must be a boolean, got {test_type}"
+            )
+
+        env: Environment = Environment(self.env)
+        body_returned: bool = self.process_block(stmt.body, env)
+        else_returned: bool = self.process_block(stmt.orelse, env)
+        self.env.return_types.extend(env.return_types)
+        if body_returned and else_returned:
+            raise ReturnException()
+
+    def visit_pass(self, stmt: p.Pass) -> None:
+        pass
+
+    def visit_for_stmt(self, stmt: p.ForStmt) -> None:
+        item_type: Optional[Type] = self._get_iterator_type(stmt.iterator)
+        if item_type is None:
+            iterator_type: Type = self.compute_type(stmt.iterator)
+            self.reporter.error(
+                stmt.iterator.location, f"{iterator_type} is not iterable"
+            )
+            item_type = UnknownType()
+
+        self._assign(stmt.location, stmt.target, item_type)
+        self.judge(stmt.target, item_type)
+        env: Environment = Environment(self.env)
+        body_returned: bool = self.process_block(stmt.body, env)
+        if body_returned:
+            raise ReturnException()
+
+    def visit_raw_stmt(self, stmt: p.RawStmt) -> None:
+        pass
+
+    def visit_binary_expr(self, expr: p.BinaryExpr) -> Type:
+        method: Optional[str] = OPERATOR_METHODS.get(expr.operator.__class__)
+        if method is None:
+            self.logger.warning(f"Unsupported operator {expr.operator}")
+            self.reporter.warning(
+                expr.location, f"Unsupported operator {expr.operator}"
+            )
+            return UnknownType()
+
+        return self._visit_binary_expr(expr.location, expr.left, expr.right, method)
+
+    def visit_compare_expr(self, expr: p.CompareExpr) -> Type:
+        method: Optional[str] = COMPARATOR_METHODS.get(expr.operator.__class__)
+        if method is None:
+            self.logger.warning(f"Unsupported operator {expr.operator}")
+            self.reporter.warning(
+                expr.location, f"Unsupported operator {expr.operator}"
+            )
+            return UnknownType()
+
+        return self._visit_binary_expr(expr.location, expr.left, expr.right, method)
+
+    def _visit_binary_expr(
+        self, location: Location, left_expr: p.Expr, right_expr: p.Expr, method: str
+    ) -> Type:
+        left: Type = self.type_of(left_expr)
+        right: Type = self.type_of(right_expr)
+
+        operation: Optional[Type] = self.types.lookup_member(left, method)
+        if operation is None:
+            self.reporter.error(
+                location,
+                f"Undefined operation {method} between {left} and {right}",
+            )
+            return UnknownType()
+
+        result: Optional[Type] = self._get_call_result(
+            location,
+            operation,
+            [(right_expr, right)],
+            {},
+        )
+        return result or UnknownType()
+
+    def visit_unary_expr(self, expr: p.UnaryExpr) -> Type:
+        method: Optional[str] = UNARY_METHODS.get(expr.operator.__class__)
+        if method is None:
+            self.logger.warning(f"Unsupported operator {expr.operator}")
+            self.reporter.warning(
+                expr.location, f"Unsupported operator {expr.operator}"
+            )
+            return UnknownType()
+
+        operand: Type = self.type_of(expr.right)
+        operation: Optional[Type] = self.types.lookup_member(operand, method)
+        if operation is None:
+            self.reporter.error(
+                expr.location,
+                f"Undefined operation {method} for {operand}",
+            )
+            return UnknownType()
+
+        result: Optional[Type] = self._get_call_result(
+            expr.location,
+            operation,
+            [],
+            {},
+        )
+        return result or UnknownType()
+
+    def visit_call_expr(self, expr: p.CallExpr) -> Type:
+        callee: Type = self.type_of(expr.callee)
+        positional: list[TypedExpr] = [
+            (arg, self.type_of(arg)) for arg in expr.arguments
+        ]
+        keywords: dict[str, TypedExpr] = {
+            name: (arg, self.type_of(arg)) for name, arg in expr.keywords.items()
+        }
+        return (
+            self._get_call_result(
+                location=expr.location,
+                callee=callee,
+                positional=positional,
+                keywords=keywords,
+            )
+            or UnknownType()
+        )
+
+    def visit_get_expr(self, expr: p.GetExpr) -> Type:
+        object: Type = self.type_of(expr.object)
+        member: Optional[Type] = self.types.lookup_member(object, expr.name)
+        if member is None:
+            self.reporter.error(
+                expr.location, f"Unknown member '{expr.name}' of {object}"
+            )
+            return UnknownType()
+        self.logger.debug(f"Member '{expr.name}' of {object} has type {member}")
+        return member
+
+    def visit_literal_expr(self, expr: p.LiteralExpr) -> Type:
+        match expr.value:
+            case bool():  # Must be before int
+                return self.types.get_type("bool")
+            case int():
+                return self.types.get_type("int")
+            case float():
+                return self.types.get_type("float")
+            case str():
+                return self.types.get_type("str")
+            case _:
+                self.reporter.warning(expr.location, f"Unknown literal {expr}")
+                return UnknownType()
+
+    def visit_variable_expr(self, expr: p.VariableExpr) -> Type:
+        type: Optional[Type] = self.look_up_variable(expr.name, expr)
+        if type is None:
+            self.logger.debug(f"Unknown variable {expr.name} in {self.env.flat_dict()}")
+            self.reporter.warning(expr.location, "Unknown variable")
+        return type or UnknownType()
+
+    def visit_logical_expr(self, expr: p.LogicalExpr) -> Type:
+        left: Type = self.type_of(expr.left)
+        right: Type = self.type_of(expr.right)
+
+        if self.is_subtype(left, right):
+            return right
+        if self.is_subtype(right, left):
+            return left
+
+        self.reporter.error(
+            expr.location,
+            f"Incompatible operand types, {left=} and {right=}",
+        )
+        return UnknownType()
+
+    def visit_cast_expr(self, expr: p.CastExpr) -> Type:
+        return self.resolve_type_expr(expr.type)
+
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> Type:
+        test_type: Type = self.type_of(expr.test)
+
+        # TODO Allow subtypes or any type
+        if test_type != self.types.get_type("bool"):
+            self.reporter.error(
+                expr.test.location, f"If test must be a boolean, got {test_type}"
+            )
+
+        true_type: Type = self.type_of(expr.if_true)
+        false_type: Type = self.type_of(expr.if_false)
+        if self.is_subtype(true_type, false_type):
+            return false_type
+        if self.is_subtype(false_type, true_type):
+            return true_type
+
+        self.reporter.error(
+            expr.location,
+            f"Incompatible types in ternary if branches: true={true_type} and false={false_type}",
+        )
+        return UnknownType()
+
+    def visit_list_expr(self, expr: p.ListExpr) -> Type:
+        list_type: Type = self.types.get_type("list")
+        item_types: list[Type] = [self.type_of(item) for item in expr.items]
+        item_types = self.types.reduce_types(item_types)
+
+        if len(item_types) == 0:
+            return list_type
+
+        if len(item_types) == 1:
+            item_type: Type = item_types[0]
+            return self.types.apply_generic(list_type, [item_type])
+        self.reporter.error(
+            expr.location,
+            f"Heterogeneous list items: {item_types}",
+        )
+        return self.types.apply_generic(list_type, [UnknownType()])
+
+    def visit_dict_expr(self, expr: p.DictExpr) -> Type:
+        dict_type: Type = self.types.get_type("dict")
+
+        key_types: list[Type] = []
+        value_types: list[Type] = []
+        for key, value in zip(expr.keys, expr.values):
+            if key is None:
+                self.reporter.warning(
+                    value.location, "Dictionary unpacking not supported"
+                )
+                continue
+            key_types.append(self.type_of(key))
+            value_types.append(self.type_of(value))
+
+        key_types = self.types.reduce_types(key_types)
+        value_types = self.types.reduce_types(value_types)
+
+        if len(key_types) == 0 or len(value_types) == 0:
+            return dict_type
+
+        key_type: Type = UnknownType()
+        value_type: Type = UnknownType()
+
+        if len(key_types) == 1:
+            key_type = key_types[0]
+        else:
+            self.reporter.error(
+                expr.location,
+                f"Heterogeneous dict keys: {key_types}",
+            )
+
+        if len(value_types) == 1:
+            value_type = value_types[0]
+        else:
+            self.reporter.error(
+                expr.location,
+                f"Heterogeneous dict values: {value_types}",
+            )
+        return self.types.apply_generic(dict_type, [key_type, value_type])
+
+    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)
+        return (
+            self._get_call_result(expr.location, operation, [(expr.index, index)], {})
+            or UnknownType()
+        )
+
+    def visit_slice_expr(self, expr: p.SliceExpr) -> Type:
+        return self.types.get_type("slice")
+
+    def visit_raw_expr(self, expr: p.RawExpr) -> Type:
+        return UnknownType()
+
+    def visit_base_type(self, node: p.BaseType) -> Type:
+        base: Type
+        try:
+            base = self.types.get_type(node.base)
+        except NameError:
+            self.reporter.warning(node.location, f"Unknown type '{node.base}'")
+            return UnknownType()
+
+        if node.param is not None:
+            param: Type = self.resolve_type_expr(node.param)
+            return self.types.apply_generic(base, [param])
+        return base
+
+    def visit_constraint_type(self, node: p.ConstraintType) -> Type:
+        self.reporter.warning(node.location, "ConstraintType not yet supported")
+        return UnknownType()
+
+    def visit_frame_column(self, node: p.FrameColumn) -> Type:
+        self.reporter.warning(node.location, "FrameColumn not yet supported")
+        return UnknownType()
+
+    def visit_frame_type(self, node: p.FrameType) -> Type:
+        self.reporter.warning(node.location, "FrameType not yet supported")
+        return UnknownType()
+
+    def _get_call_result(
+        self,
+        location: Location,
+        callee: Type,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+        report_errors: bool = True,
+    ) -> Optional[Type]:
+        """Get the result type of a function call
+
+        If the function has overloads, the function will try to resolve the
+        appropriate signature.
+        Argument types are matched to the defined parameters.
+        The function doesn't take the raw expression as a parameter to accomodate
+        for desugared calls such as for operators.
+
+        Args:
+            location (Location): the call location
+            callee (Type): the called function
+            positional (list[TypedExpr]): the list positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+            report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
+
+        Returns:
+            Type: the return type of the call, or `None` if either
+            the call is invalid or no overload matched the arguments uniquely
+        """
+        match callee:
+            case Function() as function:
+                valid: bool
+                mapped: list[MappedArgument]
+                valid, mapped = self.map_call_arguments(
+                    function, location, positional, keywords
+                )
+                valid = valid and self._are_arguments_valid(mapped, report_errors)
+                if not valid:
+                    return None
+                return function.returns
+
+            case OverloadedFunction(overloads=overloads):
+                function = self._match_overload(
+                    overloads, location, positional, keywords, report_errors
+                )
+                if function is None:
+                    return None
+                return function.returns
+
+            case AppliedType(body=body):
+                return self._get_call_result(
+                    location, body, positional, keywords, report_errors
+                )
+
+            case UnknownType():
+                return UnknownType()
+
+            case _:
+                if report_errors:
+                    self.reporter.error(location, f"{callee} is not callable")
+                return None
+
+    def _are_arguments_valid(
+        self,
+        arguments: list[MappedArgument],
+        report_errors: bool = True,
+    ) -> bool:
+        """Check whether the passed argument types correspond to their matched parameter definitions
+
+        Args:
+            arguments (list[MappedArgument]): the list of argument/parameter pairs
+            report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
+
+        Returns:
+            bool: True if all arguments fit the matching parameter definitions, False otherwise
+        """
+        valid: bool = True
+        for arg in arguments:
+            if not self.is_subtype(arg.type, arg.argument.type):
+                if report_errors:
+                    self.reporter.error(
+                        arg.expr.location,
+                        f"Wrong type for argument '{arg.argument.name}', expected {arg.argument.type}, got {arg.type}",
+                    )
+                valid = False
+        return valid
+
+    def _match_overload(
+        self,
+        overloads: list[Type],
+        location: Location,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+        report_errors: bool = True,
+    ) -> Optional[Function]:
+        """Try and resolve the appropriate overload for the given arguments
+
+        Args:
+            overloads (list[Type]): the list of possible overloads
+            location (Location): the call location
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keywords arguments
+            report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
+
+        Returns:
+            Optional[Function]: the resolved function signature if it can be
+            determined unambigously, or `None`.
+        """
+        candidates: list[OverloadCandidate] = []
+        for overload in overloads:
+            function: Type = unfold_type(overload)
+            if not isinstance(function, Function):
+                if report_errors:
+                    self.logger.error(
+                        f"Overload is not a function: {overload} is {function}"
+                    )
+                continue
+            valid, mapped = self.map_call_arguments(
+                function=function,
+                location=location,
+                positional=positional,
+                keywords=keywords,
+                report_errors=False,
+            )
+            if valid and self._are_arguments_valid(mapped, report_errors=False):
+                candidates.append(
+                    OverloadCandidate(
+                        function=function,
+                        mapped=mapped,
+                    )
+                )
+
+        pos_types: str = ", ".join(str(type) for _, type in positional)
+        kw_types: str = ", ".join(
+            f"{name}: {type}" for name, (_, type) in keywords.items()
+        )
+        for_args: str = f"for arguments pos=[{pos_types}] and kw={{{kw_types}}}"
+
+        n_candidates: int = len(candidates)
+
+        # Exactly 1 match -> return it
+        if n_candidates == 1:
+            return candidates[0].function
+
+        # No match -> invalid call
+        if n_candidates == 0:
+            overloads_str: str = ", ".join(map(str, overloads))
+            if report_errors:
+                self.reporter.error(
+                    location,
+                    f"No matching overload in [{overloads_str}] {for_args}",
+                )
+            return None
+
+        # Multiple matches -> see if one <: all others (more specific)
+        for i1, c1 in enumerate(candidates):
+            mapped1: list[MappedArgument] = c1.mapped
+            best_match: bool = True
+            for i2, c2 in enumerate(candidates):
+                if i1 == i2:
+                    continue
+                mapped2: list[MappedArgument] = c2.mapped
+                if not self._are_mapped_subtypes(mapped1, mapped2):
+                    best_match = False
+                    break
+                self.logger.debug(f"{c1.function} is a full overload of {c2.function}")
+            if best_match:
+                return c1.function
+
+        candidates_str: str = ", ".join(
+            str(candidate.function) for candidate in candidates
+        )
+        if report_errors:
+            self.reporter.error(
+                location,
+                f"Multiple matching overloads {for_args}: {candidates_str}",
+            )
+        return None
+
+    def map_call_arguments(
+        self,
+        function: Function,
+        location: Location,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+        report_errors: bool = True,
+    ) -> tuple[bool, list[MappedArgument]]:
+        """Map call arguments to a function's parameters as defined in its signature
+
+        This method maps positional-only, keyword-only and mixed parameter definitions
+        with the arguments passed at the call site
+
+        Any mismatched, missing or unexpected argument is reported as a diagnostic,
+        unless `report_errors` is set to `False`
+
+        Args:
+            function (Function): the function definition
+            location (Location): the call location
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+            report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
+
+        Returns:
+            tuple[bool, list[MappedArgument]]: a boolean reporting whether
+            the call is valid and the list of mapped arguments
+        """
+        set_args: set[str] = set()
+
+        required_positional: list[str] = [
+            arg.name for arg in function.pos_args + function.args if arg.required
+        ]
+        required_keyword: list[str] = [
+            arg.name for arg in function.kw_args if arg.required
+        ]
+
+        mapped: list[MappedArgument] = []
+
+        pos_params: list[Function.Argument] = list(function.pos_args)
+        mixed_params: list[Function.Argument] = list(function.args)
+        kw_params: dict[str, Function.Argument] = {
+            arg.name: arg for arg in function.kw_args
+        }
+
+        valid_call: bool = True
+
+        # TODO: handle *args and **kwargs sinks
+        for arg in positional:
+            param: Function.Argument
+            if len(pos_params) != 0:
+                param = pos_params.pop(0)
+            elif len(mixed_params) != 0:
+                param = mixed_params.pop(0)
+            else:
+                if report_errors:
+                    self.reporter.error(
+                        arg[0].location, "Too many positional arguments"
+                    )
+                valid_call = False
+                break
+            name: str = param.name
+            if name in required_positional:
+                required_positional.remove(name)
+            if name in required_keyword:
+                required_keyword.remove(name)
+            set_args.add(name)
+            mapped.append(
+                MappedArgument(
+                    expr=arg[0],
+                    type=arg[1],
+                    argument=param,
+                )
+            )
+
+        kw_params.update({arg.name: arg for arg in mixed_params})
+        for name, arg in keywords.items():
+            param: Function.Argument
+            if name not in kw_params:
+                if report_errors:
+                    if name in set_args:
+                        self.reporter.error(
+                            arg[0].location, f"Multiple values for argument '{name}'"
+                        )
+                    else:
+                        self.reporter.error(
+                            arg[0].location, f"Unknown keyword argument '{name}'"
+                        )
+                valid_call = False
+                continue
+            param = kw_params.pop(name)
+            if name in required_positional:
+                required_positional.remove(name)
+            if name in required_keyword:
+                required_keyword.remove(name)
+            set_args.add(name)
+            mapped.append(
+                MappedArgument(
+                    expr=arg[0],
+                    type=arg[1],
+                    argument=param,
+                )
+            )
+
+        def join_args(args: list[str]) -> str:
+            args = list(map(lambda a: f"'{a}'", args))
+            if len(args) == 0:
+                return ""
+            if len(args) == 1:
+                return args[0]
+            return ", ".join(args[:-1]) + " and " + args[-1]
+
+        if len(required_positional) != 0:
+            plural: str = "" if len(required_positional) == 1 else "s"
+            args: str = join_args(required_positional)
+            if report_errors:
+                self.reporter.error(
+                    location,
+                    f"Missing required positional argument{plural}: {args}",
+                )
+            valid_call = False
+
+        if len(required_keyword) != 0:
+            plural: str = "" if len(required_keyword) == 1 else "s"
+            args: str = join_args(required_keyword)
+            if report_errors:
+                self.reporter.error(
+                    location,
+                    f"Missing required keyword argument{plural}: {args}",
+                )
+            valid_call = False
+
+        return valid_call, mapped
+
+    def _are_mapped_subtypes(
+        self, mapped1: list[MappedArgument], mapped2: list[MappedArgument]
+    ) -> bool:
+        """Check whether the given argument mappings are subtype/supertype of one another
+
+        This function checks whether the argument mappings `mapped1` are subtypes
+        of `mapped2`. If any of the parameter type in `mapped1` is not a subtype
+        of the corresponding parameter in `mapped2`, `False` is returned.
+
+        This is used to check whether a given overload is
+        a more specific function/ a subtype of another.
+
+        Args:
+            mapped1 (list[MappedArgument]): the first argument mappings (subtype)
+            mapped2 (list[MappedArgument]): the second argument mappings (supertype)
+
+        Returns:
+            bool: `True` if `mapped1` is a subtype of `mapped2`, `False` otherwise
+        """
+        by_expr: dict[p.Expr, Type] = {}
+        for arg in mapped1:
+            by_expr[arg.expr] = arg.argument.type
+
+        for arg in mapped2:
+            type2: Type = arg.argument.type
+            type1: Type = by_expr[arg.expr]
+            if not self.is_subtype(type1, type2):
+                return False
+        return True
+
+    def _get_iterator_type(self, expr: p.Expr) -> Optional[Type]:
+        # TODO: lookup __iter__
+        type: Type = self.type_of(expr)
+        getitem: Optional[Type] = self.types.lookup_member(type, "__getitem__")
+        if getitem is None:
+            return None
+
+        index: p.Expr = p.LiteralExpr(location=expr.location, value=0)
+        index_type: Type = self.compute_type(index)
+        result: Optional[Type] = self._get_call_result(
+            location=expr.location,
+            callee=getitem,
+            positional=[(index, index_type)],
+            keywords={},
+            report_errors=False,
+        )
+        return result

midas/checker/registry.py

@@ -0,0 +1,347 @@
+import logging
+from typing import Optional
+
+from midas.checker.builtins import BUILTIN_SUBTYPES
+from midas.checker.types import (
+    AliasType,
+    AppliedType,
+    BaseType,
+    ComplexType,
+    ExtensionType,
+    Function,
+    GenericType,
+    OverloadedFunction,
+    TopType,
+    Type,
+    TypeVar,
+    UnknownType,
+    substitute_typevars,
+)
+
+
+class TypesRegistry:
+    def __init__(self) -> None:
+        self.logger: logging.Logger = logging.getLogger("TypesRegistry")
+        self._types: dict[str, Type] = {}
+        self._members: dict[str, dict[str, Type]] = {}
+
+    def get_type(self, name: str) -> Type:
+        """Get a type from its name
+
+        Args:
+            name (str): the name of the type
+
+        Raises:
+            NameError: if the type is not defined
+
+        Returns:
+            Type: the type
+        """
+        if name in self._types:
+            return self._types[name]
+        raise NameError(f"Undefined type {name}")
+
+    def define_type(self, name: str, type: Type) -> Type:
+        """Define a type in the registry
+
+        Args:
+            name (str): the name of the type
+            type (Type): the type to define
+
+        Raises:
+            ValueError: if a type is already defined with that name
+
+        Returns:
+            Type: the defined type
+        """
+        if name in self._types:
+            raise ValueError(f"Type {name} already defined")
+        self._types[name] = type
+        return type
+
+    def define_member(
+        self, type_name: str, member_name: str, member_type: Type, is_method: bool
+    ):
+        members: dict[str, Type] = self._members.setdefault(type_name, {})
+        if member_name in members:
+            if not is_method:
+                self.logger.error(
+                    f"Member '{member_name}' already defined for type {type_name}"
+                )
+                return
+            current: Type = members[member_name]
+            combined: Type
+            match current:
+                case OverloadedFunction(overloads=overloads):
+                    combined = OverloadedFunction(overloads=overloads + [member_type])
+                case _:
+                    combined = OverloadedFunction(overloads=[current, member_type])
+            members[member_name] = combined
+
+        else:
+            members[member_name] = member_type
+
+    def is_subtype(self, type1: Type, type2: Type) -> bool:
+        """Check whether `type1` is a subtype of `type2`
+
+        For more details on the rules checked here, see TAPL Chap. 15-16-17
+
+        Args:
+            type1 (Type): the potential subtype
+            type2 (Type): the potential supertype
+
+        Returns:
+            bool: whether `type1` is a subtype of `type2`
+        """
+
+        if type1 == type2:
+            return True
+
+        match (type1, type2):
+            case (_, TopType()):
+                return True
+
+            case (AliasType(type=base1), _):
+                return self.is_subtype(base1, type2)
+
+            case (BaseType(name=name1), BaseType(name=name2)):
+                return name1 in BUILTIN_SUBTYPES.get(name2, set())
+
+            case (ComplexType(properties=props1), ComplexType(properties=props2)):
+                for k, t in props2.items():
+                    if k not in props1:
+                        return False
+                    if not self.is_subtype(props1[k], t):
+                        return False
+                return True
+
+            case (Function(), Function()):
+                return self.is_func_subtype(type1, type2)
+
+            case (TypeVar(bound=bound), _):
+                if bound is None:
+                    return False
+                return self.is_subtype(bound, type2)
+
+        return False
+
+    # TODO: verify the logic in here
+    def is_func_subtype(self, func1: Function, func2: Function) -> bool:
+        """Check whether a function is a subtype of another
+
+        Args:
+            func1 (Function): the potential function subtype
+            func2 (Function): the potential function supertype
+
+        Returns:
+            bool: whether `func1` is a subtype of `func2`
+        """
+        if not self.is_subtype(func1.returns, func2.returns):
+            return False
+
+        pos1: list[Function.Argument] = func1.pos_args
+        mixed1: list[Function.Argument] = func1.args
+        kw1: dict[str, Function.Argument] = {a.name: a for a in func1.kw_args}
+        pos2: list[Function.Argument] = func2.pos_args
+        mixed2: list[Function.Argument] = func2.args
+        kw2: dict[str, Function.Argument] = {a.name: a for a in func2.kw_args}
+
+        mixed_by_pos: dict[int, Function.Argument] = {arg.pos: arg for arg in mixed2}
+        mixed_by_name: dict[str, Function.Argument] = {arg.name: arg for arg in mixed2}
+
+        def is_arg_subtype(sub: Function.Argument, sup: Function.Argument) -> bool:
+            if not self.is_subtype(sub.type, sup.type):
+                return False
+            if not sup.required and sub.required:
+                return False
+            return True
+
+        for arg1 in pos1:
+            arg2: Function.Argument
+            if arg1.pos < len(pos2):
+                arg2 = pos2[arg1.pos]
+            elif arg1.pos in mixed_by_pos:
+                arg2 = mixed_by_pos[arg1.pos]
+            elif not arg1.required:
+                continue
+            else:
+                return False
+            if not is_arg_subtype(arg2, arg1):
+                return False
+
+        for name, arg1 in kw1.items():
+            arg2: Function.Argument
+            if name in kw2:
+                arg2 = kw2[name]
+            elif name in mixed_by_name:
+                arg2 = mixed_by_name[name]
+            elif not arg1.required:
+                continue
+            else:
+                return False
+            if not is_arg_subtype(arg2, arg1):
+                return False
+
+        for arg1 in mixed1:
+            pos_arg2: Optional[Function.Argument] = None
+            kw_arg2: Optional[Function.Argument] = None
+            if arg1.name in kw2:
+                kw_arg2 = kw2[arg1.name]
+            elif arg1.name in mixed_by_name:
+                kw_arg2 = mixed_by_name[arg1.name]
+            if arg1.pos < len(pos2):
+                pos_arg2 = pos2[arg1.pos]
+            elif arg1.pos in mixed_by_pos:
+                pos_arg2 = mixed_by_pos[arg1.pos]
+
+            # No match in func2 and arg is required
+            if pos_arg2 is None and kw_arg2 is None and arg1.required:
+                return False
+
+            # Matching keyword argument
+            if kw_arg2 is not None and not is_arg_subtype(kw_arg2, arg1):
+                return False
+
+            # Matching positional argument
+            if pos_arg2 is not None and not is_arg_subtype(pos_arg2, arg1):
+                return False
+
+        mixed_positions: set[int] = {a.pos for a in mixed1}
+        mixed_names: set[str] = {a.name for a in mixed1}
+        for arg2 in pos2:
+            if not arg2.required:
+                continue
+            if arg2.pos >= len(pos1) and arg2.pos not in mixed_positions:
+                return False
+
+        for name, arg2 in kw2.items():
+            if not arg2.required:
+                continue
+            if name not in kw1 and name not in mixed_names:
+                return False
+
+        for arg2 in mixed2:
+            if arg2.required:
+                continue
+            pos_match: bool = arg2.pos < len(pos1) or arg2.pos in mixed_positions
+            kw_match: bool = arg2.name in kw1 or arg2.name in mixed_names
+            if not pos_match or not kw_match:
+                return False
+
+        return True
+
+    def apply_generic(self, type: Type, args: list[Type]) -> Type:
+        match type:
+            case AliasType(name=name, type=base):
+                return AliasType(name=name, type=self.apply_generic(base, args))
+
+            case GenericType(name=name, params=type_vars, body=body):
+                n_args: int = len(args)
+                n_type_vars: int = len(type_vars)
+                if n_args < n_type_vars:
+                    raise ValueError(
+                        f"Missing type arguments, expected {n_type_vars} but only {n_args} provided"
+                    )
+                if n_args > n_type_vars:
+                    raise ValueError(
+                        f"Too many type arguments, expected {n_type_vars} but {n_args} provided"
+                    )
+                substitutions: dict[str, Type] = {}
+                for arg, type_var in zip(args, type_vars):
+                    if type_var.bound is not None and not self.is_subtype(
+                        arg, type_var.bound
+                    ):
+                        raise ValueError(
+                            f"Type argument {arg} is not a subtype of {type_var.bound}"
+                        )
+                    substitutions[type_var.name] = arg
+                return AppliedType(
+                    name=name,
+                    args=args,
+                    body=substitute_typevars(body, substitutions),
+                )
+
+            case _:
+                raise ValueError(f"{type} is not a generic type")
+
+    def reduce_types(self, types: list[Type]) -> list[Type]:
+        """Reduce a list of types to remove subtypes and only keep the highest types
+
+        Args:
+            types (list[Type]): the types to reduce
+
+        Returns:
+            list[Type]: the reduced list of types
+        """
+
+        reduced: bool = True
+        keep: list[int] = list(range(len(types)))
+        while reduced:
+            reduced = False
+            for i, i1 in enumerate(keep):
+                type1: Type = types[i1]
+                for i2 in keep[i + 1 :]:
+                    type2 = types[i2]
+                    if self.is_subtype(type1, type2):
+                        keep.remove(i1)
+                    elif self.is_subtype(type2, type1):
+                        keep.remove(i2)
+                    else:
+                        continue
+                    reduced = True
+                    break
+        return [types[i] for i in keep]
+
+    def lookup_member(self, type: Type, member_name: str) -> Optional[Type]:
+        match type:
+            case BaseType(name=name):
+                if name in self._members:
+                    if member_name in self._members[name]:
+                        return self._members[name][member_name]
+                return None
+
+            case AliasType(name=name, type=base):
+                if name in self._members:
+                    if member_name in self._members[name]:
+                        return self._members[name][member_name]
+                return self.lookup_member(base, member_name)
+
+            case AppliedType(name=name, body=body, args=args):
+                generic: Type = self.get_type(name)
+
+                if not isinstance(generic, GenericType):
+                    raise ValueError("AppliedType not derived from a GenericType")
+
+                substitutions = {
+                    type_var.name: arg for arg, type_var in zip(args, generic.params)
+                }
+                if name in self._members:
+                    if member_name in self._members[name]:
+                        member_type: Type = self._members[name][member_name]
+                        return substitute_typevars(member_type, substitutions)
+
+                member_type2: Optional[Type] = self.lookup_member(body, member_name)
+                if member_type2 is not None:
+                    member_type2 = substitute_typevars(member_type2, substitutions)
+                return member_type2
+
+            case ComplexType(members=members):
+                if member_name in members:
+                    return members[member_name]
+                self.logger.debug(f"No member '{member_name}' in {type}")
+                return None
+
+            case ExtensionType(base=base, extension=ComplexType(members=members)):
+                if member_name in members:
+                    return members[member_name]
+                self.logger.debug(
+                    f"No member '{member_name}' on {type}, looking up in base"
+                )
+                return self.lookup_member(base, member_name)
+
+            case UnknownType():
+                return UnknownType()
+
+            case _:
+                self.logger.debug(f"Can't get member on {type}")
+                return None

midas/checker/reporter.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from midas.ast.location import Location
+from midas.checker.diagnostic import Diagnostic, DiagnosticType
+
+
+class Reporter:
+    def __init__(self):
+        self.diagnostics: list[Diagnostic] = []
+
+    def report(
+        self,
+        path: Optional[str],
+        type: DiagnosticType,
+        location: Location,
+        message: str,
+    ):
+        self.diagnostics.append(
+            Diagnostic(
+                file_path=path,
+                location=location,
+                type=type,
+                message=message,
+            )
+        )
+
+    def for_file(self, path: Optional[str]) -> FileReporter:
+        return FileReporter(self, path)
+
+
+class FileReporter:
+    def __init__(self, base_reporter: Reporter, path: Optional[str]) -> None:
+        self.base_reporter: Reporter = base_reporter
+        self.path: Optional[str] = path
+
+    def for_file(self, path: Optional[str]) -> FileReporter:
+        return FileReporter(self.base_reporter, path)
+
+    def report(self, type: DiagnosticType, location: Location, message: str):
+        self.base_reporter.report(self.path, type, location, message)
+
+    def error(self, location: Location, message: str):
+        self.report(
+            type=DiagnosticType.ERROR,
+            location=location,
+            message=message,
+        )
+
+    def warning(self, location: Location, message: str):
+        self.report(
+            type=DiagnosticType.WARNING,
+            location=location,
+            message=message,
+        )
+
+    def info(self, location: Location, message: str):
+        self.report(
+            type=DiagnosticType.INFO,
+            location=location,
+            message=message,
+        )

midas/checker/resolver.py

@@ -0,0 +1,236 @@
+import midas.ast.python as p
+
+
+class ResolverError(Exception): ...
+
+
+class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
+    """A variable assignment and reference resolver
+
+    This class keeps track of which scope a variable is defined in and which
+    scope is referred to when a variable is referenced
+    """
+
+    def __init__(self):
+        self.locals: dict[p.Expr, int] = {}
+        self.scopes: list[dict[str, bool]] = [{}]
+
+    def resolve(self, *objects: p.Stmt | p.Expr) -> None:
+        """Resolve the given statements or expressions"""
+
+        for obj in objects:
+            obj.accept(self)
+
+    def begin_scope(self):
+        """Begin a new scope inside the current one"""
+        self.scopes.append({})
+
+    def end_scope(self):
+        """Close the current scope"""
+        self.scopes.pop()
+
+    def declare(self, name: str) -> None:
+        """Declare a variable in the current scope
+
+        This method must be called *before* evaluating the variable initializer
+
+        Args:
+            name (str): the name of the variable
+
+        Raises:
+            ResolverError: if the variable has already been declared in the current scope
+        """
+        if len(self.scopes) == 0:
+            return
+        scope: dict[str, bool] = self.scopes[-1]
+        if name in scope:
+            raise ResolverError(
+                f"A variable with the name {name} is already declared in this scope"
+            )
+        scope[name] = False
+
+    def define(self, name: str) -> None:
+        """Define a variable in the current scope
+
+        This method must be called *after* evaluating the variable initializer
+
+        Args:
+            name (str): the name of the variable
+        """
+        if len(self.scopes) == 0:
+            return
+        self.scopes[-1][name] = True
+
+    def resolve_local(self, expr: p.Expr, name: str) -> None:
+        """Resolve a variable reference and store the scope distance
+
+        This method associates to the variable expression a number representing
+        the "distance" of the variable declaration, i.e. the number of scope
+        levels to go "up" to find the closest declaration for that variable.
+
+        Args:
+            expr (p.Expr): the variable expression
+            name (str): the name of the variable
+        """
+        for i, scope in enumerate(reversed(self.scopes)):
+            if name in scope:
+                self.locals[expr] = i
+                return
+
+    def is_defined(self, name: str) -> bool:
+        for scope in self.scopes:
+            if name in scope:
+                return True
+        return False
+
+    def resolve_function(self, function: p.Function) -> None:
+        """Resolve a function definition
+
+        This method creates a new scope for the function, resolves all the
+        parameter declarations and then the body.
+
+        Args:
+            function (p.Function): the function to resolve
+        """
+        self.begin_scope()
+        for param in function.all_args:
+            self.declare(param.name)
+            self.define(param.name)
+        self.resolve(*function.body)
+        self.end_scope()
+
+    def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
+        stmt.expr.accept(self)
+
+    def visit_function(self, stmt: p.Function) -> None:
+        # Declare before resolving body to allow recursion
+        self.declare(stmt.name)
+        self.define(stmt.name)
+        self.resolve_function(stmt)
+
+    def visit_type_assign(self, stmt: p.TypeAssign) -> None:
+        self.declare(stmt.name)
+        # NOTE: resolve type here?
+        self.define(stmt.name)
+
+    def visit_assign_stmt(self, stmt: p.AssignStmt) -> None:
+        self.resolve(stmt.value)
+        for target in stmt.targets:
+            self._visit_assign(target)
+
+    def _visit_assign(self, target: p.Expr):
+        match target:
+            case p.VariableExpr(name=name):
+                if not self.is_defined(name):
+                    self.declare(name)
+                    self.define(name)
+                target.accept(self)
+
+            case p.GetExpr():
+                target.accept(self)
+            case _:
+                raise Exception(f"Unsupported assignment to {target}")
+
+    def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
+        if stmt.value is not None:
+            self.resolve(stmt.value)
+
+    def visit_if_stmt(self, stmt: p.IfStmt) -> None:
+        # Not resolved in sub-environment because assignments in the test leak out of the if
+        # For example:
+        # if (m := 1 + 1) < 2:
+        #     ...
+        # print(m)  # <- m is still defined
+        self.resolve(stmt.test)
+
+        # Body
+        self.begin_scope()
+        self.resolve(*stmt.body)
+        self.end_scope()
+
+        # Else
+        self.begin_scope()
+        self.resolve(*stmt.orelse)
+        self.end_scope()
+
+    def visit_pass(self, stmt: p.Pass) -> None:
+        pass
+
+    def visit_for_stmt(self, stmt: p.ForStmt) -> None:
+        self.resolve(stmt.iterator)
+        self._visit_assign(stmt.target)
+        self.begin_scope()
+        self.resolve(*stmt.body)
+        self.end_scope()
+
+    def visit_raw_stmt(self, stmt: p.RawStmt) -> None:
+        pass
+
+    def visit_binary_expr(self, expr: p.BinaryExpr) -> None:
+        self.resolve(expr.left)
+        self.resolve(expr.right)
+
+    def visit_compare_expr(self, expr: p.CompareExpr) -> None:
+        self.resolve(expr.left)
+        self.resolve(expr.right)
+
+    def visit_unary_expr(self, expr: p.UnaryExpr) -> None:
+        self.resolve(expr.right)
+
+    def visit_call_expr(self, expr: p.CallExpr) -> None:
+        self.resolve(expr.callee)
+        for arg in expr.arguments:
+            self.resolve(arg)
+        for arg in expr.keywords.values():
+            self.resolve(arg)
+
+    def visit_get_expr(self, expr: p.GetExpr) -> None:
+        self.resolve(expr.object)
+
+    def visit_literal_expr(self, expr: p.LiteralExpr) -> None:
+        pass
+
+    def visit_variable_expr(self, expr: p.VariableExpr) -> None:
+        if len(self.scopes) != 0 and self.scopes[-1].get(expr.name) is False:
+            raise ResolverError(
+                f"Cannot use local variable '{expr.name}' in its own initializer"
+            )  # aka. UnboundLocalError
+        self.resolve_local(expr, expr.name)
+
+    def visit_logical_expr(self, expr: p.LogicalExpr) -> None:
+        self.resolve(expr.left)
+        self.resolve(expr.right)
+
+    def visit_cast_expr(self, expr: p.CastExpr) -> None:
+        self.resolve(expr.expr)
+
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> None:
+        self.resolve(expr.test)
+        self.resolve(expr.if_true)
+        self.resolve(expr.if_false)
+
+    def visit_list_expr(self, expr: p.ListExpr) -> None:
+        for item in expr.items:
+            self.resolve(item)
+
+    def visit_dict_expr(self, expr: p.DictExpr) -> None:
+        for key in expr.keys:
+            if key is not None:
+                self.resolve(key)
+        for value in expr.values:
+            self.resolve(value)
+
+    def visit_subscript_expr(self, expr: p.SubscriptExpr) -> None:
+        self.resolve(expr.object)
+        self.resolve(expr.index)
+
+    def visit_slice_expr(self, expr: p.SliceExpr) -> None:
+        if expr.lower is not None:
+            self.resolve(expr.lower)
+        if expr.upper is not None:
+            self.resolve(expr.upper)
+        if expr.step is not None:
+            self.resolve(expr.step)
+
+    def visit_raw_expr(self, expr: p.RawExpr) -> None:
+        pass

midas/checker/types.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass(frozen=True, kw_only=True)
+class TopType:
+    def __str__(self) -> str:
+        return "Any"
+
+
+@dataclass(frozen=True, kw_only=True)
+class BaseType:
+    name: str
+
+    def __str__(self) -> str:
+        return self.name
+
+
+@dataclass(frozen=True, kw_only=True)
+class AliasType:
+    name: str
+    type: Type
+
+    def __str__(self) -> str:
+        return self.name
+
+
+@dataclass(frozen=True, kw_only=True)
+class UnknownType:
+    def __str__(self) -> str:
+        return "<Unknown>"
+
+
+@dataclass(frozen=True, kw_only=True)
+class UnitType:
+    def __str__(self) -> str:
+        return "None"
+
+
+@dataclass(frozen=True, kw_only=True)
+class Function:
+    pos_args: list[Argument] = field(default_factory=list)
+    args: list[Argument] = field(default_factory=list)
+    kw_args: list[Argument] = field(default_factory=list)
+    returns: Type
+
+    def __str__(self) -> str:
+        args: list[str] = []
+        if len(self.pos_args) != 0:
+            args += list(map(str, self.pos_args))
+            args.append("/")
+
+        if len(self.args) != 0:
+            args += list(map(str, self.args))
+
+        if len(self.kw_args) != 0:
+            args.append("*")
+            args += list(map(str, self.kw_args))
+
+        return f"({', '.join(args)}) -> {self.returns}"
+
+    @dataclass(frozen=True, kw_only=True)
+    class Argument:
+        pos: int
+        name: str
+        type: Type
+        required: bool
+
+        def __str__(self) -> str:
+            opt: str = "" if self.required else "?"
+            return f"{self.name}: {self.type}{opt}"
+
+
+@dataclass(frozen=True, kw_only=True)
+class OverloadedFunction:
+    overloads: list[Type]
+
+    def __str__(self) -> str:
+        return "<overloaded function>"
+
+
+@dataclass(frozen=True, kw_only=True)
+class ComplexType:
+    members: dict[str, Type]
+
+    def __str__(self) -> str:
+        props: list[str] = [f"{name}: {type}" for name, type in self.members.items()]
+        return f"{{{', '.join(props)}}}"
+
+
+@dataclass(frozen=True, kw_only=True)
+class ExtensionType:
+    base: Type
+    extension: ComplexType
+
+    def __str__(self) -> str:
+        return f"{self.base} & {self.extension}"
+
+
+@dataclass(frozen=True, kw_only=True)
+class TypeVar:
+    name: str
+    bound: Optional[Type]
+
+    def __str__(self) -> str:
+        if self.bound is not None:
+            return f"{self.name} <: {self.bound}"
+        return self.name
+
+
+@dataclass(frozen=True, kw_only=True)
+class GenericType:
+    name: str
+    params: list[TypeVar]
+    body: Type
+
+    def __str__(self) -> str:
+        return f"{self.name}[{', '.join(map(str, self.params))}]"
+
+
+@dataclass(frozen=True, kw_only=True)
+class AppliedType:
+    name: str
+    args: list[Type]
+    body: Type
+
+    def __str__(self) -> str:
+        return f"{self.name}[{', '.join(map(str, self.args))}]"
+
+
+def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
+    def sub_argument(arg: Function.Argument):
+        return Function.Argument(
+            pos=arg.pos,
+            name=arg.name,
+            type=substitute_typevars(arg.type, substitutions),
+            required=arg.required,
+        )
+
+    match type:
+        case BaseType(name=name) if name in substitutions:
+            return substitutions[name]
+
+        case BaseType():
+            return type
+
+        case AliasType(name=name, type=type2):
+            return AliasType(name=name, type=substitute_typevars(type2, substitutions))
+
+        case Function(
+            pos_args=pos_args,
+            args=args,
+            kw_args=kw_args,
+            returns=returns,
+        ):
+            return Function(
+                pos_args=list(map(sub_argument, pos_args)),
+                args=list(map(sub_argument, args)),
+                kw_args=list(map(sub_argument, kw_args)),
+                returns=substitute_typevars(returns, substitutions),
+            )
+
+        case OverloadedFunction(overloads=overloads):
+            return OverloadedFunction(
+                overloads=[
+                    substitute_typevars(overload, substitutions)
+                    for overload in overloads
+                ]
+            )
+
+        case ComplexType(members=members):
+            members2: dict[str, Type] = {
+                name: substitute_typevars(prop, substitutions)
+                for name, prop in members.items()
+            }
+            return ComplexType(members=members2)
+
+        case ExtensionType(base=base, extension=ComplexType(members=members)):
+            return ExtensionType(
+                base=substitute_typevars(base, substitutions),
+                extension=ComplexType(
+                    members={
+                        name: substitute_typevars(prop, substitutions)
+                        for name, prop in members.items()
+                    }
+                ),
+            )
+
+        case AppliedType(name=name, args=args, body=body):
+            return AppliedType(
+                name=name,
+                args=[substitute_typevars(arg, substitutions) for arg in args],
+                body=substitute_typevars(body, substitutions),
+            )
+
+        case TypeVar(name=name):
+            if name in substitutions:
+                return substitutions[name]
+            raise ValueError(f"Missing TypeVar substitution for {name}")
+
+        case UnknownType() | UnitType():
+            return type
+
+        case _:
+            raise NotImplementedError(f"Unsupported type {type}")
+
+
+def unfold_type(type: Type) -> Type:
+    match type:
+        case AliasType(type=ref_type):
+            return unfold_type(ref_type)
+        case _:
+            return type
+
+
+Type = (
+    TopType
+    | BaseType
+    | AliasType
+    | UnknownType
+    | UnitType
+    | Function
+    | OverloadedFunction
+    | ComplexType
+    | ExtensionType
+    | TypeVar
+    | GenericType
+    | AppliedType
+)

midas/cli/ansi.py

@@ -0,0 +1,41 @@
+class Ansi:
+    CTRL = "\x1b["
+    RESET = CTRL + "0m"
+    BOLD = CTRL + "1m"
+    DIM = CTRL + "2m"
+    ITALIC = CTRL + "3m"
+    UNDERLINE = CTRL + "4m"
+
+    BLACK = 0
+    RED = 1
+    GREEN = 2
+    YELLOW = 3
+    BLUE = 4
+    MAGENTA = 5
+    CYAN = 6
+    WHITE = 7
+
+    BRIGHT_BLACK = 60
+    BRIGHT_RED = 61
+    BRIGHT_GREEN = 62
+    BRIGHT_YELLOW = 63
+    BRIGHT_BLUE = 64
+    BRIGHT_MAGENTA = 65
+    BRIGHT_CYAN = 66
+    BRIGHT_WHITE = 67
+
+    @classmethod
+    def FG(cls, col: int) -> str:
+        return f"{cls.CTRL}{30 + col}m"
+
+    @classmethod
+    def BG(cls, col: int) -> str:
+        return f"{cls.CTRL}{40 + col}m"
+
+    @classmethod
+    def FG_RGB(cls, r: int, g: int, b: int) -> str:
+        return f"{cls.CTRL}38;2;{r};{g};{b}m"
+
+    @classmethod
+    def BG_RGB(cls, r: int, g: int, b: int) -> str:
+        return f"{cls.CTRL}48;2;{r};{g};{b}m"

midas/cli/commands/__init__.py

@@ -0,0 +1,8 @@
+from .check import check as check
+from .compile import compile as compile
+from .format import format as format
+from .highlight import highlight as highlight
+from .parse import parse as parse
+from .registry import dump_registry as dump_registry
+from .types import types as types
+from .validate import validate as validate

midas/cli/commands/check.py

@@ -0,0 +1,41 @@
+# **Run type checker and report diagnostics**
+# ```shell
+# midas check <file.py> [--types <file.midas>]
+# ```
+
+from pathlib import Path
+from typing import Optional, TextIO
+
+import click
+
+from midas.checker.checker import TypeChecker
+from midas.checker.diagnostic import Diagnostic
+from midas.cli.highlighter import DiagnosticsHighlighter
+from midas.cli.utils import DiagnosticPrinter
+
+
+@click.command(help="Run type checker and report diagnostics")
+@click.argument("file", type=click.File("r"))
+@click.option("-t", "--types", type=click.File("r"), multiple=True)
+@click.option("-l", "--highlight", type=click.File("w"))
+def check(
+    file: TextIO,
+    types: tuple[TextIO],
+    highlight: Optional[TextIO],
+):
+    source_path: Path = Path(file.name).resolve()
+
+    checker = TypeChecker()
+    for types_file in types:
+        checker.import_midas(Path(types_file.name).resolve())
+
+    checker.type_check(source_path)
+    diagnostics: list[Diagnostic] = checker.diagnostics.copy()
+    printer = DiagnosticPrinter()
+    printer.print_all(diagnostics)
+
+    if highlight is not None:
+        source: str = file.read()
+        highlighter = DiagnosticsHighlighter(source)
+        highlighter.highlight(diagnostics)
+        highlighter.dump(highlight)

midas/cli/commands/compile.py

@@ -0,0 +1,42 @@
+# **Compile source**
+# ```shell
+# midas compile <file.py> [--types <file.midas>] [-o <output>] [--assertions|--strict|--no-checks]
+# ```
+
+import sys
+from pathlib import Path
+from typing import TextIO
+
+import click
+
+from midas.checker.checker import TypeChecker
+from midas.checker.diagnostic import Diagnostic, DiagnosticType
+from midas.cli.utils import DiagnosticPrinter
+from midas.generator.generator import Generator
+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)
+def compile(
+    file: TextIO,
+    types: tuple[TextIO],
+):
+    source: str = file.read()
+    source_path: Path = Path(file.name).resolve()
+
+    checker = TypeChecker()
+    for types_file in types:
+        checker.import_midas(Path(types_file.name).resolve())
+
+    typed_ast: TypedAST = checker.type_check_source(source, str(source_path))
+    diagnostics: list[Diagnostic] = checker.diagnostics.copy()
+    printer = DiagnosticPrinter()
+    printer.print_all(diagnostics)
+
+    if any(map(lambda d: d.type == DiagnosticType.ERROR, diagnostics)):
+        sys.exit(1)
+
+    generator = Generator(workdir=source_path.parent)
+    generator.generate(typed_ast, source_path)

midas/cli/commands/format.py

@@ -0,0 +1,25 @@
+from typing import TextIO
+
+import click
+
+import midas.ast.midas as m
+from midas.ast.printer import MidasPrinter
+from midas.lexer.midas import MidasLexer
+from midas.lexer.token import Token
+from midas.parser.midas import MidasParser
+
+
+@click.command(help="Parse and pretty print a Midas file")
+@click.argument("file", type=click.File("r"))
+@click.option("-o", "--output", type=click.File("w"), default="-")
+def format(file: TextIO, output: TextIO):
+    source: str = file.read()
+    printer = MidasPrinter()
+    lexer = MidasLexer(source, file=file.name)
+    tokens: list[Token] = lexer.process()
+    parser = MidasParser(tokens)
+    stmts: list[m.Stmt] = parser.parse()
+    for err in parser.errors:
+        print(err.get_report())
+    for stmt in stmts:
+        output.write(printer.print(stmt) + "\n")

midas/cli/commands/highlight.py

@@ -0,0 +1,66 @@
+import ast
+from typing import TextIO
+
+import click
+
+import midas.ast.midas as m
+import midas.ast.python as p
+from midas.cli.highlighter import (
+    Highlighter,
+    LocatableToken,
+    MidasHighlighter,
+    PythonHighlighter,
+)
+from midas.lexer.midas import MidasLexer
+from midas.lexer.token import Token, TokenType
+from midas.parser.midas import MidasParser
+from midas.parser.python import PythonParser
+
+
+def highlight_python(source: str, path: str) -> Highlighter:
+    tree: ast.Module = ast.parse(source, filename=path)
+    parser = PythonParser()
+    stmts: list[p.Stmt] = parser.parse_module(tree)
+    highlighter = PythonHighlighter(source)
+    for stmt in stmts:
+        highlighter.highlight(stmt)
+    return highlighter
+
+
+def highlight_midas(source: str, path: str) -> Highlighter:
+    lexer = MidasLexer(source, file=path)
+    tokens: list[Token] = lexer.process()
+    parser = MidasParser(tokens)
+    stmts: list[m.Stmt] = parser.parse()
+    highlighter = MidasHighlighter(source)
+    for err in parser.errors:
+        print(err.get_report())
+
+    for stmt in stmts:
+        highlighter.highlight(stmt)
+    for token in tokens:
+        if token.type == TokenType.COMMENT:
+            highlighter.wrap(LocatableToken(token), "comment")
+        elif token.is_keyword:
+            highlighter.wrap(LocatableToken(token), "keyword")
+    return highlighter
+
+
+@click.command(
+    help="Parse a Python or Midas file and produce a highlighted version showing AST node types inline",
+    short_help="Parse and highlight a Python or Midas file",
+)
+@click.argument("file", type=click.File("r"))
+@click.option("-o", "--output", type=click.File("w"), default="-")
+def highlight(output: TextIO, file: TextIO):
+    source: str = file.read()
+    highlighter: Highlighter
+
+    if file.name.endswith(".py"):
+        highlighter = highlight_python(source, file.name)
+    elif file.name.endswith(".midas"):
+        highlighter = highlight_midas(source, file.name)
+    else:
+        raise ValueError("Unsupported file type")
+
+    highlighter.dump(output)

midas/cli/commands/parse.py

@@ -0,0 +1,66 @@
+# **Parse and pretty-print AST**
+# ```shell
+# midas parse <file.midas / file.py>
+# ```
+
+import ast
+from typing import TextIO
+
+import click
+
+import midas.ast.midas as m
+import midas.ast.python as p
+from midas.ast.printer import MidasAstPrinter, PythonAstPrinter
+from midas.lexer.midas import MidasLexer
+from midas.lexer.token import Token
+from midas.parser.midas import MidasParser
+from midas.parser.python import PythonParser
+
+
+def dump_python_ast(tree: ast.Module) -> str:
+    parser = PythonParser()
+    stmts: list[p.Stmt] = parser.parse_module(tree)
+    printer = PythonAstPrinter()
+    dump: str = ""
+    for stmt in stmts:
+        dump += printer.print(stmt)
+        dump += "\n"
+    return dump
+
+
+def dump_midas_ast(source: str, filename: str) -> str:
+    lexer = MidasLexer(source, file=filename)
+    tokens: list[Token] = lexer.process()
+    parser = MidasParser(tokens)
+    stmts: list[m.Stmt] = parser.parse()
+    if len(parser.errors) != 0:
+        for err in parser.errors:
+            print(err.get_report())
+        raise RuntimeError("A parsing error occurred")
+    printer = MidasAstPrinter()
+    dump: str = ""
+    for stmt in stmts:
+        dump += printer.print(stmt)
+        dump += "\n"
+    return dump
+
+
+@click.command(help="Parse a Python or Midas file and pretty-print its AST")
+@click.argument("file", type=click.File("r"))
+@click.option("--raw", is_flag=True)
+def parse(file: TextIO, raw: bool):
+    source: str = file.read()
+
+    dump: str
+    if file.name.endswith(".py"):
+        tree: ast.Module = ast.parse(source, filename=file.name)
+        if raw:
+            dump = ast.dump(tree, indent=4)
+        else:
+            dump = dump_python_ast(tree)
+    elif file.name.endswith(".midas"):
+        dump = dump_midas_ast(source, file.name)
+    else:
+        raise ValueError("Unsupported file type")
+
+    click.echo(dump)

midas/cli/commands/registry.py

@@ -0,0 +1,44 @@
+# **Dump types registry**
+# ```shell
+# midas dump-registry [--types <file.midas>]
+# ```
+
+from pathlib import Path
+from typing import TextIO
+
+import click
+
+from midas.checker.checker import TypeChecker
+from midas.checker.types import AliasType, AppliedType, BaseType, GenericType, Type
+
+
+def base_type(type: Type) -> Type:
+    match type:
+        case BaseType():
+            return type
+        case AliasType(type=base):
+            return base
+        case AppliedType(body=body):
+            return body
+        case GenericType(body=body):
+            return body
+        case _:
+            return type
+
+
+@click.command(help="Dump types registry")
+@click.option("-t", "--types", type=click.File("r"), multiple=True)
+def dump_registry(
+    types: tuple[TextIO],
+):
+    checker = TypeChecker()
+    for types_file in types:
+        checker.import_midas(Path(types_file.name).resolve())
+
+    for name, type in checker.types._types.items():
+        members: dict[str, Type] = checker.types._members.get(name, {})
+        print(f"{name} = {base_type(type)}")
+        if len(members) != 0:
+            print(" " * 4 + "Members:")
+            for member_name, member_type in members.items():
+                print(" " * 8 + f"{member_name}: {member_type}")

midas/cli/commands/types.py

@@ -0,0 +1,51 @@
+# **Print judgements**
+# ```shell
+# midas types <file.py> [--types <file.midas>]
+# ```
+
+from pathlib import Path
+from typing import Optional, TextIO
+
+import click
+
+from midas.checker.checker import TypeChecker
+from midas.checker.diagnostic import Diagnostic, DiagnosticType
+from midas.cli.highlighter import DiagnosticsHighlighter
+from midas.cli.utils import DiagnosticPrinter
+
+
+@click.command(help="Print typing judgements")
+@click.argument("file", type=click.File("r"))
+@click.option("-t", "--types", type=click.File("r"), multiple=True)
+@click.option("-l", "--highlight", type=click.File("w"))
+def types(
+    file: TextIO,
+    types: tuple[TextIO],
+    highlight: Optional[TextIO],
+):
+    source_path: Path = Path(file.name).resolve()
+
+    checker = TypeChecker()
+    for types_file in types:
+        checker.import_midas(Path(types_file.name).resolve())
+
+    checker.type_check(source_path)
+
+    diagnostics: list[Diagnostic] = []
+    for expr, type in checker.python_typer.judgements:
+        diagnostics.append(
+            Diagnostic(
+                file_path=str(source_path),
+                location=expr.location,
+                type=DiagnosticType.INFO,
+                message=f"Type: {type}",
+            )
+        )
+    printer = DiagnosticPrinter()
+    printer.print_all(diagnostics)
+
+    if highlight is not None:
+        source: str = file.read()
+        highlighter = DiagnosticsHighlighter(source)
+        highlighter.highlight(diagnostics)
+        highlighter.dump(highlight)

midas/cli/commands/validate.py

@@ -0,0 +1,37 @@
+# **Validate midas definitions**
+# ```shell
+# midas validate <file.midas>
+# ```
+
+from pathlib import Path
+from typing import Optional, TextIO
+
+import click
+
+from midas.checker.checker import TypeChecker
+from midas.checker.diagnostic import Diagnostic
+from midas.cli.highlighter import DiagnosticsHighlighter
+from midas.cli.utils import DiagnosticPrinter
+
+
+@click.command(help="Validate Midas definitions")
+@click.argument("file", type=click.File("r"))
+@click.option("-l", "--highlight", type=click.File("w"))
+def validate(
+    file: TextIO,
+    highlight: Optional[TextIO],
+):
+    source_path: Path = Path(file.name).resolve()
+
+    checker = TypeChecker()
+    checker.import_midas(source_path)
+
+    diagnostics: list[Diagnostic] = checker.diagnostics.copy()
+    printer = DiagnosticPrinter()
+    printer.print_all(diagnostics)
+
+    if highlight is not None:
+        source: str = file.read()
+        highlighter = DiagnosticsHighlighter(source)
+        highlighter.highlight(diagnostics)
+        highlighter.dump(highlight)

midas/cli/highlight.css

@@ -0,0 +1,58 @@
+html,
+body {
+    margin: 0;
+    font-size: 14pt;
+}
+
+* {
+    box-sizing: border-box;
+}
+
+#code {
+    display: flex;
+    flex-direction: column;
+    font-family: monospace;
+    white-space: pre-wrap;
+}
+
+.line {
+    display: flex;
+
+    &:nth-child(odd) {
+        background-color: rgb(247, 247, 247);
+    }
+
+    .no {
+        width: 4em;
+        text-align: right;
+        padding: 0.2em 0.4em;
+        border-right: solid black 1px;
+        flex-shrink: 0;
+    }
+
+    .txt {
+        flex-grow: 1;
+        padding: 0.2em 0.8em;
+    }
+}
+
+span {
+    --col: transparent;
+    --opacity: 0.1;
+    --border: 0px;
+    background-color: rgba(var(--col), var(--opacity));
+    outline: solid rgb(var(--col)) var(--border);
+    outline-offset: 2px;
+    border-radius: 2px;
+
+    &:hover:not(:has(*:hover)) {
+        --opacity: 0.8;
+        --border: 2px;
+        z-index: 10;
+    }
+
+    &.keyword {
+        color: rgb(211, 72, 9);
+        pointer-events: none;
+    }
+}

midas/cli/highlighter.py

@@ -0,0 +1,336 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Generic, Optional, Protocol, TextIO, TypeVar
+
+import midas.ast.midas as m
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.diagnostic import Diagnostic
+from midas.lexer.token import Token
+
+H = TypeVar("H", bound="Highlighter", contravariant=True)
+
+
+class Highlightable(Protocol, Generic[H]):
+    def accept(self, visitor: H): ...
+
+
+class Locatable(Protocol):
+    @property
+    @abstractmethod
+    def location(self) -> Optional[Location]: ...
+
+
+@dataclass(frozen=True)
+class LocatableToken:
+    token: Token
+
+    @property
+    def location(self) -> Location:
+        return self.token.get_location()
+
+
+class Highlighter(ABC):
+    BASE_CSS_PATH: Path = Path(__file__).parent / "highlight.css"
+    EXTRA_CSS_PATH: Optional[Path] = None
+
+    def __init__(self, source: str) -> None:
+        self.source: str = source
+        self.lines: list[str] = self.source.splitlines()
+        self.openings: dict[tuple[int, int], list[str]] = {}
+        self.closings: dict[tuple[int, int], list[str]] = {}
+
+    def format_css(self, path: Path) -> list[str]:
+        css: str = path.read_text()
+        css = "\n".join(("        " + line).rstrip() for line in css.splitlines())
+        return [
+            "    <style>",
+            css,
+            "    </style>",
+        ]
+
+    def dump(self, buf: TextIO):
+        base_css: list[str] = self.format_css(self.BASE_CSS_PATH)
+        extra_css: list[str] = (
+            self.format_css(self.EXTRA_CSS_PATH)
+            if self.EXTRA_CSS_PATH is not None
+            else []
+        )
+        lines: list[str] = [
+            "<!DOCTYPE html>",
+            '<html lang="en">',
+            "<head>",
+            '    <meta charset="UTF-8">',
+            '    <meta name="viewport" content="width=device-width, initial-scale=1.0">',
+            "    <title>Highlighted file</title>",
+            *base_css,
+            *extra_css,
+            "</head>",
+            "<body>",
+            '    <div id="code">',
+        ]
+        for l, line in enumerate(self.lines):
+            lineno: int = l + 1
+            line_buf: str = (
+                f'<div class="line" id="l{lineno}"><div class="no">{lineno}</div><div class="txt">'
+            )
+            for c, char in enumerate(line):
+                pos: tuple[int, int] = (lineno, c)
+                closings: list[str] = self.closings.get(pos, [])
+                openings: list[str] = self.openings.get(pos, [])
+                line_buf += "".join(closings + openings)
+                line_buf += char
+            line_buf += "".join(self.closings.get((lineno, len(line)), []))
+            line_buf += "</div></div>"
+            lines.append("        " + line_buf)
+        lines.extend(
+            [
+                "    </div>",
+                "</body>",
+                "</html>",
+            ]
+        )
+
+        buf.write("\n".join(lines))
+
+    def wrap(self, node: Locatable, cls: str, message: Optional[str] = None):
+        if node.location is None:
+            return
+        if node.location.end_lineno is None or node.location.end_col_offset is None:
+            return
+        start_pos: tuple[int, int] = (node.location.lineno, node.location.col_offset)
+        end_pos: tuple[int, int] = (
+            node.location.end_lineno,
+            node.location.end_col_offset,
+        )
+        opening: str = f'<span class="{cls}" title="{cls}">'
+        closing: str = "</span>"
+        if message is not None:
+            opening = f'<span class="with-msg">{opening}'
+            closing = f'{closing}<span class="message">{message}</span></span>'
+
+        self.openings.setdefault(start_pos, []).append(opening)
+        self.closings.setdefault(end_pos, []).insert(0, closing)
+        if start_pos[0] != end_pos[0]:
+            for l in range(start_pos[0], end_pos[0]):
+                c: int = len(self.lines[l - 1])
+                self.closings.setdefault((l, c), []).insert(0, closing)
+                self.openings.setdefault((l + 1, 0), []).append(opening)
+
+
+class PythonHighlighter(
+    Highlighter,
+    p.MidasType.Visitor[None],
+    p.Stmt.Visitor[None],
+    p.Expr.Visitor[None],
+):
+    EXTRA_CSS_PATH: Optional[Path] = Path(__file__).parent / "hl_python.css"
+
+    def highlight(self, node: Highlightable[PythonHighlighter]):
+        node.accept(self)
+
+    def visit_base_type(self, node: p.BaseType) -> None:
+        self.wrap(node, "base-type")
+        if node.param is not None:
+            self.wrap(node.param, "param")
+            node.param.accept(self)
+
+    def visit_constraint_type(self, node: p.ConstraintType) -> None:
+        self.wrap(node, "constraint-type")
+        node.type.accept(self)
+
+    def visit_frame_column(self, node: p.FrameColumn) -> None:
+        self.wrap(node, "frame-column")
+        if node.type is not None:
+            node.type.accept(self)
+
+    def visit_frame_type(self, node: p.FrameType) -> None:
+        self.wrap(node, "frame-type")
+        for column in node.columns:
+            column.accept(self)
+
+    def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
+        stmt.expr.accept(self)
+
+    def visit_function(self, stmt: p.Function) -> None:
+        self.wrap(stmt, "function")
+        for arg in stmt.posonlyargs + stmt.args + stmt.kwonlyargs:
+            self._highlight_function_argument(arg)
+        for body_stmt in stmt.body:
+            body_stmt.accept(self)
+
+    def _highlight_function_argument(self, arg: p.Function.Argument) -> None:
+        self.wrap(arg, "argument")
+        if arg.type is not None:
+            arg.type.accept(self)
+
+    def visit_type_assign(self, stmt: p.TypeAssign) -> None:
+        stmt.type.accept(self)
+
+    def visit_assign_stmt(self, stmt: p.AssignStmt) -> None:
+        for target in stmt.targets:
+            target.accept(self)
+        stmt.value.accept(self)
+
+    def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
+        self.wrap(stmt, "return")
+        if stmt.value is not None:
+            stmt.value.accept(self)
+
+    def visit_if_stmt(self, stmt: p.IfStmt) -> None:
+        self.wrap(stmt, "if")
+        stmt.test.accept(self)
+        for body_stmt in stmt.body:
+            body_stmt.accept(self)
+        for else_stmt in stmt.orelse:
+            else_stmt.accept(self)
+
+    def visit_pass(self, stmt: p.Pass) -> None:
+        pass
+
+    def visit_for_stmt(self, stmt: p.ForStmt) -> None:
+        self.wrap(stmt, "for")
+        stmt.iterator.accept(self)
+        stmt.target.accept(self)
+        for body_stmt in stmt.body:
+            body_stmt.accept(self)
+
+    def visit_binary_expr(self, expr: p.BinaryExpr) -> None: ...
+
+    def visit_compare_expr(self, expr: p.CompareExpr) -> None: ...
+
+    def visit_unary_expr(self, expr: p.UnaryExpr) -> None: ...
+
+    def visit_call_expr(self, expr: p.CallExpr) -> None:
+        self.wrap(expr, "call")
+        expr.callee.accept(self)
+        for arg in expr.arguments:
+            arg.accept(self)
+        for arg in expr.keywords.values():
+            arg.accept(self)
+
+    def visit_get_expr(self, expr: p.GetExpr) -> None: ...
+
+    def visit_literal_expr(self, expr: p.LiteralExpr) -> None: ...
+
+    def visit_variable_expr(self, expr: p.VariableExpr) -> None: ...
+
+    def visit_logical_expr(self, expr: p.LogicalExpr) -> None: ...
+
+    def visit_cast_expr(self, expr: p.CastExpr) -> None: ...
+
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> None: ...
+
+    def visit_list_expr(self, expr: p.ListExpr) -> None:
+        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)
+
+    def visit_slice_expr(self, expr: p.SliceExpr) -> None:
+        if expr.lower is not None:
+            expr.lower.accept(self)
+        if expr.upper is not None:
+            expr.upper.accept(self)
+        if expr.step is not None:
+            expr.step.accept(self)
+
+
+class MidasHighlighter(
+    Highlighter, m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[None]
+):
+    EXTRA_CSS_PATH: Optional[Path] = Path(__file__).parent / "hl_midas.css"
+
+    def highlight(self, node: Highlightable[MidasHighlighter]):
+        node.accept(self)
+
+    def visit_type_stmt(self, stmt: m.TypeStmt) -> None:
+        self.wrap(stmt, "type-stmt")
+        self.wrap(LocatableToken(stmt.name), "type-name")
+        stmt.type.accept(self)
+
+    def visit_member_stmt(self, stmt: m.MemberStmt) -> None:
+        self.wrap(stmt, "member")
+        stmt.type.accept(self)
+
+    def visit_extend_stmt(self, stmt: m.ExtendStmt) -> None:
+        self.wrap(stmt, "extend")
+        for member in stmt.members:
+            member.accept(self)
+
+    def visit_predicate_stmt(self, stmt: m.PredicateStmt) -> None:
+        self.wrap(stmt, "predicate")
+        self.wrap(LocatableToken(stmt.name), "predicate-name")
+        stmt.type.accept(self)
+        stmt.condition.accept(self)
+
+    def visit_logical_expr(self, expr: m.LogicalExpr) -> None:
+        self.wrap(expr, "logical-expr")
+        expr.left.accept(self)
+        expr.right.accept(self)
+
+    def visit_binary_expr(self, expr: m.BinaryExpr) -> None:
+        self.wrap(expr, "binary-expr")
+        expr.left.accept(self)
+        expr.right.accept(self)
+
+    def visit_unary_expr(self, expr: m.UnaryExpr) -> None:
+        self.wrap(expr, "unary-expr")
+        expr.right.accept(self)
+
+    def visit_get_expr(self, expr: m.GetExpr) -> None:
+        self.wrap(expr, "get-expr")
+        expr.expr.accept(self)
+
+    def visit_variable_expr(self, expr: m.VariableExpr) -> None:
+        self.wrap(expr, "variable")
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr) -> None:
+        expr.expr.accept(self)
+
+    def visit_literal_expr(self, expr: m.LiteralExpr) -> None: ...
+
+    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> None: ...
+
+    def visit_named_type(self, type: m.NamedType) -> None:
+        self.wrap(type, "named-type")
+
+    def visit_generic_type(self, type: m.GenericType) -> None:
+        self.wrap(type, "generic-type")
+        type.type.accept(self)
+        for arg in type.args:
+            arg.accept(self)
+
+    def visit_constraint_type(self, type: m.ConstraintType) -> None:
+        self.wrap(type, "constraint-type")
+        type.type.accept(self)
+        type.constraint.accept(self)
+
+    def visit_complex_type(self, type: m.ComplexType) -> None:
+        self.wrap(type, "complex-type")
+        for member in type.members:
+            member.accept(self)
+
+    def visit_function_type(self, type: m.FunctionType) -> None:
+        self.wrap(type, "function")
+        for arg in type.pos_args + type.args + type.kw_args:
+            arg.type.accept(self)
+        type.returns.accept(self)
+
+    def visit_extension_type(self, type: m.ExtensionType) -> None:
+        self.wrap(type, "extension")
+        type.base.accept(self)
+        type.extension.accept(self)
+
+
+class DiagnosticsHighlighter(Highlighter):
+    EXTRA_CSS_PATH: Optional[Path] = Path(__file__).parent / "hl_diagnostic.css"
+
+    def highlight(self, diagnostics: list[Diagnostic]):
+        for diagnostic in diagnostics:
+            self.wrap(diagnostic, str(diagnostic.type).lower(), diagnostic.message)

midas/cli/hl_diagnostic.css

@@ -0,0 +1,39 @@
+span {
+    --opacity: 0.4;
+
+    &.error {
+        --col: 255, 0, 0;
+    }
+    &.warning {
+        --col: 250, 160, 0;
+    }
+    &.info {
+        --col: 150, 190, 250;
+    }
+
+    &.with-msg {
+        position: relative;
+
+        .message {
+            display: none;
+        }
+
+        &:hover:not(:has(.with-msg:hover)) {
+            .message {
+                display: inline-block;
+            }
+        }
+
+        .message {
+            position: absolute;
+            top: calc(100% + 0.2em);
+            left: -.2em;
+            background-color: black;
+            color: white;
+            padding: 0.2em 0.4em;
+            border-radius: .2em;
+            z-index: 10;
+            width: 300%;
+        }
+    }
+}

midas/cli/hl_midas.css

@@ -0,0 +1,52 @@
+span {
+    &.comment {
+        --col: 200, 200, 200;
+        color: rgb(110, 110, 110);
+        font-style: italic;
+    }
+
+    &.named-type,
+    &.generic-type,
+    &.constraint-type,
+    &.complex-type {
+        --col: 150, 150, 150;
+    }
+
+    &.constraint {
+        --col: 233, 108, 108;
+    }
+
+    &.property {
+        --col: 233, 108, 176;
+    }
+
+    &.extend {
+        --col: 108, 197, 233;
+    }
+
+    &.op {
+        --col: 108, 148, 233;
+    }
+
+    &.predicate {
+        --col: 193, 108, 233;
+    }
+
+    &.logical-expr,
+    &.binary-expr,
+    &.unary-expr,
+    &.get-expr {
+        --col: 123, 215, 193;
+    }
+
+    &.template {
+        --col: 163, 117, 71;
+    }
+
+    &.type-name,
+    &.op-name,
+    &.predicate-name {
+        --col: 200, 200, 200;
+        font-weight: bold;
+    }
+}

midas/cli/hl_python.css

@@ -0,0 +1,29 @@
+span {
+    &.base-type {
+        --col: 108, 233, 108;
+    }
+
+    &.param {
+        --col: 103, 192, 224;
+    }
+
+    &.constraint-type {
+        --col: 174, 200, 195;
+    }
+
+    &.frame-column {
+        --col: 216, 231, 81;
+    }
+
+    &.frame-type {
+        --col: 231, 46, 40;
+    }
+
+    &.function {
+        --col: 215, 103, 224;
+    }
+
+    &.argument {
+        --col: 103, 192, 224;
+    }
+}

midas/cli/main.py

@@ -0,0 +1,25 @@
+import logging
+
+import click
+
+from midas.cli import commands
+
+
+@click.group()
+@click.option("-v", "--verbose", is_flag=True)
+def midas(verbose: bool):
+    logging.basicConfig(level=logging.DEBUG if verbose else logging.WARN)
+
+
+midas.add_command(commands.check)
+midas.add_command(commands.compile)
+midas.add_command(commands.format)
+midas.add_command(commands.highlight)
+midas.add_command(commands.parse)
+midas.add_command(commands.dump_registry)
+midas.add_command(commands.types)
+midas.add_command(commands.validate)
+
+
+if __name__ == "__main__":
+    midas()

midas/cli/utils.py

@@ -0,0 +1,78 @@
+from pathlib import Path
+from typing import Optional
+
+from midas.ast.location import Location
+from midas.checker.diagnostic import Diagnostic, DiagnosticType
+from midas.cli.ansi import Ansi
+
+
+class DiagnosticPrinter:
+    def __init__(self) -> None:
+        self.files: dict[Optional[str], list[str]] = {}
+
+    def get_lines(self, filename: Optional[str]) -> list[str]:
+        if filename is None:
+            return []
+        if filename not in self.files:
+            path: Path = Path(filename)
+            if path.exists() and path.is_file():
+                self.files[filename] = path.read_text().split("\n")
+            else:
+                self.files[filename] = []
+        return self.files[filename]
+
+    def print_all(self, diagnostics: list[Diagnostic], indent: int = 4):
+        for diagnostic in diagnostics:
+            filename: Optional[str] = diagnostic.file_path
+            lines = self.get_lines(filename)
+            self.print(lines, diagnostic, indent=indent)
+
+    def print(self, lines: list[str], diagnostic: Diagnostic, indent: int = 4):
+        """Pretty-print a diagnostic, showing some context if possible
+
+        If the diagnostic concerns a specific part of one line, the line is shown
+        with the affected part highlighted. The message is clearly printed under the
+        line with an underline further indicating the target expression.
+
+        If multiple lines are concerned, no context is shown, only the
+        diagnostic type, location and message
+
+        Args:
+            lines (list[str]): source code lines
+            diagnostic (Diagnostic): the diagnostic to print
+            indent (int, optional): the number of spaces added before the target line to indent if from the location header. Defaults to 4.
+        """
+
+        loc: Location = diagnostic.location
+        if loc.lineno != loc.end_lineno:
+            print(diagnostic)
+            return
+
+        start_offset: int = loc.col_offset
+        end_offset: int = loc.end_col_offset or (start_offset + 1)
+
+        line: str = lines[loc.lineno - 1]
+        before: str = line[:start_offset]
+        after: str = line[end_offset:]
+
+        color: int = {
+            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 = (
+            " " * start_offset
+            + Ansi.FG(color)
+            + "~" * (end_offset - start_offset)
+            + "> "
+            + diagnostic.message
+            + Ansi.RESET
+        )
+
+        indent_str: str = " " * indent
+        print(diagnostic.location_str + ":")
+        print(indent_str + before + subject + after)
+        print(indent_str + cursor)
+        print()

midas/generator/generator.py

@@ -0,0 +1,341 @@
+import ast
+import shutil
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.types import (
+    AliasType,
+    AppliedType,
+    BaseType,
+    ComplexType,
+    ExtensionType,
+    Function,
+    GenericType,
+    OverloadedFunction,
+    TopType,
+    Type,
+    TypeVar,
+    UnitType,
+)
+from midas.utils import TypedAST
+
+
+@dataclass
+class Scope:
+    pre_assertions: list[ast.stmt] = field(default_factory=list)
+    aliases: list[str] = field(default_factory=list)
+
+
+class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
+    def __init__(self, workdir: Path) -> None:
+        self.workdir: Path = workdir.resolve()
+        self.build_dir: Path = self.workdir / "build" / "midas"
+        if self.build_dir.exists():
+            shutil.rmtree(self.build_dir)
+        self.build_dir.mkdir(parents=True, exist_ok=True)
+        self.rel_src_path: Path = Path()
+
+        self._typed_ast: TypedAST = TypedAST(
+            stmts=[],
+            judgements=[],
+        )
+        self._alias_count: int = 0
+        self._scopes: list[Scope] = []
+
+    def generate_ast(self, typed_ast: TypedAST, src_path: Path) -> ast.AST:
+        self.rel_src_path = src_path.relative_to(self.workdir)
+        self._typed_ast = typed_ast
+        body: list[ast.stmt] = self._visit_body(typed_ast.stmts)
+        module = ast.Module(body=body, type_ignores=[])
+        module = ast.fix_missing_locations(module)
+        return module
+
+    def generate(
+        self, typed_ast: TypedAST, src_path: Path, out_path: Optional[Path] = None
+    ) -> Path:
+        module: ast.AST = self.generate_ast(typed_ast, src_path)
+        compiled: str = ast.unparse(module)
+        if out_path is None:
+            out_path = (self.build_dir / self.rel_src_path).resolve()
+            try:
+                _ = out_path.relative_to(self.build_dir)
+            except ValueError:
+                raise ValueError(
+                    f"Directory traversal, {self.rel_src_path} points outside of parent directory"
+                )
+            out_path.parent.mkdir(parents=True, exist_ok=True)
+        out_path.write_text(compiled)
+        return out_path
+
+    def visit_binary_expr(self, expr: p.BinaryExpr) -> ast.expr:
+        return ast.BinOp(
+            left=expr.left.accept(self),
+            op=expr.operator,
+            right=expr.right.accept(self),
+        )
+
+    def visit_compare_expr(self, expr: p.CompareExpr) -> ast.expr:
+        return ast.Compare(
+            left=expr.left.accept(self),
+            ops=[expr.operator],
+            comparators=[expr.right.accept(self)],
+        )
+
+    def visit_unary_expr(self, expr: p.UnaryExpr) -> ast.expr:
+        return ast.UnaryOp(
+            op=expr.operator,
+            operand=expr.right.accept(self),
+        )
+
+    def visit_call_expr(self, expr: p.CallExpr) -> ast.expr:
+        return ast.Call(
+            func=expr.callee.accept(self),
+            args=[arg.accept(self) for arg in expr.arguments],
+            keywords=[
+                ast.keyword(arg=name, value=arg.accept(self))
+                for name, arg in expr.keywords.items()
+            ],
+        )
+
+    def visit_get_expr(self, expr: p.GetExpr) -> ast.expr:
+        return ast.Attribute(
+            value=expr.object.accept(self),
+            attr=expr.name,
+        )
+
+    def visit_literal_expr(self, expr: p.LiteralExpr) -> ast.expr:
+        return ast.Constant(value=expr.value)
+
+    def visit_variable_expr(self, expr: p.VariableExpr) -> ast.expr:
+        return ast.Name(id=expr.name)
+
+    def visit_logical_expr(self, expr: p.LogicalExpr) -> ast.expr:
+        return ast.BoolOp(
+            op=expr.operator,
+            values=[expr.left.accept(self), expr.right.accept(self)],
+        )
+
+    def visit_cast_expr(self, expr: p.CastExpr) -> ast.expr:
+        expr2: ast.expr = expr.expr.accept(self)
+        alias: ast.expr = self._make_alias(expr2)
+
+        type: Type = self._get_expr_type(expr)
+        self._make_cast_asserts(expr.location, alias, type)
+
+        return alias
+
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> ast.expr:
+        return ast.IfExp(
+            test=expr.test.accept(self),
+            body=expr.if_true.accept(self),
+            orelse=expr.if_false.accept(self),
+        )
+
+    def visit_list_expr(self, expr: p.ListExpr) -> ast.expr:
+        return ast.List(
+            elts=[item.accept(self) for item in expr.items],
+        )
+
+    def visit_dict_expr(self, expr: p.DictExpr) -> ast.expr:
+        return ast.Dict(
+            keys=[key.accept(self) if key is not None else None for key in expr.keys],
+            values=[value.accept(self) for value in expr.values],
+        )
+
+    def visit_subscript_expr(self, expr: p.SubscriptExpr) -> ast.expr:
+        return ast.Subscript(
+            value=expr.object.accept(self),
+            slice=expr.index.accept(self),
+        )
+
+    def visit_slice_expr(self, expr: p.SliceExpr) -> ast.expr:
+        return ast.Slice(
+            lower=expr.lower.accept(self) if expr.lower is not None else None,
+            upper=expr.upper.accept(self) if expr.upper is not None else None,
+            step=expr.step.accept(self) if expr.step is not None else None,
+        )
+
+    def visit_raw_expr(self, expr: p.RawExpr) -> ast.expr:
+        return expr.expr
+
+    def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> ast.stmt:
+        return ast.Expr(
+            value=stmt.expr.accept(self),
+        )
+
+    def visit_function(self, stmt: p.Function) -> ast.stmt:
+        return ast.FunctionDef(
+            name=stmt.name,
+            args=ast.arguments(
+                posonlyargs=[ast.arg(arg=arg.name) for arg in stmt.posonlyargs],
+                vararg=None,
+                args=[ast.arg(arg=arg.name) for arg in stmt.args],
+                kwonlyargs=[ast.arg(arg=arg.name) for arg in stmt.kwonlyargs],
+                kwarg=None,
+                defaults=[
+                    arg.default.accept(self)
+                    for arg in stmt.posonlyargs + stmt.args
+                    if arg.default is not None
+                ],
+                kw_defaults=[
+                    arg.default.accept(self) if arg.default is not None else None
+                    for arg in stmt.kwonlyargs
+                ],
+            ),
+            body=self._visit_body(stmt.body),
+            decorator_list=[],
+        )
+
+    def visit_type_assign(self, stmt: p.TypeAssign) -> ast.stmt:
+        # TODO: is that ok?
+        return ast.Pass()
+
+    def visit_assign_stmt(self, stmt: p.AssignStmt) -> ast.stmt:
+        return ast.Assign(
+            targets=[target.accept(self) for target in stmt.targets],
+            value=stmt.value.accept(self),
+        )
+
+    def visit_return_stmt(self, stmt: p.ReturnStmt) -> ast.stmt:
+        return ast.Return(
+            value=stmt.value.accept(self) if stmt.value is not None else None,
+        )
+
+    def visit_if_stmt(self, stmt: p.IfStmt) -> ast.stmt:
+        return ast.If(
+            test=stmt.test.accept(self),
+            body=self._visit_body(stmt.body),
+            orelse=self._visit_body(stmt.orelse),
+        )
+
+    def visit_pass(self, stmt: p.Pass) -> ast.stmt:
+        return ast.Pass()
+
+    def visit_for_stmt(self, stmt: p.ForStmt) -> ast.stmt:
+        return ast.For(
+            target=stmt.target.accept(self),
+            iter=stmt.iterator.accept(self),
+            body=self._visit_body(stmt.body),
+            orelse=[],
+        )
+
+    def visit_raw_stmt(self, stmt: p.RawStmt) -> ast.stmt:
+        return stmt.stmt
+
+    def _visit_body(self, stmts: list[p.Stmt]) -> list[ast.stmt]:
+        generated: list[ast.stmt] = []
+        for stmt in stmts:
+            scope = Scope()
+            self._scopes.append(scope)
+
+            stmt2 = stmt.accept(self)
+            generated.extend(scope.pre_assertions)
+            generated.append(stmt2)
+            if len(scope.aliases) != 0:
+                generated.append(
+                    ast.Delete(targets=[ast.Name(id=alias) for alias in scope.aliases])
+                )
+            self._scopes.pop()
+
+        # Remove redundant pass statements
+        if len(generated) > 1:
+            generated = [stmt for stmt in generated if not isinstance(stmt, ast.Pass)]
+        return generated
+
+    def _make_alias(self, expr: ast.expr) -> ast.expr:
+        name: str = f"__midas_alias_{self._alias_count}__"
+        alias = ast.Name(id=name)
+        self._alias_count += 1
+        self._scopes[-1].aliases.append(name)
+        self._scopes[-1].pre_assertions.append(
+            ast.Assign(
+                targets=[alias],
+                value=expr,
+            )
+        )
+        return alias
+
+    def _add_assert(self, expr: ast.expr, message: str | ast.expr):
+        if isinstance(message, str):
+            message = ast.Constant(value=message)
+        self._scopes[-1].pre_assertions.append(
+            ast.Assert(
+                test=expr,
+                msg=message,
+            )
+        )
+
+    def _get_expr_type(self, query: p.Expr) -> Type:
+        for expr, type in self._typed_ast.judgements:
+            if expr == query:
+                return type
+        raise RuntimeError(f"Cannot get type judgement for {query}")
+
+    def _make_cast_asserts(self, src_location: Location, expr: ast.expr, type: Type):
+        match type:
+            case BaseType(name=name):
+                self._add_assert(
+                    ast.Call(
+                        func=ast.Name(id="isinstance"),
+                        args=[expr, ast.Name(id=name)],
+                        keywords=[],
+                    ),
+                    self._make_cast_assert_message(src_location, expr, type),
+                )
+
+            case AliasType(type=base):
+                self._make_cast_asserts(src_location, expr, base)
+
+            case UnitType():
+                self._add_assert(
+                    ast.Compare(
+                        left=expr,
+                        ops=[ast.Is()],
+                        comparators=[
+                            ast.Constant(value=None),
+                        ],
+                    ),
+                    self._make_cast_assert_message(src_location, expr, type),
+                )
+
+            case AppliedType():
+                self._make_cast_asserts(src_location, expr, type.body)
+
+            case (
+                TopType()
+                | Function()
+                | OverloadedFunction()
+                | ComplexType()
+                | ExtensionType()
+                | GenericType()
+            ):
+                raise NotImplementedError(f"Can't make assertion for type {type}")
+
+            case TypeVar():
+                raise RuntimeError("Unexpected TypeVar")
+
+    def _make_cast_assert_message(
+        self, location: Location, expr: ast.expr, type: Type
+    ) -> ast.expr:
+        loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
+        # f"file.py:L1:1: CastError: Cannot cast {type(expr).__name__} to Type"
+        return ast.JoinedStr(
+            values=[
+                ast.Constant(f"{loc_str}: CastError: Cannot cast "),
+                ast.FormattedValue(
+                    value=ast.Attribute(
+                        value=ast.Call(
+                            func=ast.Name(id="type"),
+                            args=[expr],
+                            keywords=[],
+                        ),
+                        attr="__name__",
+                    ),
+                    conversion=-1,
+                ),
+                ast.Constant(f" to {type}"),
+            ]
+        )

midas/lexer/base.py

@@ -1,8 +1,15 @@
 from abc import ABC, abstractmethod
 from typing import Any, Callable, Optional
 
-from lexer.position import Position
-from lexer.token import Token, TokenType
+from midas.lexer.position import Position
+from midas.lexer.token import Token, TokenType
+
+
+class MidasSyntaxError(Exception):
+    def __init__(self, pos: Position, message: str):
+        super().__init__(f"[ERROR] Error at {pos}: {message}")
+        self.pos: Position = pos
+        self.message: str = message
 
 
 class Lexer(ABC):
@@ -38,9 +45,9 @@ class Lexer(ABC):
             msg (str): the error message
 
         Raises:
-            SyntaxError
+            MidasSyntaxError
         """
-        raise SyntaxError(f"[ERROR] Error at {self.start_pos}: {msg}")
+        raise MidasSyntaxError(self.start_pos, msg)
 
     def process(self) -> list[Token]:
         """Scan tokens out of the source text
@@ -49,7 +56,7 @@ class Lexer(ABC):
             list[Token]: all the tokens that could be scanned
 
         Raises:
-            SyntaxError: if a syntax error is found
+            MidasSyntaxError: if a syntax error is found
         """
         self.scan_tokens()
         self.tokens.append(Token(TokenType.EOF, "", None, self.get_position()))

midas/lexer/midas.py

@@ -1,6 +1,5 @@
-from lexer.base import Lexer
-from lexer.keyword import MIDAS_KEYWORDS
-from lexer.token import TokenType
+from midas.lexer.base import Lexer
+from midas.lexer.token import KEYWORDS, TokenType
 
 
 class MidasLexer(Lexer):
@@ -31,30 +30,34 @@ class MidasLexer(Lexer):
                 self.add_token(
                     TokenType.EQUAL_EQUAL if self.match("=") else TokenType.EQUAL
                 )
-            case "!":
-                if self.match("="):
-                    self.add_token(TokenType.BANG_EQUAL)
-                else:
-                    self.error("Unexpected single bang. Did you mean '!=' ?")
+            case "!" if self.match("="):
+                self.add_token(TokenType.BANG_EQUAL)
             case ":":
                 self.add_token(TokenType.COLON)
+            case ".":
+                self.add_token(TokenType.DOT)
+            case "&":
+                self.add_token(TokenType.AND)
+            case "?":
+                self.add_token(TokenType.QMARK)
             case ",":
                 self.add_token(TokenType.COMMA)
-            case "_":
+            case "_" if not self.is_identifier_char(self.peek_next(), start=False):
                 self.add_token(TokenType.UNDERSCORE)
-            case "+":
-                self.add_token(TokenType.PLUS)
+            case "-" if self.match(">"):
+                self.add_token(TokenType.ARROW)
+            # case "+":
+            #     self.add_token(TokenType.PLUS)
             case "-":
                 self.add_token(TokenType.MINUS)
             case "*":
                 self.add_token(TokenType.STAR)
+            case "/" if self.match("/"):
+                self.scan_comment()
+            case "/" if self.match("*"):
+                self.scan_comment_multiline()
             case "/":
-                if self.match("/"):
-                    self.scan_comment()
-                elif self.match("*"):
-                    self.scan_comment_multiline()
-                else:
-                    self.add_token(TokenType.SLASH)
+                self.add_token(TokenType.SLASH)
             case "\n":
                 self.add_token(TokenType.NEWLINE)
             case " " | "\r" | "\t":
@@ -69,7 +72,7 @@ class MidasLexer(Lexer):
             case _:
                 if char.isdigit():
                     self.scan_number()
-                elif char.isalpha():
+                elif self.is_identifier_char(char, start=True):
                     self.scan_identifier()
                 else:
                     self.error("Unexpected character")
@@ -98,11 +101,11 @@ class MidasLexer(Lexer):
         An identifier starts with a letter, followed by any number of
         alphanumerical characters or underscores
         """
-        while self.peek().isalnum() or self.peek() == "_":
+        while self.is_identifier_char(self.peek(), start=False):
             self.advance()
 
         lexeme: str = self.source[self.start : self.idx]
-        token_type: TokenType = MIDAS_KEYWORDS.get(lexeme, TokenType.IDENTIFIER)
+        token_type: TokenType = KEYWORDS.get(lexeme, TokenType.IDENTIFIER)
         self.add_token(token_type)
 
     def scan_comment(self):
@@ -129,3 +132,12 @@ class MidasLexer(Lexer):
         if not self.is_at_end():
             self.advance()
         self.add_token(TokenType.COMMENT)
+
+    def is_identifier_char(self, char: str, *, start: bool) -> bool:
+        if char == "_":
+            return True
+        if char.isalpha():
+            return True
+        if not start and char.isdigit():
+            return True
+        return False

midas/lexer/position.py

@@ -5,6 +5,7 @@ from typing import Optional
 @dataclass(frozen=True)
 class Position:
     """A simple structure to store the position of a token"""
+
     file: Optional[str]
     line: int
     column: int

midas/lexer/token.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import Any
+
+from midas.ast.location import Location
+from midas.lexer.position import Position
+
+
+class TokenType(Enum):
+    # Punctuation
+    LEFT_PAREN = auto()
+    RIGHT_PAREN = auto()
+    LEFT_BRACKET = auto()
+    RIGHT_BRACKET = auto()
+    LEFT_BRACE = auto()
+    RIGHT_BRACE = auto()
+    COLON = auto()
+    COMMA = auto()
+    UNDERSCORE = auto()
+    ARROW = auto()
+    AND = auto()
+    QMARK = auto()
+    DOT = auto()
+
+    # Operators
+    # PLUS = auto()
+    MINUS = auto()
+    STAR = auto()
+    SLASH = auto()
+    GREATER = auto()
+    GREATER_EQUAL = auto()
+    LESS = auto()
+    LESS_EQUAL = auto()
+    EQUAL = auto()
+    EQUAL_EQUAL = auto()
+    BANG_EQUAL = auto()
+
+    # Literals
+    IDENTIFIER = auto()
+    NUMBER = auto()
+    TRUE = auto()
+    FALSE = auto()
+    NONE = auto()
+
+    # Keywords
+    TYPE = auto()
+    PREDICATE = auto()
+    EXTEND = auto()
+    WHERE = auto()
+    PROP = auto()
+    DEF = auto()
+    FUNC = auto()
+
+    # Misc
+    COMMENT = auto()
+    WHITESPACE = auto()
+    EOF = auto()
+    NEWLINE = auto()
+
+
+KEYWORDS: dict[str, TokenType] = {
+    "type": TokenType.TYPE,
+    "predicate": TokenType.PREDICATE,
+    "extend": TokenType.EXTEND,
+    "where": TokenType.WHERE,
+    "true": TokenType.TRUE,
+    "false": TokenType.FALSE,
+    "none": TokenType.NONE,
+    "prop": TokenType.PROP,
+    "def": TokenType.DEF,
+    "fn": TokenType.FUNC,
+}
+
+
+@dataclass(frozen=True)
+class Token:
+    """A scanned token"""
+
+    type: TokenType
+    lexeme: str
+    value: Any
+    position: Position
+
+    def get_location(self) -> Location:
+        lineno: int = self.position.line
+        col_offset: int = self.position.column - 1
+        end_lineno = lineno
+        end_col_offset = col_offset
+        for c in self.lexeme:
+            end_col_offset += 1
+            if c == "\n":
+                end_lineno += 1
+                end_col_offset = 0
+        return Location(
+            lineno=lineno,
+            col_offset=col_offset,
+            end_lineno=end_lineno,
+            end_col_offset=end_col_offset,
+        )
+
+    def location_to(self, to: Token) -> Location:
+        return Location.span(self.get_location(), to.get_location())
+
+    @property
+    def is_keyword(self) -> bool:
+        return self.lexeme in KEYWORDS

midas/parser/base.py

@@ -2,8 +2,8 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Generic, TypeVar
 
-from lexer.token import Token, TokenType
-from parser.errors import ParsingError
+from midas.lexer.token import Token, TokenType
+from midas.parser.errors import ParsingError
 
 
 @dataclass(frozen=True)

midas/parser/midas.py

@@ -0,0 +1,539 @@
+from typing import Optional
+
+from midas.ast.location import Location
+from midas.ast.midas import (
+    BinaryExpr,
+    ComplexType,
+    ConstraintType,
+    Expr,
+    ExtendStmt,
+    ExtensionType,
+    FunctionType,
+    GenericType,
+    GetExpr,
+    GroupingExpr,
+    LiteralExpr,
+    LogicalExpr,
+    MemberKind,
+    MemberStmt,
+    NamedType,
+    PredicateStmt,
+    Stmt,
+    Type,
+    TypeParam,
+    TypeStmt,
+    UnaryExpr,
+    VariableExpr,
+    WildcardExpr,
+)
+from midas.lexer.token import KEYWORDS, Token, TokenType
+from midas.parser.base import Parser
+from midas.parser.errors import ParsingError
+
+
+class MidasParser(Parser):
+    """A simple parser for midas type definitions"""
+
+    SYNC_BOUNDARY: set[TokenType] = {
+        TokenType.TYPE,
+        TokenType.EXTEND,
+        TokenType.PREDICATE,
+        TokenType.PROP,
+        TokenType.FUNC,
+    }
+
+    def parse(self) -> list[Stmt]:
+        statements: list[Stmt] = []
+        while not self.is_at_end():
+            stmt: Optional[Stmt] = self.declaration()
+            if stmt is None:
+                print("Early stop")
+                break
+            statements.append(stmt)
+        return statements
+
+    def synchronize(self):
+        """Skip tokens until a synchronization boundary is found
+
+        This method allows gracefully recovering from a parse error
+        to a safe place and continue parsing
+        """
+        self.advance()
+        while not self.is_at_end():
+            if self.previous().type == TokenType.NEWLINE:
+                return
+            if self.peek().type in self.SYNC_BOUNDARY:
+                return
+            self.advance()
+
+    def declaration(self) -> Optional[Stmt]:
+        """Try and parse a declaration
+
+        Any parsing error is caught and None is returned
+
+        Returns:
+            Optional[Stmt]: the parsed Midas statement, or None if a ParsingError was raised
+        """
+        try:
+            if self.match(TokenType.TYPE):
+                return self.type_declaration()
+            if self.match(TokenType.EXTEND):
+                return self.extend_declaration()
+            if self.match(TokenType.PREDICATE):
+                return self.predicate_declaration()
+            raise self.error(self.peek(), "Unexpected token")
+        except ParsingError:
+            self.synchronize()
+            return None
+
+    def type_declaration(self) -> TypeStmt:
+        """Parse a type declaration
+
+        A type declaration can either be a simple type alias or a new complex type.
+        In either case, it can have an optional template expression after its name, wrapped in brackets.
+        A simple type alias is derived from a base type expression, and can have a optional constraint expression preceded by the `where` keyword.
+        A full simple type alias is thus written:
+        ```
+        type Name[Template](TypeExpr) where Condition
+        ```
+
+        A new complex type has a set of properties which are named, have a type and an optional constraint expression (also preceded by the `where` keyword).
+        A full complex type definition is thus written:
+        ```
+        type Name[Template] {
+            prop1: TypeExpr1 where Condition1
+            prop2: TypeExpr2 where Condition2
+            ...
+        }
+        ```
+
+        Returns:
+            TypeStmt: the parsed type declaration statement
+        """
+        keyword: Token = self.previous()
+        name: Token = self.consume_identifier("Expected type name")
+        params: list[TypeParam] = self.type_params()
+
+        self.consume(TokenType.EQUAL, "Expected '=' before type definition")
+
+        type: Type = self.type_expr()
+
+        return TypeStmt(
+            location=keyword.location_to(self.previous()),
+            name=name,
+            params=params,
+            type=type,
+        )
+
+    def type_params(self) -> list[TypeParam]:
+        """Parse a list of type parameters
+
+        Type parameters are a comma-separated list of type variables wrapped in brackets.
+        Each type variable is either a simple variable, or a bounded variable written `S <: T`
+
+        Returns:
+            list[TypeParam]: the list of type parameters, if any, or an empty list
+        """
+        if not self.match(TokenType.LEFT_BRACKET):
+            return []
+
+        params: list[TypeParam] = []
+        while not self.is_at_end() and not self.check(TokenType.RIGHT_BRACKET):
+            name: Token = self.consume_identifier("Expected type variable")
+            bound: Optional[Type] = None
+            if self.match(TokenType.LESS):
+                self.consume(TokenType.COLON, "Expected ':' after '<'")
+                bound = self.type_expr()
+            params.append(
+                TypeParam(
+                    location=name.location_to(self.previous()),
+                    name=name,
+                    bound=bound,
+                )
+            )
+            if not self.match(TokenType.COMMA):
+                break
+        self.consume(TokenType.RIGHT_BRACKET, "Missing ']' after type parameters")
+        return params
+
+    def type_expr(self) -> Type:
+        """Parse a type expression
+
+        A type is an identifier, optionally followed by a template expression.
+        It can also optionally be followed by a '?' to indicate a nullable type
+
+        Returns:
+            TypeExpr: the parsed type expression
+        """
+        base: Type
+        if self.match(TokenType.FUNC):
+            base = self.function()
+        else:
+            base = self.constraint_type()
+        if self.match(TokenType.AND):
+            extension: ComplexType = self.complex_type()
+            return ExtensionType(
+                location=Location.span(base.location, extension.location),
+                base=base,
+                extension=extension,
+            )
+        return base
+
+    def constraint_type(self) -> Type:
+        type: Type = self.base_type()
+        if self.match(TokenType.WHERE):
+            constraint: Expr = self.constraint()
+            return ConstraintType(
+                location=Location.span(type.location, constraint.location),
+                type=type,
+                constraint=constraint,
+            )
+        return type
+
+    def base_type(self) -> Type:
+        if self.match(TokenType.LEFT_PAREN):
+            type: Type = self.type_expr()
+            self.consume(TokenType.RIGHT_PAREN, "Unclosed parenthesis")
+            return type
+
+        if self.check(TokenType.LEFT_BRACE):
+            return self.complex_type()
+
+        return self.generic_type()
+
+    def generic_type(self) -> Type:
+        type: Type = self.named_type()
+        if self.check(TokenType.LEFT_BRACKET):
+            args: list[Type] = self.type_args()
+            return GenericType(
+                location=Location.span(type.location, self.previous().get_location()),
+                type=type,
+                args=args,
+            )
+        return type
+
+    def type_args(self) -> list[Type]:
+        args: list[Type] = []
+        self.consume(TokenType.LEFT_BRACKET, "Missing '[' before generic arguments")
+        while not self.is_at_end() and not self.check(TokenType.RIGHT_BRACKET):
+            args.append(self.type_expr())
+            if not self.match(TokenType.COMMA):
+                break
+        self.consume(TokenType.RIGHT_BRACKET, "Missing ']' after generic arguments")
+        return args
+
+    def named_type(self) -> Type:
+        name: Token = self.consume_identifier("Expected type name")
+        return NamedType(
+            location=name.get_location(),
+            name=name,
+        )
+
+    def complex_type(self) -> ComplexType:
+        """Parse a type definition body
+
+        A type definition body is a set of whitespace-separated
+        property statements enclosed in curly braces
+
+        Returns:
+            ComplexType: the parsed complex type
+        """
+        left: Token = self.consume(
+            TokenType.LEFT_BRACE, "Expected '{' to start type body"
+        )
+        members: list[MemberStmt] = []
+        # TODO: add keyword to differentiate properties and methods,
+        # and allow multiple methods with the same name but not properties
+        names: set[str] = set()
+        while not self.check(TokenType.RIGHT_BRACE) and not self.is_at_end():
+            member: MemberStmt = self.member_stmt()
+            # if member.name.lexeme in names:
+            #    raise self.error(member.name, "Duplicate property")
+            # names.add(member.name.lexeme)
+            members.append(member)
+        right: Token = self.consume(TokenType.RIGHT_BRACE, "Unclosed type body")
+        return ComplexType(
+            location=left.location_to(right),
+            members=members,
+        )
+
+    def constraint(self) -> Expr:
+        """Parse a constraint
+
+        A constraint is basically a logical predicate
+
+        Returns:
+            Expr: the parsed constraint expression
+        """
+        return self.and_()
+
+    def and_(self) -> Expr:
+        """Parse a logical AND expression or a simpler expression
+
+        Returns:
+            Expr: the parsed expression
+        """
+        expr: Expr = self.equality()
+        while self.match(TokenType.AND):
+            operator: Token = self.previous()
+            right: Expr = self.equality()
+            location: Location = Location.span(expr.location, right.location)
+            expr = LogicalExpr(
+                location=location, left=expr, operator=operator, right=right
+            )
+        return expr
+
+    def equality(self) -> Expr:
+        """Parse a logical equality expression or a simpler expression
+
+        Returns:
+            Expr: the parsed expression
+        """
+        expr: Expr = self.comparison()
+        while self.match(TokenType.BANG_EQUAL, TokenType.EQUAL_EQUAL):
+            operator: Token = self.previous()
+            right: Expr = self.comparison()
+            location: Location = Location.span(expr.location, right.location)
+            expr = BinaryExpr(
+                location=location, left=expr, operator=operator, right=right
+            )
+        return expr
+
+    def comparison(self) -> Expr:
+        """Parse a logical comparison expression or a simpler expression
+
+        Returns:
+            Expr: the parsed expression
+        """
+        expr: Expr = self.unary()
+        while self.match(
+            TokenType.LESS,
+            TokenType.LESS_EQUAL,
+            TokenType.GREATER,
+            TokenType.GREATER_EQUAL,
+        ):
+            operator: Token = self.previous()
+            right: Expr = self.unary()
+            location: Location = Location.span(expr.location, right.location)
+            expr = BinaryExpr(
+                location=location, left=expr, operator=operator, right=right
+            )
+        return expr
+
+    def unary(self) -> Expr:
+        """Parse a unary expression or a simpler expression
+
+        Returns:
+            Expr: the parsed expression
+        """
+        if self.match(TokenType.MINUS):
+            operator: Token = self.previous()
+            right: Expr = self.unary()
+            location: Location = Location.span(operator.get_location(), right.location)
+            return UnaryExpr(location=location, operator=operator, right=right)
+        return self.reference()
+
+    def reference(self) -> Expr:
+        """Parse an attribute access expression or a simpler expression
+
+        Returns:
+            Expr: the parsed expression
+        """
+        expr: Expr = self.primary()
+        while self.match(TokenType.DOT):
+            name: Token = self.consume_identifier("Expected property name after '.'")
+            location: Location = Location.span(expr.location, name.get_location())
+            expr = GetExpr(location=location, expr=expr, name=name)
+        return expr
+
+    def primary(self) -> Expr:
+        """Parse a primary expression
+
+        This includes literals (booleans, numbers, etc.), wildcards, identifiers and grouped expressions
+
+        Returns:
+            Expr: the parsed expression
+        """
+        token: Token = self.peek()
+        if self.match(TokenType.FALSE):
+            return LiteralExpr(location=token.get_location(), value=False)
+        if self.match(TokenType.TRUE):
+            return LiteralExpr(location=token.get_location(), value=True)
+        if self.match(TokenType.NONE):
+            return LiteralExpr(location=token.get_location(), value=None)
+
+        if self.match(TokenType.NUMBER):
+            return LiteralExpr(location=token.get_location(), value=token.value)
+
+        if self.match_identifier():
+            return VariableExpr(location=token.get_location(), name=token)
+
+        if self.match(TokenType.UNDERSCORE):
+            return WildcardExpr(location=token.get_location(), token=token)
+
+        if self.match(TokenType.LEFT_PAREN):
+            expr: Expr = self.constraint()
+            right: Token = self.consume(TokenType.RIGHT_PAREN, "Unclosed parenthesis")
+            return GroupingExpr(location=token.location_to(right), expr=expr)
+
+        raise self.error(self.peek(), "Expected expression")
+
+    def consume_identifier(self, message: str = "Expected identifier") -> Token:
+        if not self.match_identifier():
+            raise self.error(self.peek(), message)
+        return self.previous()
+
+    def match_identifier(self) -> bool:
+        return self.match(TokenType.IDENTIFIER, *KEYWORDS.values())
+
+    def check_identifier(self) -> bool:
+        for tt in [TokenType.IDENTIFIER, *KEYWORDS.values()]:
+            if self.check(tt):
+                return True
+        return False
+
+    def member_stmt(self) -> MemberStmt:
+        """Parse a member statement
+
+        A type member statement is written `prop name: Type` or `def name: Type`
+
+        Returns:
+            MemberStmt: the parsed member statement
+        """
+        kind: MemberKind
+        if self.match(TokenType.PROP):
+            kind = MemberKind.PROPERTY
+        elif self.match(TokenType.DEF):
+            kind = MemberKind.METHOD
+        else:
+            raise self.error(self.peek(), "Expected 'prop' or 'def'")
+
+        name: Token = self.consume_identifier("Expected member name")
+        self.consume(TokenType.COLON, "Expected ':' after member name")
+
+        type: Type = self.type_expr()
+        return MemberStmt(
+            location=name.location_to(self.previous()),
+            name=name,
+            type=type,
+            kind=kind,
+        )
+
+    def extend_declaration(self) -> ExtendStmt:
+        """Parse an extension definition
+
+        An extension is written `extend Type { operations }` or `extend[S <: T, U] Type { operations }`
+
+        Returns:
+            ExtendStmt: the parsed extension statement
+        """
+        keyword: Token = self.previous()
+        name: Token = self.consume_identifier("Expected type name")
+        params: list[TypeParam] = self.type_params()
+
+        self.consume(TokenType.LEFT_BRACE, "Expected '{' to start extend body")
+        members: list[MemberStmt] = []
+        while not self.is_at_end() and not self.check(TokenType.RIGHT_BRACE):
+            members.append(self.member_stmt())
+        self.consume(TokenType.RIGHT_BRACE, "Unclosed extend body")
+        location: Location = keyword.location_to(self.previous())
+        return ExtendStmt(
+            location=location,
+            name=name,
+            params=params,
+            members=members,
+        )
+
+    def predicate_declaration(self) -> PredicateStmt:
+        """Parse a predicate declaration
+
+        A predicate is written `predicate Name(subject: Type) = constraint_expression`
+
+        Returns:
+            PredicateStmt: the parsed predicate declaration statement
+        """
+        keyword: Token = self.previous()
+        name: Token = self.consume_identifier("Expected predicate name")
+        self.consume(TokenType.LEFT_PAREN, "Expected '(' before predicate subject")
+        subject: Token = self.consume_identifier("Expected subject name")
+        self.consume(TokenType.COLON, "Expected ':' after subject name")
+        type: Type = self.type_expr()
+        self.consume(TokenType.RIGHT_PAREN, "Expected ')' after predicate subject")
+        self.consume(TokenType.EQUAL, "Expected '=' after predicate subject")
+        condition: Expr = self.constraint()
+        return PredicateStmt(
+            location=keyword.location_to(self.previous()),
+            name=name,
+            subject=subject,
+            type=type,
+            condition=condition,
+        )
+
+    def function(self) -> FunctionType:
+        l_paren: Token = self.consume(
+            TokenType.LEFT_PAREN, "Expected '(' before function parameters"
+        )
+        pos_args: list[FunctionType.Argument] = []
+        args: list[FunctionType.Argument] = []
+        kw_args: list[FunctionType.Argument] = []
+
+        args_first_tokens: list[Token] = []
+
+        section: int = 0
+        while not self.is_at_end() and not self.check(TokenType.RIGHT_PAREN):
+            match section:
+                case 0 if self.match(TokenType.SLASH):
+                    pos_args = args
+                    args = []
+                    args_first_tokens = []
+                    section = 1
+                case 0 | 1 if self.match(TokenType.STAR):
+                    section = 2
+                case _:
+                    # Record first token of mixed argument for errors if unnamed
+                    if section != 2:
+                        args_first_tokens.append(self.peek())
+
+                    name: Optional[Token] = None
+                    if section == 2:
+                        name = self.consume_identifier("Expected keyword argument name")
+                        self.consume(
+                            TokenType.COLON, "Expected ':' after argument name"
+                        )
+                    elif self.check_identifier() and self.check_next(TokenType.COLON):
+                        name = self.advance()
+                        self.advance()
+
+                    type: Type = self.type_expr()
+                    optional: bool = self.match(TokenType.QMARK)
+                    arg = FunctionType.Argument(
+                        location=None,
+                        name=name,
+                        type=type,
+                        required=not optional,
+                    )
+                    if section == 2:
+                        kw_args.append(arg)
+                    else:
+                        args.append(arg)
+
+            if not self.match(TokenType.COMMA):
+                break
+
+        for arg, token in zip(args, args_first_tokens):
+            if arg.name is None:
+                # Not raised because we can keep parsing
+                self.error(token, "Unnamed mixed argument")
+
+        self.consume(TokenType.RIGHT_PAREN, "Expected ')' after function parameters")
+
+        self.consume(TokenType.ARROW, "Expected '->' before result type")
+        result: Type = self.type_expr()
+
+        return FunctionType(
+            location=l_paren.location_to(self.previous()),
+            pos_args=pos_args,
+            args=args,
+            kw_args=kw_args,
+            returns=result,
+        )

midas/parser/python.py

@@ -0,0 +1,560 @@
+import ast
+from typing import Optional
+
+from midas.ast.location import Location
+from midas.ast.python import (
+    AssignStmt,
+    BaseType,
+    BinaryExpr,
+    CallExpr,
+    CastExpr,
+    CompareExpr,
+    ConstraintType,
+    DictExpr,
+    Expr,
+    ExpressionStmt,
+    ForStmt,
+    FrameColumn,
+    FrameType,
+    Function,
+    GetExpr,
+    IfStmt,
+    ListExpr,
+    LiteralExpr,
+    LogicalExpr,
+    MidasType,
+    RawExpr,
+    RawStmt,
+    ReturnStmt,
+    SliceExpr,
+    Stmt,
+    SubscriptExpr,
+    TernaryExpr,
+    TypeAssign,
+    UnaryExpr,
+    VariableExpr,
+)
+
+
+class InvalidSyntaxError(Exception):
+    pass
+
+
+class UnsupportedSyntaxError(Exception):
+    def __init__(self, expr: ast.expr) -> None:
+        super().__init__(
+            f"Unsupported syntax at L{expr.lineno}:{expr.col_offset}: {ast.unparse(expr)}"
+        )
+
+
+class PythonParser:
+    CAST_FUNCTION = "cast"
+
+    def parse_module(self, node: ast.Module) -> list[Stmt]:
+        statements: list[Stmt] = []
+        for stmt in node.body:
+            try:
+                parsed: None | Stmt | list[Stmt] = self.parse_stmt(stmt)
+                if isinstance(parsed, Stmt):
+                    statements.append(parsed)
+                elif parsed is not None:
+                    statements.extend(parsed)
+            except UnsupportedSyntaxError as e:
+                print(f"{e}, skipping")
+                continue
+        return statements
+
+    def parse_stmt(self, node: ast.stmt) -> None | Stmt | list[Stmt]:
+        location: Location = Location.from_ast(node)
+        match node:
+            case ast.AnnAssign():
+                return self.parse_annotation_assign(node)
+
+            case ast.Assign():
+                return self.parse_assign(node)
+
+            case ast.AugAssign():
+                return self.parse_aug_assign(node)
+
+            case ast.FunctionDef():
+                return self.parse_function(node)
+
+            case ast.Expr(value=expr):
+                return ExpressionStmt(
+                    location=location,
+                    expr=self.parse_expr(expr),
+                )
+
+            case ast.Return(value=value):
+                return ReturnStmt(
+                    location=location,
+                    value=self.parse_expr(value) if value is not None else None,
+                )
+
+            case ast.If():
+                return self.parse_if(node)
+
+            case ast.Pass():
+                return None
+
+            case ast.For(orelse=[]):
+                return self.parse_for(node)
+
+            case _:
+                print(f"Unsupported statement: {ast.unparse(node)}")
+                return RawStmt(location=location, stmt=node)
+
+    def parse_annotation_assign(self, node: ast.AnnAssign) -> list[Stmt]:
+        statements: list[Stmt] = []
+        loc: Location = Location.from_ast(node)
+        match node:
+            case ast.AnnAssign(
+                target=ast.Name(id=target),
+                annotation=annotation,
+                value=value,
+                simple=1,
+            ):
+                type = self._parse_type(annotation)
+                statements.append(
+                    TypeAssign(
+                        location=loc,
+                        name=target,
+                        type=type,
+                    )
+                )
+
+                if value is not None:
+                    statements.append(
+                        AssignStmt(
+                            location=loc,
+                            targets=[
+                                VariableExpr(
+                                    location=Location.from_ast(node.target), name=target
+                                ),
+                            ],
+                            value=self.parse_expr(value),
+                        ),
+                    )
+            case _:
+                print(f"Unsupported annotation: {ast.unparse(node)}")
+        return statements
+
+    def parse_assign(self, node: ast.Assign) -> AssignStmt:
+        targets: list[Expr] = []
+        for target in node.targets:
+            targets.append(self.parse_expr(target))
+        value: Expr = self.parse_expr(node.value)
+        return AssignStmt(
+            location=Location.from_ast(node),
+            targets=targets,
+            value=value,
+        )
+
+    def parse_aug_assign(self, node: ast.AugAssign) -> AssignStmt:
+        location: Location = Location.from_ast(node)
+        target: Expr = self.parse_expr(node.target)
+        value: Expr = self.parse_expr(node.value)
+        return AssignStmt(
+            location=location,
+            targets=[target],
+            value=BinaryExpr(
+                location=location,
+                left=target,
+                operator=node.op,
+                right=value,
+            ),
+        )
+
+    def parse_if(self, node: ast.If) -> IfStmt:
+        body: list[Stmt] = []
+        for stmt in node.body:
+            stmts = self.parse_stmt(stmt)
+            if isinstance(stmts, Stmt):
+                body.append(stmts)
+            elif stmts is not None:
+                body.extend(stmts)
+
+        orelse: list[Stmt] = []
+        for stmt in node.orelse:
+            stmts = self.parse_stmt(stmt)
+            if isinstance(stmts, Stmt):
+                orelse.append(stmts)
+            elif stmts is not None:
+                orelse.extend(stmts)
+
+        return IfStmt(
+            location=Location.from_ast(node),
+            test=self.parse_expr(node.test),
+            body=body,
+            orelse=orelse,
+        )
+
+    def parse_for(self, node: ast.For) -> ForStmt:
+        body: list[Stmt] = []
+        for stmt in node.body:
+            stmts = self.parse_stmt(stmt)
+            if isinstance(stmts, Stmt):
+                body.append(stmts)
+            elif stmts is not None:
+                body.extend(stmts)
+
+        return ForStmt(
+            location=Location.from_ast(node),
+            target=self.parse_expr(node.target),
+            iterator=self.parse_expr(node.iter),
+            body=body,
+        )
+
+    def parse_function(self, node: ast.FunctionDef) -> Function:
+        loc: Location = Location.from_ast(node)
+        match node:
+            case ast.FunctionDef(
+                name=name,
+                args=ast.arguments(
+                    posonlyargs=posonlyargs,
+                    args=args,
+                    vararg=sink,
+                    kwonlyargs=kwonlyargs,
+                    kwarg=kw_sink,
+                    defaults=defaults,
+                    kw_defaults=kw_defaults,
+                ),
+                returns=returns,
+                body=raw_body,
+            ):
+
+                def parse_args(
+                    args_list: list[ast.arg], defaults: list[Optional[Expr]]
+                ) -> list[Function.Argument]:
+                    return [
+                        self._parse_function_argument(arg, default)
+                        for arg, default in zip(args_list, defaults)
+                    ]
+
+                body: list[Stmt] = []
+                for stmt in raw_body:
+                    stmts = self.parse_stmt(stmt)
+                    if isinstance(stmts, Stmt):
+                        body.append(stmts)
+                    elif stmts is not None:
+                        body.extend(stmts)
+
+                parsed_defaults: list[Optional[Expr]] = [
+                    self.parse_expr(default) for default in defaults
+                ]
+                n_posargs: int = len(posonlyargs)
+                n_args: int = len(args)
+                n_all_posargs = n_posargs + n_args
+                parsed_defaults = [
+                    None,
+                ] * (n_all_posargs - len(defaults)) + parsed_defaults
+
+                posargs_defaults: list[Optional[Expr]] = parsed_defaults[:n_posargs]
+                args_defaults: list[Optional[Expr]] = parsed_defaults[n_posargs:]
+                kwargs_defaults: list[Optional[Expr]] = [
+                    self.parse_expr(default) if default is not None else None
+                    for default in kw_defaults
+                ]
+
+                return Function(
+                    location=loc,
+                    name=name,
+                    posonlyargs=parse_args(posonlyargs, posargs_defaults),
+                    args=parse_args(args, args_defaults),
+                    sink=(
+                        self._parse_function_argument(sink, None)
+                        if sink is not None
+                        else None
+                    ),
+                    kwonlyargs=parse_args(kwonlyargs, kwargs_defaults),
+                    kw_sink=(
+                        self._parse_function_argument(kw_sink, None)
+                        if kw_sink is not None
+                        else None
+                    ),
+                    returns=self._parse_type(returns) if returns is not None else None,
+                    body=body,
+                )
+            case _:
+                print(f"Unsupported function definition: {ast.unparse(node)}")
+
+    def _parse_function_argument(
+        self, arg: ast.arg, default: Optional[Expr]
+    ) -> Function.Argument:
+        loc: Location = Location.from_ast(arg)
+        name: str = arg.arg
+        type: Optional[MidasType] = None
+        if arg.annotation is not None:
+            type = self._parse_type(arg.annotation)
+        return Function.Argument(
+            location=loc,
+            name=name,
+            type=type,
+            default=default,
+        )
+
+    def _parse_type(self, type_expr: ast.expr) -> MidasType:
+        loc: Location = Location.from_ast(type_expr)
+        match type_expr:
+            case ast.Subscript(value=ast.Name(id="Frame"), slice=schema):
+                return self._parse_frame_type(schema)
+
+            case ast.Subscript(value=ast.Name(id=name), slice=param):
+                return BaseType(
+                    location=loc,
+                    base=name,
+                    param=self._parse_type(param),
+                )
+
+            case ast.Name(id=name):
+                return BaseType(
+                    location=loc,
+                    base=name,
+                    param=None,
+                )
+
+            case ast.BinOp(left=left_expr, op=ast.Add(), right=right_expr):
+                left = self._parse_type(left_expr)
+                match left:
+                    case None:
+                        raise InvalidSyntaxError()
+
+                    # If chained constraints, separate base type and rebuild constraint
+                    case ConstraintType(type=left_type, constraint=left_constraint):
+                        constraint = ast.BinOp(
+                            left=left_constraint,
+                            op=ast.Add(),
+                            right=right_expr,
+                        )
+                        ast.copy_location(constraint, type_expr)
+                        return ConstraintType(
+                            location=loc,
+                            type=left_type,
+                            constraint=constraint,
+                        )
+
+                    case _:
+                        return ConstraintType(
+                            location=loc,
+                            type=left,
+                            constraint=right_expr,
+                        )
+
+            case ast.Constant(value=None):
+                return BaseType(
+                    location=loc,
+                    base="None",
+                    param=None,
+                )
+
+            case _:
+                raise UnsupportedSyntaxError(type_expr)
+
+    def _parse_frame_type(self, schema: ast.expr) -> FrameType:
+        loc: Location = Location.from_ast(schema)
+        columns: list[FrameColumn] = []
+
+        match schema:
+            case ast.Tuple(elts=cols):
+                for col in cols:
+                    columns.append(self._parse_frame_column(col))
+
+            case ast.Slice() | ast.Name():
+                columns.append(self._parse_frame_column(schema))
+
+            case _:
+                raise UnsupportedSyntaxError(schema)
+
+        return FrameType(location=loc, columns=columns)
+
+    def _parse_frame_column(self, column: ast.expr) -> FrameColumn:
+        loc: Location = Location.from_ast(column)
+        match column:
+            case ast.Name():
+                return FrameColumn(
+                    location=loc,
+                    name=None,
+                    type=self._parse_type(column),
+                )
+
+            case ast.Slice(lower=ast.Name(id=name), upper=type_expr):
+                if name == "_":
+                    name = None
+
+                type: Optional[MidasType] = None
+                match type_expr:
+                    case None:
+                        raise InvalidSyntaxError("Missing column type")
+                    case ast.Name(id="_"):
+                        type = None
+                    case ast.expr():
+                        type = self._parse_type(type_expr)
+                    case _:
+                        raise UnsupportedSyntaxError(type_expr)
+                return FrameColumn(location=loc, name=name, type=type)
+
+            case _:
+                raise UnsupportedSyntaxError(column)
+
+    def parse_expr(self, node: ast.expr) -> Expr:
+        location: Location = Location.from_ast(node)
+        match node:
+            case ast.BoolOp():
+                return self.parse_bool_op(node)
+
+            case ast.BinOp(left=left, op=op, right=right):
+                return BinaryExpr(
+                    location=location,
+                    left=self.parse_expr(left),
+                    operator=op,
+                    right=self.parse_expr(right),
+                )
+
+            case ast.UnaryOp(op=op, operand=right):
+                return UnaryExpr(
+                    location=location,
+                    operator=op,
+                    right=self.parse_expr(right),
+                )
+
+            case ast.Compare():
+                return self.parse_compare(node)
+
+            case ast.Call(func=ast.Name(id=self.CAST_FUNCTION)):
+                return self.parse_cast(node)
+
+            case ast.Call():
+                return self.parse_call(node)
+
+            case ast.IfExp():
+                return self.parse_ternary(node)
+
+            case ast.Constant(value=value):
+                return LiteralExpr(location=location, value=value)
+
+            case ast.Attribute(value=object, attr=name):
+                return GetExpr(
+                    location=location,
+                    object=self.parse_expr(object),
+                    name=name,
+                )
+
+            case ast.Name(id=name):
+                return VariableExpr(location=location, name=name)
+
+            case ast.List(elts=items):
+                return ListExpr(
+                    location=location,
+                    items=[self.parse_expr(item) for item in items],
+                )
+
+            case ast.Dict(keys=keys, values=values):
+                return DictExpr(
+                    location=location,
+                    keys=[
+                        self.parse_expr(key) if key is not None else None
+                        for key in keys
+                    ],
+                    values=[self.parse_expr(value) for value in values],
+                )
+
+            case ast.Subscript(value=value, slice=index):
+                return SubscriptExpr(
+                    location=location,
+                    object=self.parse_expr(value),
+                    index=self.parse_expr(index),
+                )
+
+            case ast.Slice(lower=lower, upper=upper, step=step):
+                return SliceExpr(
+                    location=location,
+                    lower=self.parse_expr(lower) if lower is not None else None,
+                    upper=self.parse_expr(upper) if upper is not None else None,
+                    step=self.parse_expr(step) if step is not None else None,
+                )
+
+            case _:
+                print(f"Unsupported expression: {ast.unparse(node)}")
+                return RawExpr(location=location, expr=node)
+
+    def parse_bool_op(self, node: ast.BoolOp) -> LogicalExpr:
+        op: ast.boolop = node.op
+        rights: list[Expr] = [self.parse_expr(expr) for expr in node.values]
+        expr: LogicalExpr = LogicalExpr(
+            location=Location.span(
+                rights[0].location,
+                rights[1].location,
+            ),
+            left=rights[0],
+            operator=op,
+            right=rights[1],
+        )
+        for right in rights[2:]:
+            expr = LogicalExpr(
+                location=Location.span(expr.location, right.location),
+                left=expr,
+                operator=op,
+                right=right,
+            )
+        return expr
+
+    def parse_compare(self, node: ast.Compare) -> Expr:
+        ops: list[ast.cmpop] = node.ops
+        left: Expr = self.parse_expr(node.left)
+        rights: list[Expr] = [self.parse_expr(expr) for expr in node.comparators]
+        expr: Expr = CompareExpr(
+            location=Location.span(
+                left.location,
+                rights[0].location,
+            ),
+            left=left,
+            operator=ops[0],
+            right=rights[0],
+        )
+        for i, right in enumerate(rights[1:]):
+            comparison = CompareExpr(
+                location=Location.span(rights[i].location, right.location),
+                left=rights[i],
+                operator=ops[i],
+                right=right,
+            )
+            expr = LogicalExpr(
+                location=Location.span(expr.location, comparison.location),
+                left=expr,
+                operator=ast.And(),
+                right=comparison,
+            )
+        return expr
+
+    def parse_cast(self, node: ast.Call) -> CastExpr:
+        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),
+                )
+            case _:
+                raise InvalidSyntaxError(
+                    f"Invalid call to {self.CAST_FUNCTION}, expected type and expression"
+                )
+
+    def parse_call(self, node: ast.Call) -> CallExpr:
+        return CallExpr(
+            location=Location.from_ast(node),
+            callee=self.parse_expr(node.func),
+            arguments=[self.parse_expr(arg) for arg in node.args],
+            keywords={
+                arg.arg: self.parse_expr(arg.value)
+                for arg in node.keywords
+                if arg.arg is not None  # Should always be True, type checker happy
+            },
+        )
+
+    def parse_ternary(self, node: ast.IfExp) -> TernaryExpr:
+        return TernaryExpr(
+            location=Location.from_ast(node),
+            test=self.parse_expr(node.test),
+            if_true=self.parse_expr(node.body),
+            if_false=self.parse_expr(node.orelse),
+        )

midas/utils.py

@@ -0,0 +1,64 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Optional
+
+import midas.ast.python as p
+from midas.checker.types import Type
+
+AllowRepeat = Callable[[object], bool]
+
+
+class UniversalJSONDumper:
+    @classmethod
+    def dump(
+        cls,
+        obj: Any,
+        include_keys: Optional[list[str | tuple[str, str]]] = None,
+        allow_repeat: Optional[AllowRepeat] = None,
+    ) -> Any:
+        if include_keys is None:
+            include_keys = []
+        return cls._dump(obj, include_keys, allow_repeat, [])
+
+    @classmethod
+    def _dump(
+        cls,
+        obj: Any,
+        include_keys: list[str | tuple[str, str]],
+        allow_repeat: Optional[AllowRepeat],
+        visited: list[Any],
+    ) -> Any:
+        if obj in visited:
+            return None
+        match obj:
+            case str() | int() | float() | None:
+                return obj
+            case list() | set() | tuple():
+                return [
+                    cls._dump(child, include_keys, allow_repeat, visited)
+                    for child in obj
+                ]
+            case dict():
+                return {
+                    str(k): cls._dump(v, include_keys, allow_repeat, visited)
+                    for k, v in obj.items()
+                }
+            case object():
+                if allow_repeat is None or not allow_repeat(obj):
+                    visited.append(obj)
+                return {
+                    "_type": obj.__class__.__name__,
+                } | {
+                    k: cls._dump(v, include_keys, allow_repeat, visited)
+                    for k, v in obj.__dict__.items()
+                    if not k.startswith("_")
+                    or k in include_keys
+                    or (obj.__class__.__name__, k) in include_keys
+                }
+            case _:
+                raise ValueError(f"Unsupported value: {obj}")
+
+
+@dataclass(frozen=True, kw_only=True)
+class TypedAST:
+    stmts: list[p.Stmt]
+    judgements: list[tuple[p.Expr, Type]]

parser/annotations.py

@@ -1,152 +0,0 @@
-from typing import Optional
-
-from core.ast.annotations import (
-    AnnotationStmt,
-    ConstraintExpr,
-    Expr,
-    LiteralExpr,
-    SchemaElementExpr,
-    SchemaExpr,
-    Stmt,
-    TypeExpr,
-    WildcardExpr,
-)
-from lexer.token import Token, TokenType
-from parser.base import Parser
-from parser.errors import ParsingError
-
-
-class AnnotationParser(Parser):
-    """A simple parser for custom type annotations"""
-
-    SYNC_BOUNDARY: set[TokenType] = set()
-
-    def parse(self) -> Optional[Stmt]:
-        stmt: Optional[Stmt] = None
-        try:
-            stmt = self.annotation()
-        except ParsingError:
-            self.synchronize()
-        if not self.is_at_end():
-            self.error(self.peek(), "Extra tokens")
-        return stmt
-
-    def synchronize(self):
-        """Skip tokens until a synchronization boundary is found
-
-        This method allows gracefully recovering from a parse error
-        to a safe place and continue parsing
-        """
-        self.advance()
-        while not self.is_at_end():
-            if self.peek().type in self.SYNC_BOUNDARY:
-                return
-            self.advance()
-
-    def annotation(self) -> AnnotationStmt:
-        """Parse an annotation
-
-        An annotation is written as `Type` or `Type[Schema]`
-
-        Returns:
-            AnnotationStmt: the parsed annotation statement
-        """
-
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected type identifier")
-        schema: Optional[SchemaExpr] = None
-        if self.match(TokenType.LEFT_BRACKET):
-            schema = self.schema()
-        return AnnotationStmt(name=name, schema=schema)
-
-    def type_expr(self) -> TypeExpr:
-        """Parse a type expression
-
-        Returns:
-            TypeExpr: the parsed type expression
-        """
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected type name")
-        constraints: list[ConstraintExpr] = []
-
-        while not self.is_at_end() and self.match(TokenType.PLUS):
-            self.consume(TokenType.LEFT_PAREN, "Expected '(' before type constraint")
-            constraints.append(self.constraint_expr())
-            self.consume(TokenType.RIGHT_PAREN, "Expected ')' after type constraint")
-
-        return TypeExpr(name=name, constraints=constraints)
-
-    def constraint_expr(self) -> ConstraintExpr:
-        """Parse a type constraint
-
-        Returns:
-            ConstraintExpr: the parsed type constraint expression
-        """
-        
-        left: Expr = self.constraint_value()
-        op: Token = self.constraint_operator()
-        right: Expr = self.constraint_value()
-        return ConstraintExpr(left=left, op=op, right=right)
-
-    def constraint_value(self) -> Expr:
-        if self.match(TokenType.UNDERSCORE):
-            return WildcardExpr(self.previous())
-        return self.literal()
-
-    def literal(self) -> LiteralExpr:
-        if self.match(TokenType.FALSE):
-            return LiteralExpr(False)
-        if self.match(TokenType.TRUE):
-            return LiteralExpr(True)
-        if self.match(TokenType.NONE):
-            return LiteralExpr(None)
-        
-        if self.match(TokenType.NUMBER):
-            return LiteralExpr(self.previous().value)
-        
-        raise self.error(self.peek(), "Expected literal")
-
-    def constraint_operator(self) -> Token:
-        if self.match(TokenType.LESS, TokenType.LESS_EQUAL, TokenType.GREATER, TokenType.GREATER_EQUAL, TokenType.EQUAL_EQUAL, TokenType.BANG_EQUAL):
-            return self.previous()
-        raise self.error(self.peek(), "Expected constraint operator")
-
-    def schema(self) -> SchemaExpr:
-        """Parse a schema definition
-
-        A comma separated list of schema elements
-
-        Returns:
-            SchemaExpr: the parsed schema expression
-        """
-        left: Token = self.previous()
-        elements: list[Expr] = []
-        while not self.check(TokenType.RIGHT_BRACKET) and not self.is_at_end():
-            elements.append(self.schema_element())
-            if not self.check(TokenType.RIGHT_BRACKET):
-                self.consume(TokenType.COMMA, "Expected ',' between schema elements")
-
-        right: Token = self.consume(TokenType.RIGHT_BRACKET, "Unclosed schema")
-        return SchemaExpr(left=left, elements=elements, right=right)
-
-    def schema_element(self) -> SchemaElementExpr:
-        """Parse a schema element
-
-        An anonymous element (`_`), a type, an untyped named column (`name: _`),
-        or a named column (`name: Type`)
-
-        Returns:
-            SchemaElementExpr: the parsed schema element expression
-        """
-        if self.match(TokenType.UNDERSCORE):
-            return SchemaElementExpr(name=None, type=None)
-
-        if not self.check(TokenType.IDENTIFIER):
-            raise self.error(self.peek(), "Expected schema element")
-
-        name: Optional[Token] = None
-        type: Optional[TypeExpr] = None
-        if self.check_next(TokenType.COLON):
-            name = self.advance()
-            self.advance()
-        if not self.match(TokenType.UNDERSCORE):
-            type = self.type_expr()
-        return SchemaElementExpr(name=name, type=type)

parser/midas.py

@@ -1,217 +0,0 @@
-from typing import Optional
-
-from core.ast.midas import (
-    ConstraintExpr,
-    ConstraintStmt,
-    Expr,
-    LiteralExpr,
-    OpStmt,
-    PropertyStmt,
-    Stmt,
-    TypeBodyExpr,
-    TypeExpr,
-    TypeStmt,
-    WildcardExpr,
-)
-from lexer.token import Token, TokenType
-from parser.base import Parser
-from parser.errors import ParsingError
-
-
-class MidasParser(Parser):
-    """A simple parser for midas type definitions"""
-
-    SYNC_BOUNDARY: set[TokenType] = {TokenType.TYPE, TokenType.OP, TokenType.CONSTRAINT}
-
-    def parse(self) -> list[Stmt]:
-        statements: list[Stmt] = []
-        while not self.is_at_end():
-            stmt: Optional[Stmt] = self.declaration()
-            if stmt is None:
-                print("Early stop")
-                break
-            statements.append(stmt)
-        return statements
-
-    def synchronize(self):
-        """Skip tokens until a synchronization boundary is found
-
-        This method allows gracefully recovering from a parse error
-        to a safe place and continue parsing
-        """
-        self.advance()
-        while not self.is_at_end():
-            if self.previous().type == TokenType.NEWLINE:
-                return
-            if self.peek().type in self.SYNC_BOUNDARY:
-                return
-            self.advance()
-
-    def declaration(self) -> Optional[Stmt]:
-        """Try and parse a declaration
-
-        Any parsing error is caught and None is returned
-
-        Returns:
-            Optional[Stmt]: the parsed Midas statement, or None if a ParsingError was raised
-        """
-        try:
-            if self.match(TokenType.TYPE):
-                return self.type_declaration()
-            if self.match(TokenType.OP):
-                return self.op_declaration()
-            if self.match(TokenType.CONSTRAINT):
-                return self.constraint_declaration()
-            raise self.error(self.peek(), "Unexpected token")
-        except ParsingError:
-            self.synchronize()
-            return None
-
-    def type_declaration(self) -> TypeStmt:
-        """Parse a type declaration
-
-        A type declaration is written `type Name<TypeExpr, ...>` optionally followed by a brace-wrapped body
-
-        Returns:
-            TypeStmt: the parsed type declaration statement
-        """
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected type name")
-        self.consume(TokenType.LESS, "Expected '<' after type name")
-        bases: list[TypeExpr] = []
-        while not self.check(TokenType.GREATER) and not self.is_at_end():
-            bases.append(self.type_expr())
-            if not self.check(TokenType.GREATER):
-                self.consume(TokenType.COMMA, "Expected ',' between type bases")
-        self.consume(TokenType.GREATER, "Expected '>' after base type")
-
-        body: Optional[TypeBodyExpr] = None
-
-        if self.check(TokenType.LEFT_BRACE):
-            body = self.type_body_expr()
-        return TypeStmt(name=name, bases=bases, body=body)
-
-    def type_expr(self) -> TypeExpr:
-        """Parse a type expression
-
-        Returns:
-            TypeExpr: the parsed type expression
-        """
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected type name")
-        constraints: list[ConstraintExpr] = []
-
-        while not self.is_at_end() and self.match(TokenType.PLUS):
-            self.consume(TokenType.LEFT_PAREN, "Expected '(' before type constraint")
-            constraints.append(self.constraint_expr())
-            self.consume(TokenType.RIGHT_PAREN, "Expected ')' after type constraint")
-
-        return TypeExpr(name=name, constraints=constraints)
-
-    def constraint_expr(self) -> ConstraintExpr:
-        """Parse a type constraint
-
-        Returns:
-            ConstraintExpr: the parsed type constraint expression
-        """
-
-        left: Expr = self.constraint_value()
-        op: Token = self.constraint_operator()
-        right: Expr = self.constraint_value()
-        return ConstraintExpr(left=left, op=op, right=right)
-
-    def constraint_value(self) -> Expr:
-        if self.match(TokenType.UNDERSCORE):
-            return WildcardExpr(self.previous())
-        return self.literal()
-
-    def literal(self) -> LiteralExpr:
-        if self.match(TokenType.FALSE):
-            return LiteralExpr(False)
-        if self.match(TokenType.TRUE):
-            return LiteralExpr(True)
-        if self.match(TokenType.NONE):
-            return LiteralExpr(None)
-
-        if self.match(TokenType.NUMBER):
-            return LiteralExpr(self.previous().value)
-
-        raise self.error(self.peek(), "Expected literal")
-
-    def constraint_operator(self) -> Token:
-        if self.match(
-            TokenType.LESS,
-            TokenType.LESS_EQUAL,
-            TokenType.GREATER,
-            TokenType.GREATER_EQUAL,
-            TokenType.EQUAL_EQUAL,
-            TokenType.BANG_EQUAL,
-        ):
-            return self.previous()
-        raise self.error(self.peek(), "Expected constraint operator")
-
-    def type_body_expr(self) -> TypeBodyExpr:
-        """Parse a type definition body
-
-        A type definition body is a set of whitespace-separated
-        property statements enclosed in curly braces
-
-        Returns:
-            TypeBodyExpr: the parsed type body expression
-        """
-        self.consume(TokenType.LEFT_BRACE, "Expected '{' to start type body")
-        properties: list[PropertyStmt] = []
-        while not self.check(TokenType.RIGHT_BRACE) and not self.is_at_end():
-            properties.append(self.property_stmt())
-        self.consume(TokenType.RIGHT_BRACE, "Unclosed type body")
-        return TypeBodyExpr(properties=properties)
-
-    def property_stmt(self) -> PropertyStmt:
-        """Parse a property statement
-
-        A type property statement is written `name: Type`
-
-        Returns:
-            PropertyStmt: the parsed property statement
-        """
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected property name")
-        self.consume(TokenType.COLON, "Expected ':' after property name")
-        type: TypeExpr = self.type_expr()
-        return PropertyStmt(name=name, type=type)
-
-    def op_declaration(self) -> OpStmt:
-        """Parse an operation definition
-
-        An operation is written `op <Type1> operator <Type2> = <Type3>` where `operator` can be any single token
-
-        Returns:
-            OpStmt: the parsed operation statement
-        """
-        self.consume(TokenType.LESS, "Expected '<' before first type")
-        left: TypeExpr = self.type_expr()
-        self.consume(TokenType.GREATER, "Expected '>' after first type")
-
-        op: Token = self.advance()
-
-        self.consume(TokenType.LESS, "Expected '<' before second type")
-        right: TypeExpr = self.type_expr()
-        self.consume(TokenType.GREATER, "Expected '>' after second type")
-
-        self.consume(TokenType.EQUAL, "Expected '=' after second type")
-
-        self.consume(TokenType.LESS, "Expected '<' before result type")
-        result: TypeExpr = self.type_expr()
-        self.consume(TokenType.GREATER, "Expected '>' after result type")
-
-        return OpStmt(left=left, op=op, right=right, result=result)
-
-    def constraint_declaration(self) -> ConstraintStmt:
-        """Parse a type constraint declaration
-
-        A constraint is written `constraint Name = constraint_expression`
-
-        Returns:
-            ConstraintStmt: the parsed constraint declaration statement
-        """
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected constraint name")
-        self.consume(TokenType.EQUAL, "Expected '=' after constraint name")
-        constraint: ConstraintExpr = self.constraint_expr()
-        return ConstraintStmt(name=name, constraint=constraint)

pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "midas"
+version = "0.1.0"
+description = "A static-first type checking framework for Python data-frames"
+readme = "README.md"
+requires-python = ">=3.11"
+authors = [
+    { name = "Louis Heredero", email = "louis.heredero@students.hevs.ch" },
+]
+classifiers = ["Programming Language :: Python :: 3"]
+dependencies = ["click>=8.4.1"]
+
+[project.urls]
+Homepage = "https://git.kbk28.ch/HEL/midas"
+Repository = "https://git.kbk28.ch/HEL/midas"
+
+[project.scripts]
+midas = "midas.cli.main:midas"
+
+[build-system]
+requires = ['hatchling']
+build-backend = 'hatchling.build'

syntax/midas.ebnf

@@ -1,26 +1,43 @@
-identifier ::= '[a-zA-Z][a-zA-Z_]*'
+// W3C EBNF syntax definition for Midas
+Identifier ::= [a-zA-Z_] [a-zA-Z_0-9]*
 
-integer ::= '\d+'
-number ::= integer ["." integer]
-boolean ::= "False" | "True"
-none ::= "None"
+Integer ::= '\d+'
+Number ::= "-"? Integer ("." Integer)?
+Boolean ::= "False" | "True"
+None ::= "None"
 
-value ::= number | boolean | none
-lambda-value ::= "_" | value
-lambda-operator ::= ">" | "<" | ">=" | "<=" | "==" | "!="
-lambda ::= lambda-value lambda-operator lambda-value
+Value ::= Number | Boolean | None
 
-constraint ::= identifier | "(" lambda ")"
-base-type ::= identifier
-type ::= base-type { "+" constraint }
+ComparisonOp ::= ">" | "<" | ">=" | "<="
+EqualityOp ::= "==" | "!="
 
-type-property ::= 'identifier' ":" 'type'
-type-body ::= "{" { 'type-property' } "}"
+Grouping ::= "(" Constraint ")"
+Primary ::= "_" | Value | Identifier | Grouping
+Reference ::= Primary ("." Identifier)*
+Unary ::= "-"? Unary | Reference
+Comparison ::= Unary (ComparisonOp Unary)*
+Equality ::= Comparison (EqualityOp Comparison)*
+Constraint ::= Equality ("&" Equality)*
 
-operation-type ::= "<" 'type' ">"
+TemplateParam ::= Identifier ("<:" Type)?
+Template ::= "[" (TemplateParam ("," TemplateParam)*)? "]"
 
-type-statement ::= "type" 'identifier' "<" 'type' {"," 'type'} ">" ['type-body']
-operation-statement ::= "op" 'operation-type' 'operator' 'operation-type' "=" 'operation-type'
-constraint-statement ::= "constraint" 'identifier' "=" 'lambda'
 
-statement ::= type-statement | operation-statement | constraint-statement
+TypeProperty ::= Identifier ":" Type
+ComplexType ::= "{" TypeProperty* "}"
+NamedType ::= Identifier
+TypeParams ::= "[" (Type ("," Type)*)? "]"
+GenericType ::= NamedType TypeParams?
+GroupedType ::= "(" Type ")"
+BaseType ::= GroupedType | ComplexType | GenericType
+ConstraintType ::= BaseType ("where" Constraint)?
+Type ::= ConstraintType
+
+OpDefinition ::= "op" Identifier "(" Type ")" "->" Type
+ExtendBody ::= "{" OpDefinition* "}"
+
+TypeStatement ::= "type" Identifier Template? "=" Type
+ExtendStatement ::= "extend" Type ExtendBody
+PredicateStatement ::= "predicate" Identifier "(" Identifier ":" Type ")" "=" Constraint
+
+Statement ::= TypeStatement | ExtendStatement | PredicateStatement

syntax/midas.typ

@@ -1,4 +1,11 @@
-#import "@preview/fervojo:0.1.1": render
+#import "@preview/fervojo:0.1.1": default-css, render
+
+#let extra-css = ```css
+svg.railroad .terminal rect {
+  fill: #F7DCD4;
+}
+```
+#let css = default-css() + bytes(extra-css.text)
 
 #let value = ```
 {[`value` <
@@ -8,90 +15,193 @@
 >]}
 ```
 
-#let constraint = ```
-{[`constraint` <"_", 'value'> <">", "<", ">=", "<=", "==", "!="> <"_", 'value'>]}
+#let grouping = ```
+{[`grouping` "(" 'constraint' ")"]}
 ```
 
-#let type-with-constraints = ```
-{[`type-with-constraints` 'identifier' <!, ["+" "(" 'constraint' ")"] * !>]}
+#let primary = ```
+{[`primary` <"_", 'value', 'identifier', 'grouping'>]}
+```
+
+#let reference = ```
+{[`reference` 'primary' <!, ["." 'identifier']*!>]}
+```
+
+#let unary = ```
+{[`unary` <[<!, "-"> 'unary'], 'reference'>]}
+```
+
+#let comparison = ```
+{[`comparison` 'unary'*<">", "<", ">=", "<=">]}
+```
+
+#let equality = ```
+{[`equality` 'comparison'*<"==", "!=">]}
+```
+
+#let constraint = ```
+{[`constraint` 'equality'*"&"]}
+```
+
+#let template-param = ```
+{[`template-param` 'identifier' <!, ["<:" 'type']>]}
+```
+
+#let template = ```
+{[`template` "[" <!, 'template-param'*","> "]"]}
 ```
 
 #let type-property = ```
-{[`type-property` 'identifier' ":" 'type-with-constraints']}
+{[`type-property` 'identifier' ":" 'type']}
 ```
 
-#let type-body = ```
-{[`type-body` "{" <!, 'type-property'*!> "}"]}
+#let complex-type = ```
+{[`complex-type` "{" <!, 'type-property'*!> "}"]}
 ```
 
-#let operation-type = ```
-{[`operation-type` "<" 'type-with-constraints' ">"]}
+#let named-type = ```
+{[`named-type` 'identifier']}
+```
+
+#let type-params = ```
+{[`type-params` "[" <!, 'type'*","> "]"]}
+```
+
+#let generic-type = ```
+{[`generic-type` 'named-type' <!, 'type-params'>]}
+```
+
+#let grouped-type = ```
+{[`grouped-type` "(" 'type' ")"]}
+```
+
+#let base-type = ```
+{[`base-type` <'grouped-type', 'complex-type', 'generic-type'>]}
+```
+
+#let constraint-type = ```
+{[`constraint-type` 'base-type' <!, ["where" 'constraint']>]}
+```
+
+#let type = ```
+{[`type` 'constraint-type']}
 ```
 
 #let type-statement = ```
-{[`type-statement` "type" 'identifier' "<" 'type-with-constraints'*"," ">" <!, 'type-body'>]}
+{[`type-statement` "type" 'identifier' <!, 'template'> "=" 'type']}
 ```
 
-#let operation-statement = ```
-{[`operation-statement` "op" 'operation-type' "operator" 'operation-type' "=" 'operation-type']}
+#let op-definition = ```
+{[`op-definition` "op" 'identifier' "(" 'type' ")" "->" 'type']}
 ```
 
-#let constraint-statement = ```
-{[`constraint-statement` "constraint" 'identifier' "=" 'constraint']}
+#let extend-statement = ```
+{[`extend-statement` "extend" 'type' "{" <!, 'op-definition'*!> "}"]}
+```
+
+#let predicate-statement = ```
+{[`predicate-statement` "predicate" 'identifier' "(" 'identifier' ":" 'type' ")" "=" 'constraint']}
 ```
 
 #let statement = ```
-{[`statement` <'type-statement', 'operation-statement', 'constraint-statement'>]}
+{[`statement` <'type-statement', 'extend-statement', 'predicate-statement'>]}
 ```
 
 #let rules = (
-  value,
-  constraint,
-  type-with-constraints,
-  type-property,
-  type-body,
-  operation-type,
-  type-statement,
-  operation-statement,
-  constraint-statement,
-  statement,
+  value: value,
+  grouping: grouping,
+  primary: primary,
+  reference: reference,
+  unary: unary,
+  comparison: comparison,
+  equality: equality,
+  constraint: constraint,
+  template-param: template-param,
+  template: template,
+  type-property: type-property,
+  complex-type: complex-type,
+  named-type: named-type,
+  type-params: type-params,
+  generic-type: generic-type,
+  grouped-type: grouped-type,
+  base-type: base-type,
+  constraint-type: constraint-type,
+  type: type,
+  type-statement: type-statement,
+  op-definition: op-definition,
+  extend-statement: extend-statement,
+  predicate-statement: predicate-statement,
+  statement: statement,
+)
+
+#let inline = (
+  "grouping",
+  "value",
+  "template-param",
+  "template",
+  "type-property",
+  "complex-type",
+  "type-params",
+  "named-type",
+  "grouped-type",
+  "generic-type",
+  "base-type",
+  "constraint-type",
+  "op-definition",
+  "type-statement",
+  "extend-statement",
+  "predicate-statement",
 )
 
 #set text(font: "Source Sans 3")
 
-= Midas type definition syntax
+#title[Midas type definition syntax]
 
-#for rule in rules {
-  render(rule)
-}
+= Outline
 
-/*
-#let by-name = (
-  value: value,
-  constraint: constraint,
-  type-with-constraints: type-with-constraints,
-  type-property: type-property,
-  type-body: type-body,
-  operation-type: operation-type,
-  type-statement: type-statement,
-  operation-statement: operation-statement,
-  constraint-statement: constraint-statement,
+#box(
+  columns(
+    2,
+    outline(title: none),
+  ),
+  height: 9cm,
+  stroke: 1pt,
+  inset: 1em,
 )
 
+= Statements and expressions
+
+#for (name, rule) in rules.pairs().rev() {
+  [== #name]
+  render(rule, css: css)
+}
+
 #let substitute(base-rule) = {
   let new-rule = base-rule
-  for (key, rule) in by-name.pairs() {
-    new-rule = new-rule.replace("'" + key + "'", rule.text.slice(1, -1))
+  for name in inline {
+    let rule = rules.at(name)
+    let replacement = rule.text.slice(1, -1).replace(regex("\[`.*?`"), "[")
+    replacement = "[" + replacement + "#`" + name + "`]"
+    new-rule = new-rule.replace(
+      "'" + name + "'",
+      replacement,
+    )
   }
   if new-rule != base-rule {
     new-rule = substitute(new-rule)
   }
-  return new-rule.replace(regex("`.*?`"), "")
+  return new-rule
 }
 
-#let combined = raw(substitute(statement.text))
-
-
 #set page(flipped: true)
-#render(combined)
-*/
+
+= Combined rules
+
+#for (name, rule) in rules.pairs() {
+  if not name in inline {
+    [== #name]
+    let combined = substitute(rule.text)
+    render(raw(combined), css: css)
+    //raw(block: true, combined)
+  }
+}

test.py

@@ -1,40 +1,21 @@
-import importlib
+import json
 from pathlib import Path
 
-from core.ast.printer import AnnotationAstPrinter, MidasAstPrinter
-from lexer.annotations import AnnotationLexer
-from lexer.midas import MidasLexer
-from lexer.token import Token
-from parser.annotations import AnnotationParser
-from parser.midas import MidasParser
-
-
-def test_annotation():
-    # Frame annotation
-    mod = importlib.import_module("examples.00_syntax_prototype.01_simple_types")
-
-    annotation: str = mod.__annotations__["df"]
-    lexer: AnnotationLexer = AnnotationLexer(annotation, "01_simple_types.py")
-    tokens: list[Token] = lexer.process()
-    # print([f"{t.type.name}('{t.lexeme}')" for t in tokens])
-
-    parser = AnnotationParser(tokens)
-    parsed = parser.parse()
-    print(parsed)
-    for err in parser.errors:
-        print(err.get_report())
-    printer = AnnotationAstPrinter()
-    if parsed is not None:
-        print(printer.print(parsed))
+from midas.ast.printer import MidasAstPrinter
+from midas.lexer.midas import MidasLexer
+from midas.lexer.token import Token
+from midas.parser.midas import MidasParser
 
 
 def test_midas():
     # Midas type definitions
-    path: Path = Path("examples") / "00_syntax_prototype" / "02_custom_types.midas"
+    path: Path = Path("examples") / "00_syntax_prototype" / "03_custom_types_v2.midas"
     definitions: str = path.read_text()
     midas_lexer: MidasLexer = MidasLexer(definitions, path.name)
     tokens: list[Token] = midas_lexer.process()
     # print([f"{t.type.name}('{t.lexeme}')" for t in tokens])
+    with open("tokens.json", "w") as f:
+        json.dump([f"{t.type.name}('{t.lexeme}')" for t in tokens], f, indent=4)
 
     parser = MidasParser(tokens)
     parsed = parser.parse()

tests/base.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import argparse
+import difflib
+import sys
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Iterator, Protocol
+
+
+class CaseResult(Protocol):
+    def dumps(self) -> str: ...
+
+
+class Tester(ABC):
+    """A test runner to check for regressions in the lexer and parser"""
+
+    CASES_DIR: Path = Path(__file__).parent / "cases"
+
+    @property
+    @abstractmethod
+    def namespace(self) -> str: ...
+
+    @property
+    def extension(self) -> str:
+        return "json"
+
+    @property
+    def base_dir(self) -> Path:
+        return self.CASES_DIR / self.namespace
+
+    @abstractmethod
+    def _list_tests(self) -> list[Path]: ...
+
+    def run_all_tests(self) -> bool:
+        paths: list[Path] = sorted(self._list_tests())
+        return self.run_tests(paths)
+
+    def run_tests(self, tests: list[Path]) -> bool:
+        rule: str = "-" * 80
+        n: int = len(tests)
+        successes: int = 0
+        failures: int = 0
+
+        print(rule)
+        for i, test in enumerate(tests):
+            print(f"Case {i+1}/{n}: {test.resolve().relative_to(self.CASES_DIR)}")
+            success: bool = self._run_test(test)
+            if success:
+                successes += 1
+            else:
+                failures += 1
+
+        print(rule)
+        print(f"Success: {successes}/{n}")
+        print(f"Failed: {failures}/{n}")
+        print(rule)
+        return failures == 0
+
+    def _run_test(self, path: Path) -> bool:
+        result_path: Path = self._result_path(path)
+        if not result_path.exists():
+            print("Missing snapshot. Please run the update command first")
+            return False
+        result: CaseResult = self._exec_case(path)
+        expected: str = result_path.read_text()
+        actual: str = result.dumps()
+
+        if expected == actual:
+            return True
+
+        diff = difflib.unified_diff(
+            expected.splitlines(keepends=True),
+            actual.splitlines(keepends=True),
+            fromfile="Snapshot",
+            tofile="Result",
+        )
+        self._print_diff(diff)
+        return False
+
+    @abstractmethod
+    def _exec_case(self, path: Path) -> CaseResult: ...
+
+    def update_all_tests(self):
+        paths: list[Path] = sorted(self._list_tests())
+        return self.update_tests(paths)
+
+    def update_tests(self, tests: list[Path]):
+        updated: int = 0
+        for test in tests:
+            if self._update_test(test):
+                updated += 1
+        print(f"Updated {updated}/{len(tests)} tests")
+
+    def _update_test(self, path: Path) -> bool:
+        result: CaseResult = self._exec_case(path)
+        result_path: Path = self._result_path(path)
+        current: str = result_path.read_text() if result_path.exists() else ""
+        new: str = result.dumps()
+        if current == new:
+            return False
+        result_path.write_text(new)
+        return True
+
+    def _result_path(self, test_path: Path) -> Path:
+        return test_path.parent / (test_path.name + f".ref.{self.extension}")
+
+    def _print_diff(self, diff: Iterator[str]):
+        for line in diff:
+            if line.startswith("+") and not line.startswith("+++"):
+                print(f"\033[92m{line}\033[0m", end="")
+            elif line.startswith("-") and not line.startswith("---"):
+                print(f"\033[91m{line}\033[0m", end="")
+            else:
+                print(line, end="")
+        print()
+
+    @classmethod
+    def main(cls):
+        parser = argparse.ArgumentParser()
+        subparsers = parser.add_subparsers(dest="subcommand")
+
+        update = subparsers.add_parser("update")
+        update.add_argument("-a", "--all", action="store_true")
+        update.add_argument("FILE", type=Path, nargs="*")
+
+        run = subparsers.add_parser("run")
+        run.add_argument("-a", "--all", action="store_true")
+        run.add_argument("FILE", type=Path, nargs="*")
+        args = parser.parse_args()
+
+        tester: Tester = cls()
+
+        match args.subcommand:
+            case "update":
+                if args.all:
+                    tester.update_all_tests()
+                else:
+                    tester.update_tests(args.FILE)
+            case "run":
+                success: bool
+                if args.all:
+                    success = tester.run_all_tests()
+                else:
+                    success = tester.run_tests(args.FILE)
+                if not success:
+                    sys.exit(1)
+            case None:
+                print("No subcommand provided. Available subcommands: run, update")
+                sys.exit(1)
+            case _:
+                print(f"Unknown subcommand '{args.subcommand}'")
+                sys.exit(1)

tests/cases/checker/01_simple_types.py

@@ -0,0 +1,14 @@
+# type: ignore
+# ruff: disable[F821]
+from __future__ import annotations
+
+df: Frame[
+    verified: bool,
+    birth_year: int,
+    height: float + ( _ > 0 ) + ( _ < 250 ),
+    name: str,
+    date: datetime,
+    float,
+    unknown: _,
+    _
+]

tests/cases/checker/01_simple_types.py.ref.json

@@ -0,0 +1,19 @@
+{
+  "diagnostics": [
+    {
+      "type": "Warning",
+      "location": {
+        "start": [
+          6,
+          4
+        ],
+        "end": [
+          13,
+          5
+        ]
+      },
+      "message": "FrameType not yet supported"
+    }
+  ],
+  "judgments": []
+}

tests/cases/checker/02_simple_operations.py

@@ -0,0 +1,11 @@
+a: int = 3
+b: int = 4
+
+c = a + b
+
+c = "invalid"
+
+d = True
+e = d + d
+
+f: float = a

tests/cases/checker/02_simple_operations.py.ref.json

@@ -0,0 +1,191 @@
+{
+  "diagnostics": [
+    {
+      "type": "Error",
+      "location": {
+        "start": [
+          6,
+          0
+        ],
+        "end": [
+          6,
+          13
+        ]
+      },
+      "message": "Cannot assign str to variable 'c' of type int"
+    },
+    {
+      "type": "Error",
+      "location": {
+        "start": [
+          9,
+          4
+        ],
+        "end": [
+          9,
+          9
+        ]
+      },
+      "message": "Undefined operation __add__ between bool and bool"
+    }
+  ],
+  "judgments": [
+    {
+      "location": {
+        "from": "L1:9",
+        "to": "L1:10"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 3
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L2:9",
+        "to": "L2:10"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 4
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L4:4",
+        "to": "L4:5"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L4:8",
+        "to": "L4:9"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "b"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L4:4",
+        "to": "L4:9"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "a"
+        },
+        "operator": "+",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "b"
+        }
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L6:4",
+        "to": "L6:13"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": "invalid"
+      },
+      "type": {
+        "name": "str"
+      }
+    },
+    {
+      "location": {
+        "from": "L8:4",
+        "to": "L8:8"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": true
+      },
+      "type": {
+        "name": "bool"
+      }
+    },
+    {
+      "location": {
+        "from": "L9:4",
+        "to": "L9:5"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "d"
+      },
+      "type": {
+        "name": "bool"
+      }
+    },
+    {
+      "location": {
+        "from": "L9:8",
+        "to": "L9:9"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "d"
+      },
+      "type": {
+        "name": "bool"
+      }
+    },
+    {
+      "location": {
+        "from": "L9:4",
+        "to": "L9:9"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "d"
+        },
+        "operator": "+",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "d"
+        }
+      },
+      "type": {}
+    },
+    {
+      "location": {
+        "from": "L11:11",
+        "to": "L11:12"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    }
+  ]
+}

tests/cases/checker/03_functions.py

@@ -0,0 +1,18 @@
+def foo(a: int, /, b: float, *, c: str):
+    return True
+
+
+r1 = foo()
+r2 = foo(1)
+r3 = foo(1, 2.0)
+r4 = foo(1, b=2.0)
+r5 = foo(1, 2.0, "test")
+r6 = foo(1, 2.0, b=3.0)
+r7 = foo(a=1)
+r8 = foo(g="test")
+
+r9a = foo(1, 2.0, c="test")
+r9b = foo(1, b=2.0, c="test")
+r9c = foo(1, c="test", b=2.0)
+
+r10 = foo("a", 3, c=False)

tests/cases/checker/04_custom_types.midas

@@ -0,0 +1,14 @@
+type Meter = float
+type Second = float
+type MeterPerSecond = float
+
+extend Meter {
+    def __add__: fn(Meter, /) -> Meter
+    def __sub__: fn(Meter, /) -> Meter
+    def __truediv__: fn(Second, /) -> MeterPerSecond
+}
+
+extend Second {
+    def __add__: fn(Second, /) -> Second
+    def __sub__: fn(Second, /) -> Second
+}

tests/cases/checker/04_custom_types.py

@@ -0,0 +1,6 @@
+# type: ignore
+# ruff: disable [F821]
+
+distance: Meter = cast(Meter, 123.45)
+time: Second = cast(Second, 6.7)
+speed = distance / time

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

@@ -0,0 +1,109 @@
+{
+  "diagnostics": [],
+  "judgments": [
+    {
+      "location": {
+        "from": "L4:18",
+        "to": "L4:37"
+      },
+      "expr": {
+        "_type": "CastExpr",
+        "type": {
+          "_type": "BaseType",
+          "base": "Meter",
+          "param": null
+        },
+        "expr": {
+          "_type": "LiteralExpr",
+          "value": 123.45
+        }
+      },
+      "type": {
+        "name": "Meter",
+        "type": {
+          "name": "float"
+        }
+      }
+    },
+    {
+      "location": {
+        "from": "L5:15",
+        "to": "L5:32"
+      },
+      "expr": {
+        "_type": "CastExpr",
+        "type": {
+          "_type": "BaseType",
+          "base": "Second",
+          "param": null
+        },
+        "expr": {
+          "_type": "LiteralExpr",
+          "value": 6.7
+        }
+      },
+      "type": {
+        "name": "Second",
+        "type": {
+          "name": "float"
+        }
+      }
+    },
+    {
+      "location": {
+        "from": "L6:8",
+        "to": "L6:16"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "distance"
+      },
+      "type": {
+        "name": "Meter",
+        "type": {
+          "name": "float"
+        }
+      }
+    },
+    {
+      "location": {
+        "from": "L6:19",
+        "to": "L6:23"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "time"
+      },
+      "type": {
+        "name": "Second",
+        "type": {
+          "name": "float"
+        }
+      }
+    },
+    {
+      "location": {
+        "from": "L6:8",
+        "to": "L6:23"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "distance"
+        },
+        "operator": "/",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "time"
+        }
+      },
+      "type": {
+        "name": "MeterPerSecond",
+        "type": {
+          "name": "float"
+        }
+      }
+    }
+  ]
+}

tests/cases/checker/05_control_flow.py

@@ -0,0 +1,25 @@
+def valid(a: int, b: int) -> int:
+    return a + b
+
+def with_if(a: int, b: int) -> int:
+    if a < b:
+        return b - a
+    else:
+        return a - b
+
+def unreachable1():
+    return
+    a = 0
+
+def unreachable2(a: int) -> int:
+    if a > 10:
+        return a - 10
+    else:
+        return a
+    b = 0
+
+def mixed(a: int, b: int):
+    if a < b:
+        return b - a
+    else:
+        return "oops"

tests/cases/checker/05_control_flow.py.ref.json

@@ -0,0 +1,450 @@
+{
+  "diagnostics": [
+    {
+      "type": "Warning",
+      "location": {
+        "start": [
+          12,
+          4
+        ],
+        "end": [
+          12,
+          9
+        ]
+      },
+      "message": "Unreachable statement"
+    },
+    {
+      "type": "Warning",
+      "location": {
+        "start": [
+          19,
+          4
+        ],
+        "end": [
+          19,
+          9
+        ]
+      },
+      "message": "Unreachable statement"
+    },
+    {
+      "type": "Error",
+      "location": {
+        "start": [
+          21,
+          0
+        ],
+        "end": [
+          25,
+          21
+        ]
+      },
+      "message": "Mixed return types: [BaseType(name='int'), BaseType(name='str')]"
+    }
+  ],
+  "judgments": [
+    {
+      "location": {
+        "from": "L2:11",
+        "to": "L2:12"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L2:15",
+        "to": "L2:16"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "b"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L2:11",
+        "to": "L2:16"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "a"
+        },
+        "operator": "+",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "b"
+        }
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L5:7",
+        "to": "L5:8"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L5:11",
+        "to": "L5:12"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "b"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L5:7",
+        "to": "L5:12"
+      },
+      "expr": {
+        "_type": "CompareExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "a"
+        },
+        "operator": "<",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "b"
+        }
+      },
+      "type": {
+        "name": "bool"
+      }
+    },
+    {
+      "location": {
+        "from": "L6:15",
+        "to": "L6:16"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "b"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L6:19",
+        "to": "L6:20"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L6:15",
+        "to": "L6:20"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "b"
+        },
+        "operator": "-",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "a"
+        }
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L8:15",
+        "to": "L8:16"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L8:19",
+        "to": "L8:20"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "b"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L8:15",
+        "to": "L8:20"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "a"
+        },
+        "operator": "-",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "b"
+        }
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L15:7",
+        "to": "L15:8"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L15:11",
+        "to": "L15:13"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 10
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L15:7",
+        "to": "L15:13"
+      },
+      "expr": {
+        "_type": "CompareExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "a"
+        },
+        "operator": ">",
+        "right": {
+          "_type": "LiteralExpr",
+          "value": 10
+        }
+      },
+      "type": {
+        "name": "bool"
+      }
+    },
+    {
+      "location": {
+        "from": "L16:15",
+        "to": "L16:16"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L16:19",
+        "to": "L16:21"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 10
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L16:15",
+        "to": "L16:21"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "a"
+        },
+        "operator": "-",
+        "right": {
+          "_type": "LiteralExpr",
+          "value": 10
+        }
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L18:15",
+        "to": "L18:16"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L22:7",
+        "to": "L22:8"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L22:11",
+        "to": "L22:12"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "b"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L22:7",
+        "to": "L22:12"
+      },
+      "expr": {
+        "_type": "CompareExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "a"
+        },
+        "operator": "<",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "b"
+        }
+      },
+      "type": {
+        "name": "bool"
+      }
+    },
+    {
+      "location": {
+        "from": "L23:15",
+        "to": "L23:16"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "b"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L23:19",
+        "to": "L23:20"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L23:15",
+        "to": "L23:20"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "b"
+        },
+        "operator": "-",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "a"
+        }
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L25:15",
+        "to": "L25:21"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": "oops"
+      },
+      "type": {
+        "name": "str"
+      }
+    }
+  ]
+}

tests/cases/checker/06_subtyping.py

@@ -0,0 +1,12 @@
+v1: int = 3
+v2: float = 4
+
+
+def maximum(a: float, b: float):
+    if b > a:
+        return b
+    return a
+
+
+v3 = maximum(v1, v2)
+v3 = v2 + v1

tests/cases/checker/06_subtyping.py.ref.json

@@ -0,0 +1,239 @@
+{
+  "diagnostics": [],
+  "judgments": [
+    {
+      "location": {
+        "from": "L1:10",
+        "to": "L1:11"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 3
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L2:12",
+        "to": "L2:13"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 4
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L6:7",
+        "to": "L6:8"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "b"
+      },
+      "type": {
+        "name": "float"
+      }
+    },
+    {
+      "location": {
+        "from": "L6:11",
+        "to": "L6:12"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "float"
+      }
+    },
+    {
+      "location": {
+        "from": "L6:7",
+        "to": "L6:12"
+      },
+      "expr": {
+        "_type": "CompareExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "b"
+        },
+        "operator": ">",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "a"
+        }
+      },
+      "type": {
+        "name": "bool"
+      }
+    },
+    {
+      "location": {
+        "from": "L7:15",
+        "to": "L7:16"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "b"
+      },
+      "type": {
+        "name": "float"
+      }
+    },
+    {
+      "location": {
+        "from": "L8:11",
+        "to": "L8:12"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "a"
+      },
+      "type": {
+        "name": "float"
+      }
+    },
+    {
+      "location": {
+        "from": "L11:5",
+        "to": "L11:12"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "maximum"
+      },
+      "type": {
+        "pos_args": [],
+        "args": [
+          {
+            "pos": 0,
+            "name": "a",
+            "type": {
+              "name": "float"
+            },
+            "required": true
+          },
+          {
+            "pos": 1,
+            "name": "b",
+            "type": {
+              "name": "float"
+            },
+            "required": true
+          }
+        ],
+        "kw_args": [],
+        "returns": {
+          "name": "float"
+        }
+      }
+    },
+    {
+      "location": {
+        "from": "L11:13",
+        "to": "L11:15"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "v1"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L11:17",
+        "to": "L11:19"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "v2"
+      },
+      "type": {
+        "name": "float"
+      }
+    },
+    {
+      "location": {
+        "from": "L11:5",
+        "to": "L11:20"
+      },
+      "expr": {
+        "_type": "CallExpr",
+        "callee": {
+          "_type": "VariableExpr",
+          "name": "maximum"
+        },
+        "arguments": [
+          {
+            "_type": "VariableExpr",
+            "name": "v1"
+          },
+          {
+            "_type": "VariableExpr",
+            "name": "v2"
+          }
+        ],
+        "keywords": {}
+      },
+      "type": {
+        "name": "float"
+      }
+    },
+    {
+      "location": {
+        "from": "L12:5",
+        "to": "L12:7"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "v2"
+      },
+      "type": {
+        "name": "float"
+      }
+    },
+    {
+      "location": {
+        "from": "L12:10",
+        "to": "L12:12"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "v1"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L12:5",
+        "to": "L12:12"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "v2"
+        },
+        "operator": "+",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "v1"
+        }
+      },
+      "type": {
+        "name": "float"
+      }
+    }
+  ]
+}

tests/cases/generator/01_simple_types.midas

@@ -0,0 +1,14 @@
+type Meter = float
+type Second = float
+type MeterPerSecond = float
+
+extend Meter {
+    def __add__: fn(Meter, /) -> Meter
+    def __sub__: fn(Meter, /) -> Meter
+    def __truediv__: fn(Second, /) -> MeterPerSecond
+}
+
+extend Second {
+    def __add__: fn(Second, /) -> Second
+    def __sub__: fn(Second, /) -> Second
+}

tests/cases/generator/01_simple_types.py

@@ -0,0 +1,5 @@
+from midas import cast, Meter, Second
+
+distance: Meter = cast(Meter, 123.45)
+time: Second = cast(Second, 6.7)
+speed = distance / time

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

@@ -0,0 +1,79 @@
+Module(
+  body=[
+    ImportFrom(
+      module='midas',
+      names=[
+        alias(name='cast'),
+        alias(name='Meter'),
+        alias(name='Second')],
+      level=0),
+    Assign(
+      targets=[
+        Name(id='__midas_alias_0__')],
+      value=Constant(value=123.45)),
+    Assert(
+      test=Call(
+        func=Name(id='isinstance'),
+        args=[
+          Name(id='__midas_alias_0__'),
+          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_alias_0__')],
+                keywords=[]),
+              attr='__name__'),
+            conversion=-1),
+          Constant(value=' to float')])),
+    Assign(
+      targets=[
+        Name(id='distance')],
+      value=Name(id='__midas_alias_0__')),
+    Delete(
+      targets=[
+        Name(id='__midas_alias_0__')]),
+    Assign(
+      targets=[
+        Name(id='__midas_alias_1__')],
+      value=Constant(value=6.7)),
+    Assert(
+      test=Call(
+        func=Name(id='isinstance'),
+        args=[
+          Name(id='__midas_alias_1__'),
+          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_alias_1__')],
+                keywords=[]),
+              attr='__name__'),
+            conversion=-1),
+          Constant(value=' to float')])),
+    Assign(
+      targets=[
+        Name(id='time')],
+      value=Name(id='__midas_alias_1__')),
+    Delete(
+      targets=[
+        Name(id='__midas_alias_1__')]),
+    Assign(
+      targets=[
+        Name(id='speed')],
+      value=BinOp(
+        left=Name(id='distance'),
+        op=Div(),
+        right=Name(id='time')))],
+  type_ignores=[])

tests/cases/midas-parser/01_simple_types.midas

@@ -0,0 +1,57 @@
+// Simple custom type derived from float
+type Custom = float
+
+// Simple custom types with constraints
+type Latitude = float where (-90 <= _ <= 90)
+type Longitude = float where (-180 <= _ <= 180)
+
+// Generic custom type (a Difference of T is derived from T, e.g. a difference of floats is a float
+type Difference[T] = T
+
+// Complex custom type, containing two values accessible through properties
+type GeoLocation = {
+    prop lat: Latitude
+    prop lon: Longitude
+}
+
+// Define operations on our custom type
+extend GeoLocation {
+    // This type is compatible with the `-` operation with another GeoLocation
+    // i.e. you can subtract a GeoLocation from another GeoLocation, resulting
+    // in a Difference of GeoLocations
+    def __sub__: fn(GeoLocation, /) -> Difference[GeoLocation]
+}
+
+// For complex generics, you need to specify how the genericity the properties
+// are handled
+type Difference[GeoLocation] = {
+    prop lat: Difference[Latitude]
+    prop lon: Difference[Longitude]
+}
+
+// Simple operation defined on our custom types
+extend Latitude {
+    def __sub__: fn(Latitude, /) -> Difference[Latitude]
+}
+
+extend Longitude {
+    def __sub__: fn(Longitude, /) -> Difference[Longitude]
+}
+
+// Predefined custom predicates that can be referenced in other definitions
+predicate Positive(v: float) = v >= 0
+predicate StrictlyPositive(v: float) = v > 0
+predicate Equatorial(loc: GeoLocation) = (-10 <= loc.lat <= 10)
+predicate Arctic(loc: GeoLocation) = (loc.lat >= 66)
+
+type Person = {
+    prop name: str
+
+    // Property with an inline constraint
+    prop age: Optional[int where (0 <= _ < 150)]
+
+    // Property referencing a predicate
+    prop height: float where StrictlyPositive
+
+    prop home: GeoLocation
+}

tests/cases/python-parser/01_simple_types.py

@@ -0,0 +1,14 @@
+# type: ignore
+# ruff: disable[F821]
+from __future__ import annotations
+
+df: Frame[
+    verified: bool,
+    birth_year: int,
+    height: float + ( _ > 0 ) + ( _ < 250 ),
+    name: str,
+    date: datetime,
+    float,
+    unknown: _,
+    _
+]

tests/cases/python-parser/01_simple_types.py.ref.json

@@ -0,0 +1,89 @@
+{
+  "stmts": [
+    {
+      "_type": "RawStmt",
+      "stmt": "from __future__ import annotations"
+    },
+    {
+      "_type": "TypeAssign",
+      "name": "df",
+      "type": {
+        "_type": "FrameType",
+        "columns": [
+          {
+            "_type": "FrameColumn",
+            "name": "verified",
+            "type": {
+              "_type": "BaseType",
+              "base": "bool",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "birth_year",
+            "type": {
+              "_type": "BaseType",
+              "base": "int",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "height",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "(_ > 0) + (_ < 250)"
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "name",
+            "type": {
+              "_type": "BaseType",
+              "base": "str",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "date",
+            "type": {
+              "_type": "BaseType",
+              "base": "datetime",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": null,
+            "type": {
+              "_type": "BaseType",
+              "base": "float",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "unknown",
+            "type": null
+          },
+          {
+            "_type": "FrameColumn",
+            "name": null,
+            "type": {
+              "_type": "BaseType",
+              "base": "_",
+              "param": null
+            }
+          }
+        ]
+      }
+    }
+  ]
+}

tests/cases/python-parser/02_custom_types.py

@@ -0,0 +1,25 @@
+# type: ignore
+# ruff: disable[F821]
+from __future__ import annotations
+
+df: Frame[
+    location: GeoLocation
+]
+
+lat: Column[GeoLocation] = df["location"].lat
+lon: Column[GeoLocation] = df["location"].lon
+
+lat + lon
+
+lat1: Latitude = lat[0]
+lat2: Latitude = lat[1]
+lat_diff: Difference[Latitude] = lat2 - lat1
+
+df2: Frame[
+    age: int + (_ >= 0),
+    height: float + (_ >= 0),
+]
+df2_bis: Frame[
+    age: int + Positive,
+    height: float + Positive,
+]