Merge pull request 'Properly check variable assignment' (#36) from feat/variable-is-defined into main

Reviewed-on: #36

.gitea/workflows/tests.yaml

@@ -0,0 +1,22 @@
+name: Tests
+on:
+  push:
+    branches:
+      - main
+      - master
+  pull_request:
+    branches:
+      - "**"
+
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+    container: catthehacker/ubuntu:act-latest
+    if: github.event.pull_request.draft == false
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: "3.11"
+      - run: python3 -m tests

.gitignore

@@ -7,3 +7,5 @@ venv
 uv.lock
 .python-version
 /out
+/examples/**/build/
+/examples/**/*.pyi

README.md

@@ -1,7 +1,163 @@
-# Midas
+<h1>Midas</h1>
 
 *Midas* is a type system to _Maintain Integrity of Data with Annotated Structures_. In Greek mythology, [Midas](https://en.wikipedia.org/wiki/Midas) was a Phrygian king who was blessed with the gift of turning everything he touched into gold.
 
 *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.
+
+<details>
+<summary><strong>Table of Contents</strong></summary>
+
+- [Requirements](#requirements)
+- [Installation](#installation)
+- [Commands](#commands)
+  - [Type Checking](#type-checking)
+  - [Compiling](#compiling)
+  - [Formatting](#formatting)
+  - [Highlighting](#highlighting)
+  - [Dumping the AST](#dumping-the-ast)
+  - [Dumping the Registry](#dumping-the-registry)
+  - [Generating Stubs](#generating-stubs)
+  - [Showing Type Judgements](#showing-type-judgements)
+  - [Validating Definitions](#validating-definitions)
+- [Tests](#tests)
+</details>
+
+
+## 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
+
+Hereafter is a description of the commands you can use with Midas. For a full description, refer to [the manual](./docs/manual.pdf) or `midas <subcommand> --help`.
+
+### Type Checking
+
+```shell
+midas check -t types.midas source.py
+```
+
+This command parses the given files and run the type checkers against the Midas definitions and Python program. Diagnostics are then printed showing warnings and errors.
+
+### Compiling
+
+```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.
+
+> [!WARNING]
+> By default, any type checking error aborts the compilation and the generator is not run. You can bypass this behaviour with the `--ignore-errors` flag.
+> Only use this flag if you know what you are doing as it will produce a possibly unsafe program and goes against the whole purpose of Midas
+
+### Formatting
+
+```shell
+midas format types.midas
+midas format types.midas -o formatted.midas
+```
+
+This command parses the given Midas file and outputs a pretty printed file from the AST.
+
+### Highlighting
+
+```shell
+midas highlight source.py
+midas highlight source.py -o highlighted.html
+midas highlight types.midas
+midas highlight types.midas -o highlighted.html
+```
+
+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 parse source.py
+midas parse types.midas
+```
+
+For debugging purposes, you can output the AST parsed from a Python or Midas file. For Python files, the `--raw` flags lets you toggle the custom AST parsing. With `--raw`, the raw AST is returned, as produced by the builtin `ast` module. This flag has no effect on Midas files.
+
+### Dumping the Registry
+
+```shell
+midas dump-registry -t types.midas
+```
+
+This command processes the given Midas definitions and dumps the contents of the types registry.
+
+### Generating Stubs
+
+```shell
+midas stubs types.midas -o stubs.pyi
+```
+
+This command generates Python stubs from a Midas definition file
+
+### Showing Type Judgements
+
+```shell
+midas types -t types.midas source.py
+```
+
+This command type checks the given Python source file and logs all typing judgements made by the type checker.
+
+### Validating Definitions
+
+```shell
+midas validate types.midas
+```
+
+This command lets you validate a Midas definition file by running the parser and type checker, verifying syntax and references.
+
+## Tests
+
+Several snapshot tests are available to assert the good behaviour of the parser, type checker and generator. 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
+uv run -m tests.generator run -a
+```
+
+Alternatively, you can run all tests by executing the `tests` module directly:
+
+```shell
+uv run -m tests
+```
+
+When running only one test group, you may use one of the following subcommands.\
+Not specifying any subcommand is equivalent to running `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 ...`

assets/icon.svg

@@ -0,0 +1,117 @@
+<?xml version="1.0" encoding="UTF-8" standalone="no"?>
+<!-- Created with Inkscape (http://www.inkscape.org/) -->
+
+<svg
+   width="128"
+   height="128"
+   viewBox="0 0 128 128"
+   version="1.1"
+   id="svg1"
+   inkscape:export-filename="logo.png"
+   inkscape:export-xdpi="96"
+   inkscape:export-ydpi="96"
+   inkscape:version="1.4.4 (1:1.4.4+202605061436+dcaf3e7d9e)"
+   sodipodi:docname="logo.svg"
+   xmlns:inkscape="http://www.inkscape.org/namespaces/inkscape"
+   xmlns:sodipodi="http://sodipodi.sourceforge.net/DTD/sodipodi-0.dtd"
+   xmlns:xlink="http://www.w3.org/1999/xlink"
+   xmlns="http://www.w3.org/2000/svg"
+   xmlns:svg="http://www.w3.org/2000/svg">
+  <sodipodi:namedview
+     id="namedview1"
+     pagecolor="#ffffff"
+     bordercolor="#000000"
+     borderopacity="0.25"
+     inkscape:showpageshadow="2"
+     inkscape:pageopacity="0.0"
+     inkscape:pagecheckerboard="0"
+     inkscape:deskcolor="#d1d1d1"
+     inkscape:document-units="mm"
+     showgrid="true"
+     inkscape:zoom="1.9332778"
+     inkscape:cx="-8.2760999"
+     inkscape:cy="112.2446"
+     inkscape:window-width="2584"
+     inkscape:window-height="1028"
+     inkscape:window-x="0"
+     inkscape:window-y="24"
+     inkscape:window-maximized="1"
+     inkscape:current-layer="layer1">
+    <inkscape:grid
+       id="grid1"
+       units="px"
+       originx="0"
+       originy="0"
+       spacingx="4"
+       spacingy="4"
+       empcolor="#0099e5"
+       empopacity="0.30196078"
+       color="#0099e5"
+       opacity="0.14901961"
+       empspacing="4"
+       enabled="true"
+       visible="true" />
+  </sodipodi:namedview>
+  <defs
+     id="defs1">
+    <linearGradient
+       inkscape:collect="always"
+       xlink:href="#linearGradient4689"
+       id="linearGradient1478"
+       gradientUnits="userSpaceOnUse"
+       gradientTransform="matrix(0.562541,0,0,0.567972,-9.399749,-5.305317)"
+       x1="26.648937"
+       y1="20.603781"
+       x2="135.66525"
+       y2="114.39767" />
+    <linearGradient
+       id="linearGradient4689">
+      <stop
+         style="stop-color:#e1be1e;stop-opacity:1;"
+         offset="0"
+         id="stop4691" />
+      <stop
+         style="stop-color:#ffeb82;stop-opacity:1;"
+         offset="1"
+         id="stop4693" />
+    </linearGradient>
+    <linearGradient
+       inkscape:collect="always"
+       xlink:href="#linearGradient4671"
+       id="linearGradient1475"
+       gradientUnits="userSpaceOnUse"
+       gradientTransform="matrix(0.562541,0,0,0.567972,-9.399749,-5.305317)"
+       x1="150.96111"
+       y1="192.35176"
+       x2="112.03144"
+       y2="137.27299" />
+    <linearGradient
+       id="linearGradient4671">
+      <stop
+         style="stop-color:#ffdc21;stop-opacity:1;"
+         offset="0"
+         id="stop4673" />
+      <stop
+         style="stop-color:#ffeb82;stop-opacity:1;"
+         offset="1"
+         id="stop4675" />
+    </linearGradient>
+  </defs>
+  <g
+     inkscape:label="Calque 1"
+     inkscape:groupmode="layer"
+     id="layer1">
+    <g
+       id="g1"
+       transform="translate(2.911719,3.414527)">
+      <path
+         style="fill:url(#linearGradient1478);fill-opacity:1"
+         d="m 60.510156,6.3979729 c -4.583653,0.021298 -8.960939,0.4122177 -12.8125,1.09375 C 36.35144,9.4962267 34.291407,13.691825 34.291406,21.429223 v 10.21875 h 26.8125 v 3.40625 h -26.8125 -10.0625 c -7.792459,0 -14.6157592,4.683717 -16.7500002,13.59375 -2.46182,10.212966 -2.5710151,16.586023 0,27.25 1.9059283,7.937852 6.4575432,13.593748 14.2500002,13.59375 h 9.21875 v -12.25 c 0,-8.849902 7.657144,-16.656248 16.75,-16.65625 h 26.78125 c 7.454951,0 13.406253,-6.138164 13.40625,-13.625 v -25.53125 c 0,-7.266339 -6.12998,-12.7247775 -13.40625,-13.9375001 -4.605987,-0.7667253 -9.385097,-1.1150483 -13.96875,-1.09375 z m -14.5,8.2187501 c 2.769547,0 5.03125,2.298646 5.03125,5.125 -2e-6,2.816336 -2.261703,5.09375 -5.03125,5.09375 -2.779476,-1e-6 -5.03125,-2.277415 -5.03125,-5.09375 -1e-6,-2.826353 2.251774,-5.125 5.03125,-5.125 z"
+         id="path1948" />
+      <path
+         style="fill:url(#linearGradient1475);fill-opacity:1"
+         d="m 91.228906,35.054223 v 11.90625 c 0,9.230755 -7.825895,16.999999 -16.75,17 h -26.78125 c -7.335833,0 -13.406249,6.278483 -13.40625,13.625 v 25.531247 c 0,7.26634 6.318588,11.54032 13.40625,13.625 8.487331,2.49561 16.626237,2.94663 26.78125,0 6.750155,-1.95439 13.406253,-5.88761 13.40625,-13.625 V 92.897973 h -26.78125 v -3.40625 h 26.78125 13.406254 c 7.79246,0 10.69625,-5.435408 13.40624,-13.59375 2.79933,-8.398886 2.68022,-16.475776 0,-27.25 -1.92578,-7.757441 -5.60387,-13.59375 -13.40624,-13.59375 z m -15.0625,64.65625 c 2.779478,3e-6 5.03125,2.277417 5.03125,5.093747 -2e-6,2.82635 -2.251775,5.125 -5.03125,5.125 -2.76955,0 -5.03125,-2.29865 -5.03125,-5.125 2e-6,-2.81633 2.261697,-5.093747 5.03125,-5.093747 z"
+         id="path1950" />
+    </g>
+  </g>
+</svg>

docs/fonts.yaml

@@ -0,0 +1,2 @@
+Source Sans Pro:
+Source Sans 3:

docs/midas.sublime-syntax

@@ -0,0 +1,211 @@
+%YAML 1.2
+---
+name: Midas
+file_extensions:
+  - midas
+scope: source.midas
+
+variables:
+  identifier: "[a-zA-Z_][a-zA-Z0-9_]*"
+
+contexts:
+  prototype:
+    - include: comments
+
+  main:
+    - include: keywords
+    - include: types
+
+  comments:
+    - match: "//"
+      scope: punctuation.definition.comment.midas
+      push:
+        - meta_scope: comment.line.midas
+        - match: $
+          pop: true
+    - match: /\*
+      scope: punctuation.definition.comment.midas
+      push:
+        - meta_scope: comment.block.midas
+        - match: \*/
+          pop: true
+
+  string:
+    - meta_include_prototype: false
+    - meta_scope: string.quoted.double.c
+    - match: '"'
+      pop: true
+
+  keywords:
+    - match: \balias\b
+      scope: keyword.declaration.midas
+      push: alias-stmt
+    - match: \btype\b
+      scope: keyword.declaration.midas
+      push: type-stmt
+    - match: \bextend\b
+      scope: keyword.declaration.midas
+      push: extend-stmt
+    - match: \bpredicate\b
+      scope: keyword.declaration.midas
+      push: predicate-stmt
+
+  alias-stmt:
+    - match: "{{identifier}}"
+      scope: entity.name.type
+    - match: "="
+      scope: keyword.operator.equal.midas
+      push: type-expr
+    - match: $
+      pop: true
+
+  type-stmt:
+    - match: "{{identifier}}"
+      scope: entity.name.type
+    - match: \[
+      push: type-params
+    - match: "="
+      scope: keyword.operator.equal.midas
+      push: type-expr
+    - match: $
+      pop: true
+
+  type-expr:
+    - match: \b(fn)\s*(\()
+      captures:
+        1: keyword.other.midas
+        2: punctuation.section.group.begin
+      push: fn-params
+    - match: \b(where)\b
+      scope: keyword.other.midas
+      set: constraint
+    - match: "Frame"
+      scope: entity.name.type
+      push:
+        - match: \[
+          push: frame-schema
+        - match: $
+          pop: true
+    - match: "{{identifier}}"
+      scope: entity.name.type
+    - match: $
+      pop: 2
+
+  fn-params:
+    - match: "({{identifier}})(:)"
+      captures:
+        1: variable.parameter.midas
+        2: punctuation.separator.annotation.midas
+      push:
+        - include: type-expr
+        - match: \?
+          scope: keyword.operator.qmark.midas
+        - match: "(?=,)"
+          scope: punctuation.separator.midas
+          pop: true
+        - match: '(?=\))'
+          pop: true
+    - include: type-expr
+    - match: '\)'
+      set:
+        - match: "->"
+          scope: keyword.operator.arrow.midas
+          set: type-expr
+
+  constraint:
+    - match: $
+      pop: 2
+
+    - match: \d+(\.\d+)?
+      scope: constant.numeric.midas
+
+    - match: \b(true|false|none)\b
+      scope: constant.language.midas
+
+    - match: '"'
+      push: string
+
+    - match: (<=|>=|<|>|==|!=|&)
+      scope: keyword.operator
+
+    - match: _
+      scope: variable.language.midas
+
+    - match: '{{identifier}}(?=\s*\()'
+      scope: variable.function.midas
+
+    - match: "{{identifier}}"
+      scope: variable.other.readwrite.midas
+
+  type-params:
+    - match: "<:"
+      scope: keyword.operator.subtype.midas
+    - match: "[a-zA-Z][a-zA-Z_0-9]*"
+      scope: entity.name.type
+    - match: "]"
+      pop: true
+
+  extend-stmt:
+    - match: "{{identifier}}"
+      scope: entity.name.type
+    - match: \[
+      push: type-params
+    - match: \{
+      scope: punctuation.section.block.begin
+      set: extend-body
+
+  extend-body:
+    - include: member-stmt
+    - match: \}
+      scope: punctuation.section.block.end
+      pop: true
+
+  member-stmt:
+    - match: \b(prop|def)\b
+      scope: keyword.other.midas
+      push:
+        - match: "{{identifier}}"
+          scope: variable.other.member
+        - match: ":"
+          push: type-expr
+        - match: $
+          pop: true
+
+  predicate-stmt:
+    - match: "{{identifier}}"
+      scope: entity.name.function.midas
+    - match: '\('
+      push: predicate-params
+    - match: "="
+      scope: keyword.operator.equal.midas
+      set: constraint
+    - match: $
+      pop: true
+
+  predicate-params:
+    - match: "({{identifier}})(:)"
+      captures:
+        1: variable.parameter.midas
+        2: punctuation.separator.annotation.midas
+      push:
+        - include: type-expr
+        - match: "(?=,)"
+          scope: punctuation.separator.midas
+          pop: true
+        - match: '(?=\))'
+          pop: true
+
+    - match: '\)'
+      pop: true
+
+  frame-schema:
+    - include: frame-column
+    - match: \]
+      # scope: punctuation.section.block.end
+      pop: true
+
+  frame-column:
+    - match: "{{identifier}}"
+      scope: variable.other.member
+    - match: ":"
+      push: type-expr

docs/requirements.json

@@ -0,0 +1,6 @@
+{
+    "preview": {},
+    "local": {
+        "codly": "https://github.com/LordBaryhobal/codly-fix-v15/archive/refs/tags/v1.3.1-fix.zip"
+    }
+}

docs/template.typ

@@ -0,0 +1,147 @@
+#import "@preview/modpattern:0.2.0": modpattern
+
+#let TODO = block(
+  width: 6em,
+  height: 3em,
+  stroke: red,
+  fill: modpattern(
+    size: (10pt, 10pt),
+    line(
+      start: (0%, 0%),
+      end: (100%, 100%),
+      stroke: gray.transparentize(60%) + 2pt,
+    ),
+  ),
+  align(
+    center + horizon,
+    text(fill: red, size: 1.5em)[*TODO*],
+  ),
+)
+
+#let _render-header(version, hash) = {
+  let last-heading = query(heading.where(level: 1).before(here())).last(default: none)
+  let next-heading = query(heading.where(level: 1).after(here())).first(default: none)
+
+  let current-heading = if next-heading != none and next-heading.location().page() == here().page() {
+    next-heading
+  } else if last-heading != none {
+    last-heading
+  } else { none }
+
+  let chapter = if current-heading != none {
+    let body = current-heading.body
+    if current-heading.numbering != none {
+      let num = counter(heading).display(current-heading.numbering, at: current-heading.location())
+      body = [#num #body]
+    }
+    body
+  } else []
+
+  grid(
+    columns: (1fr, auto, 1fr),
+    align: (left, center, right),
+    document.title, [v#version - #hash], chapter,
+  )
+}
+
+#let _unshift-prefix(prefix, content) = context {
+  pad(left: -measure(prefix).width, prefix + content)
+}
+
+#let project(
+  title: none,
+  author: none,
+  version: "0.0.1",
+  hash: none,
+  icon-path: none,
+  doc,
+) = {
+  assert(title != none, message: "Please provide a title")
+
+  set document(
+    title: title,
+    author: author,
+  )
+  set text(
+    font: "Source Sans 3",
+  )
+
+  set raw(syntaxes: path("midas.sublime-syntax"))
+
+  let front-page() = {
+    let version-name = [v#version]
+    if hash != none {
+      version-name = [#version-name - #hash]
+    }
+    align(center)[
+      #{
+        set text(size: 1.5em)
+        std.title()
+      }
+
+      #version-name
+
+      #if icon-path != none {
+        v(1cm)
+        image(icon-path)
+      }
+    ]
+    pagebreak()
+  }
+
+  let outlines() = {
+    outline()
+    pagebreak()
+    outline(
+      title: [List of Listings],
+      target: figure.where(kind: raw),
+    )
+
+    outline(
+      title: [List of Tables],
+      target: figure.where(kind: table),
+    )
+  }
+
+  let main() = {
+    // Adapted from https://github.com/hei-templates/hei-synd-thesis/blob/7d2b941197babae0bf3afd4e5914754e09a64001/lib/template-thesis.typ#L242-L261
+    show heading.where(level: 1): it => {
+      pagebreak()
+
+      set text(size: 1.5em)
+      set block(above: 1.2em, below: 1.2em)
+      if it.numbering != none {
+        let num = numbering(it.numbering, ..counter(heading).at(it.location()))
+        let prefix = num + h(1em)
+        _unshift-prefix(prefix, it.body)
+      } else {
+        it
+      }
+    }
+
+    show heading.where(level: 2): it => {
+      if it.numbering != none {
+        let num = numbering(it.numbering, ..counter(heading).at(it.location()))
+        _unshift-prefix(num + h(0.8em), it.body)
+      } else {
+        it
+      }
+    }
+
+    set page(
+      header: context _render-header(version, hash),
+      footer: context if page.numbering != none {
+        align(center, counter(page).display(page.numbering, both: true))
+      },
+      numbering: "1 / 1",
+    )
+
+    show heading: set heading(numbering: "I.1.")
+    counter(page).update(1)
+    doc
+  }
+
+  front-page()
+  outlines()
+  main()
+}

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

examples/01_simple_type_checking/01_simple_operations.py

@@ -9,3 +9,5 @@ d = True
 e = d + d
 
 f: float = a
+
+f = -f

examples/01_simple_type_checking/02_simple_types.midas

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

examples/01_simple_type_checking/02_simple_types.py

@@ -1,8 +1,6 @@
 # type: ignore
 # ruff: disable [F821]
 
-midas.using("02_simple_types.midas")
-
 distance: Meter = cast(Meter, 123.45)
 time: Second = cast(Second, 6.7)
 speed = distance / time

examples/01_simple_type_checking/03_control_flow.py

@@ -4,13 +4,20 @@ def minimum(x: int, y: int):
     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)

examples/02_demonstration/demo.midas

@@ -0,0 +1,15 @@
+predicate in_range(min: float, max: float)(v: float) = min <= v & v <= max
+predicate is_ratio = in_range(0, 1)
+
+type Currency = float
+type Price[T <: Currency] = T where _ >= 0
+
+extend Price[T <: Currency] {
+    def __add__: fn(Price[T], /) -> Price[T]
+}
+
+type EUR = Currency
+type USD = Currency
+type CHF = Currency
+
+type Discount = float where is_ratio(_)

examples/02_demonstration/demo.py

@@ -0,0 +1,35 @@
+from typing import TypeVar
+
+from demo_stubs import CHF, EUR, USD, Currency, Discount, Price
+
+from midas.typing import cast, unsafe_cast
+
+T = TypeVar("T", bound=Currency)
+
+
+def apply_discount(amount: Price[T], discount: Discount) -> Price[T]:
+    return cast(Price[T], (1.0 - discount) * amount)
+
+
+a1 = cast(Price[EUR], 3.2)
+a2 = cast(Price[USD], 10.4)
+r1 = cast(Discount, 0.2)
+
+print(apply_discount(a1, r1))
+print(apply_discount(a2, r1))
+
+a3 = a1 + a1
+a4 = a1 + a2  # cannot add euros and dollars
+a3 = a2  # cannot change variable type
+
+dyn_price = float(input("Price (CHF): "))
+dyn_discount = float(input("Discount (0.0-1.0): "))
+discounted = apply_discount(
+    cast(Price[CHF], dyn_price),
+    cast(Discount, dyn_discount),
+)
+
+print(f"Discounted: CHF {discounted}")
+
+large_data = [i * 10 for i in range(100)]
+prices = unsafe_cast(list[Price[EUR]], large_data)

examples/02_demonstration/demo_stubs.pyi

@@ -0,0 +1,14 @@
+from __future__ import annotations
+from typing import Generic, TypeVar
+
+class Currency(float): ...
+
+_T0 = TypeVar("_T0", bound=Currency, covariant=True)
+
+class Price(Currency, Generic[_T0]):
+    def __add__(self, _0: Price[_T0], /) -> Price[_T0]: ...
+
+class EUR(Currency): ...
+class USD(Currency): ...
+class CHF(Currency): ...
+class Discount(float): ...

examples/02_demonstration/weather/custom_types.midas

@@ -0,0 +1,65 @@
+predicate in_range(min: float, max: float)(v: float) = min <= v & v <= max
+predicate is_percentage = in_range(0.0, 100.0)
+
+type Celsius = float
+type Kelvin = float where _ >= 0
+type Hectopascal = float
+
+type Temperature = Celsius where in_range(-30.0, 100.0)(_)
+type Pressure = Hectopascal where in_range(800.0, 1100.0)(_)
+type Humidity = float where is_percentage(_)
+type HeatIndex = float
+
+type StationID = str where len(_) == 3 & _.isupper()
+
+type Mean[T <: float] = float
+
+extend Celsius {
+    def __add__: fn(Celsius, /) -> Celsius
+    def __sub__: fn(Celsius, /) -> Celsius
+}
+
+extend Kelvin {
+    def __add__: fn(Kelvin, /) -> Kelvin
+}
+
+extend Hectopascal {
+    def __add__: fn(Hectopascal, /) -> Hectopascal
+    def __sub__: fn(Hectopascal, /) -> Hectopascal
+}
+
+alias RawData = Frame[
+    station_id: str,
+    timestamp: str,
+    temperature: float,
+    pressure: float,
+    humidity: float,
+]
+
+alias Data = Frame[
+    station_id: StationID,
+    timestamp: Any,
+    temperature: Temperature,
+    pressure: Pressure,
+    humidity: Humidity,
+]
+
+alias DataWithHI = Frame[
+    station_id: StationID,
+    timestamp: Any,
+    temperature: Temperature,
+    pressure: Pressure,
+    humidity: Humidity,
+    heat_index: HeatIndex,
+]
+
+alias DailyAverages = Frame[
+    timestamp: Any,
+    temperature: Mean[Temperature],
+    pressure: Mean[Pressure],
+    humidity: Mean[Humidity],
+    heat_index: Mean[HeatIndex],
+]
+
+// predicate limit_amplitude(max_amp: float)(ls: list[float]) = max(ls) - min(ls) <= max_amp
+// type LowAmplitudeWave = list[float where _ >= 1] where limit_amplitude(10)(_)

examples/02_demonstration/weather/gen_data.py

@@ -0,0 +1,43 @@
+import datetime
+import random
+
+import pandas as pd
+
+stations = ["SIO", "AIG", "ZER"]
+
+start_ts = datetime.datetime(2026, 1, 1)
+end_ts = datetime.datetime(2027, 1, 1)
+delta = end_ts - start_ts
+
+min_temp, max_temp = -30.0, 100.0
+min_pres, max_pres = 800.0, 1100.0
+min_hum, max_hum = 0.0, 100.0
+
+N = 3000
+
+rows: list[tuple[str, datetime.datetime, float, float, float]] = []
+
+for _ in range(N):
+    ts = random.random() * delta + start_ts
+    rows.append(
+        (
+            random.choice(stations),
+            ts,
+            random.random() * (max_temp - min_temp) + min_temp,
+            random.random() * (max_pres - min_pres) + min_pres,
+            random.random() * (max_hum - min_hum) + min_hum,
+        )
+    )
+
+df = pd.DataFrame(
+    rows,
+    columns=[
+        "station_id",
+        "timestamp",
+        "temperature",
+        "pressure",
+        "humidity",
+    ],
+)
+df = df.sort_values(by=["timestamp", "station_id"])
+df.to_csv("data.csv", index=False)

examples/02_demonstration/weather/pipeline.py

@@ -0,0 +1,81 @@
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pandas as pd
+from custom_types import DailyAverages, Data, DataWithHI, HeatIndex, RawData
+
+from midas.typing import Column, cast, unsafe_cast
+
+
+def load_data(path: Path) -> RawData:
+    # Check base types and dataframe structure
+    return cast(RawData, pd.read_csv(path))
+
+
+def convert_data(raw_df: RawData) -> Data:
+    new_df = raw_df.copy()
+    new_df["timestamp"] = cast(
+        Column[object],
+        pd.to_datetime(new_df["timestamp"]),
+    )
+
+    # Check types and constraints at runtime, catches out-of-range values and
+    # invalid types / malformed data
+    return cast(Data, new_df)
+
+
+def compute_heat_index(df: Data):
+    # The computation's result can only be typed as `Column[float]`
+    # Casting is necessary to bring back semantic
+    df["heat_index"] = cast(
+        Column[HeatIndex],
+        (
+            df["temperature"] * 2.0
+            + df["humidity"] * 10.0
+            - df["temperature"] * df["humidity"] * 0.2
+        ),
+    )
+    return df
+
+
+def daily_avg(df: DataWithHI):
+    # Group-by and aggregation methods keep the structure of the dataframe but
+    # may erase the exact types
+    # The type checker is still very conservative and often the result of most
+    # aggregation methods as `Column[Any]`
+    return cast(
+        DailyAverages,
+        df.groupby(
+            by=[
+                "station_id",
+                df["timestamp"].dt.day.rename("day"),
+            ],
+        )
+        .mean()
+        .sort_values(by="timestamp"),
+    )
+
+
+def plot(df: DailyAverages):
+    # Some operations are not implemented in Midas but the user can still use
+    # them, they will just not be fully type-checked
+    # `unsafe_cast` can also be used to avoid trivial, redundant or costly checks
+    stations = unsafe_cast(list[str], list(df.index.get_level_values(0).unique()))
+    for station in stations:
+        sub_df = unsafe_cast(DailyAverages, df.loc[station])
+        plt.plot(sub_df["timestamp"], sub_df["heat_index"])
+    plt.show()
+
+
+def main():
+    # Assigning to annotated variables help catch errors
+    raw_df: RawData = load_data(Path("data.csv"))
+    df: Data = convert_data(raw_df)
+
+    with_hi = compute_heat_index(df)
+    dailies = daily_avg(with_hi)
+    plot(dailies)
+
+
+if __name__ == "__main__":
+    main()

gen/gen.py

@@ -1,3 +1,9 @@
+"""
+Helper script to generate AST nodes for Midas and Python.
+
+Takes in simple templates and generates full dataclasses and a visitor interface
+"""
+
 import re
 from pathlib import Path
 
@@ -30,6 +36,7 @@ from __future__ import annotations
 
 T = TypeVar("T")
 
+{preamble}
 {sections}
 """
 
@@ -57,6 +64,11 @@ IMPORTS_REGEX = re.compile(
     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("_")
@@ -88,6 +100,7 @@ def make_banner(text: str) -> str:
 
 
 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")
@@ -107,6 +120,7 @@ def make_section(full_name: str, base: str, param: str, body: str) -> str:
 
 
 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()
@@ -116,6 +130,10 @@ def generate(definitions_path: Path, out_path: Path):
     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")
@@ -129,6 +147,7 @@ def generate(definitions_path: Path, out_path: Path):
             gen_path=Path(__file__).relative_to(root_dir),
         ),
         imports=imports,
+        preamble=preamble,
         sections="\n\n\n".join(sections),
     )
     out_path.write_text(result)

gen/midas.py

@@ -4,6 +4,7 @@
 ###> 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
@@ -12,40 +13,58 @@ from midas.lexer.token import Token
 ###<
 
 
-###> Stmt | Statements
-class TypeStmt:
-    name: Token
-    params: list[Param]
-    type: Type
-
+###> Preamble
 @dataclass(frozen=True, kw_only=True)
-    class Param:
+class TypeParam:
     location: Location
     name: Token
     bound: Optional[Type]
 
 
-class PropertyStmt:
+class MemberKind(Enum):
+    PROPERTY = auto()
+    METHOD = auto()
+
+
+@dataclass(frozen=True, kw_only=True)
+class ParamSpec:
+    l_paren: Token
+    pos: list[FunctionType.Parameter]
+    mixed: list[FunctionType.Parameter]
+    kw: list[FunctionType.Parameter]
+
+
+###<
+
+
+###> Stmt | Statements
+class TypeStmt:
+    name: Token
+    params: list[TypeParam]
+    type: Type
+
+
+class AliasStmt:
     name: Token
     type: Type
 
 
+class MemberStmt:
+    name: Token
+    type: Type
+    kind: MemberKind
+
+
 class ExtendStmt:
-    type: Type
-    operations: list[OpStmt]
-
-
-class OpStmt:
     name: Token
-    operand: Type
-    result: Type
+    params: list[TypeParam]
+    members: list[MemberStmt]
 
 
 class PredicateStmt:
     name: Token
-    subject: Token
-    type: Type
-    condition: Expr
+    params: list[ParamSpec]
+    body: Expr
 
 
 ###<
@@ -71,6 +90,12 @@ class UnaryExpr:
     right: Expr
 
 
+class CallExpr:
+    callee: Expr
+    arguments: list[Expr]
+    keywords: dict[str, Expr]
+
+
 class GetExpr:
     expr: Expr
     name: Token
@@ -103,7 +128,7 @@ class NamedType:
 
 class GenericType:
     type: Type
-    params: list[Type]
+    args: list[Type]
 
 
 class ConstraintType:
@@ -111,12 +136,26 @@ class ConstraintType:
     constraint: Expr
 
 
-class UnionType:
-    types: list[Type]
+class FunctionType:
+    params: ParamSpec
+    returns: Type
+
+    @dataclass(frozen=True, kw_only=True)
+    class Parameter:
+        location: Optional[Location] = None
+        name: Optional[Token]
+        type: Type
+        required: bool
 
 
-class ComplexType:
-    properties: list[PropertyStmt]
+class FrameType:
+    columns: list[Column]
+
+    @dataclass(frozen=True, kw_only=True)
+    class Column:
+        location: Optional[Location] = None
+        name: Token
+        type: Type
 
 
 ###<

gen/python.py

@@ -12,10 +12,36 @@ from midas.ast.location import Location
 ###<
 
 
+###> Preamble
+@dataclass(frozen=True, kw_only=True)
+class ParamSpec:
+    pos: list[Function.Parameter]
+    mixed: list[Function.Parameter]
+    kw: list[Function.Parameter]
+
+    @property
+    def all(self) -> list[Function.Parameter]:
+        return self.pos + self.mixed + self.kw
+
+
+@dataclass(frozen=True, kw_only=True)
+class ImportAlias:
+    location: Location
+    name: str
+    alias: Optional[str] = None
+
+    @property
+    def imported_name(self) -> str:
+        return self.alias if self.alias is not None else self.name
+
+
+###<
+
+
 ###> MidasType | Type annotations | node
 class BaseType:
     base: str
-    param: Optional[MidasType]
+    args: tuple[MidasType, ...]
 
 
 class ConstraintType:
@@ -42,25 +68,17 @@ class ExpressionStmt:
 
 class Function:
     name: str
-    posonlyargs: list[Argument]
-    args: list[Argument]
-    sink: Optional[Argument]
-    kwonlyargs: list[Argument]
-    kw_sink: Optional[Argument]
+    params: ParamSpec
     returns: Optional[MidasType]
     body: list[Stmt]
 
     @dataclass(frozen=True, kw_only=True)
-    class Argument:
+    class Parameter:
         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
@@ -82,6 +100,30 @@ class IfStmt:
     orelse: list[Stmt]
 
 
+class Pass:
+    pass
+
+
+class ForStmt:
+    target: Expr
+    iterator: Expr
+    body: list[Stmt]
+
+
+class ImportStmt:
+    imports: list[ImportAlias]
+
+
+class FromImportStmt:
+    module: Optional[str]
+    imports: list[ImportAlias]
+    level: int
+
+
+class RawStmt:
+    stmt: ast.stmt
+
+
 ###<
 
 
@@ -128,15 +170,10 @@ class LogicalExpr:
     right: Expr
 
 
-class SetExpr:
-    object: Expr
-    name: str
-    value: Expr
-
-
 class CastExpr:
     type: MidasType
     expr: Expr
+    unsafe: bool
 
 
 class TernaryExpr:
@@ -145,4 +182,32 @@ class TernaryExpr:
     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 TupleExpr:
+    items: tuple[Expr, ...]
+
+
+class RawExpr:
+    expr: ast.expr
+
+
 ###<

midas/ast/location.py

@@ -13,6 +13,8 @@ class HasLocation(Protocol):
 
 @dataclass(frozen=True, kw_only=True)
 class Location:
+    """Information about the location of an AST node"""
+
     lineno: int
     col_offset: int
     end_lineno: Optional[int]
@@ -29,6 +31,16 @@ class Location:
 
     @staticmethod
     def span(start: Location, end: Location) -> Location:
+        """Create a new location spanning from one location to another
+
+        Args:
+            start (Location): the starting location
+            end (Location): the end location
+
+        Returns:
+            Location: a new location spanning from the start of `start`
+                to the end of `end`
+        """
         return Location(
             lineno=start.lineno,
             col_offset=start.col_offset,

midas/ast/midas.py

@@ -7,6 +7,7 @@ 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
@@ -14,6 +15,26 @@ 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()
+
+
+@dataclass(frozen=True, kw_only=True)
+class ParamSpec:
+    l_paren: Token
+    pos: list[FunctionType.Parameter]
+    mixed: list[FunctionType.Parameter]
+    kw: list[FunctionType.Parameter]
+
+
 ##############
 # Statements #
 ##############
@@ -31,14 +52,14 @@ class Stmt(ABC):
         def visit_type_stmt(self, stmt: TypeStmt) -> T: ...
 
         @abstractmethod
-        def visit_property_stmt(self, stmt: PropertyStmt) -> T: ...
+        def visit_alias_stmt(self, stmt: AliasStmt) -> T: ...
+
+        @abstractmethod
+        def visit_member_stmt(self, stmt: MemberStmt) -> T: ...
 
         @abstractmethod
         def visit_extend_stmt(self, stmt: ExtendStmt) -> T: ...
 
-        @abstractmethod
-        def visit_op_stmt(self, stmt: OpStmt) -> T: ...
-
         @abstractmethod
         def visit_predicate_stmt(self, stmt: PredicateStmt) -> T: ...
 
@@ -46,53 +67,47 @@ class Stmt(ABC):
 @dataclass(frozen=True)
 class TypeStmt(Stmt):
     name: Token
-    params: list[Param]
+    params: list[TypeParam]
     type: Type
 
-    @dataclass(frozen=True, kw_only=True)
-    class Param:
-        location: Location
-        name: Token
-        bound: Optional[Type]
-
     def accept(self, visitor: Stmt.Visitor[T]) -> T:
         return visitor.visit_type_stmt(self)
 
 
 @dataclass(frozen=True)
-class PropertyStmt(Stmt):
+class AliasStmt(Stmt):
     name: Token
     type: Type
 
     def accept(self, visitor: Stmt.Visitor[T]) -> T:
-        return visitor.visit_property_stmt(self)
+        return visitor.visit_alias_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):
-    type: Type
-    operations: list[OpStmt]
+    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 OpStmt(Stmt):
-    name: Token
-    operand: Type
-    result: Type
-
-    def accept(self, visitor: Stmt.Visitor[T]) -> T:
-        return visitor.visit_op_stmt(self)
-
-
 @dataclass(frozen=True)
 class PredicateStmt(Stmt):
     name: Token
-    subject: Token
-    type: Type
-    condition: Expr
+    params: list[ParamSpec]
+    body: Expr
 
     def accept(self, visitor: Stmt.Visitor[T]) -> T:
         return visitor.visit_predicate_stmt(self)
@@ -120,6 +135,9 @@ class Expr(ABC):
         @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: ...
 
@@ -165,6 +183,16 @@ class UnaryExpr(Expr):
         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):
     expr: Expr
@@ -229,10 +257,10 @@ class Type(ABC):
         def visit_constraint_type(self, type: ConstraintType) -> T: ...
 
         @abstractmethod
-        def visit_union_type(self, type: UnionType) -> T: ...
+        def visit_function_type(self, type: FunctionType) -> T: ...
 
         @abstractmethod
-        def visit_complex_type(self, type: ComplexType) -> T: ...
+        def visit_frame_type(self, type: FrameType) -> T: ...
 
 
 @dataclass(frozen=True)
@@ -246,7 +274,7 @@ class NamedType(Type):
 @dataclass(frozen=True)
 class GenericType(Type):
     type: Type
-    params: list[Type]
+    args: list[Type]
 
     def accept(self, visitor: Type.Visitor[T]) -> T:
         return visitor.visit_generic_type(self)
@@ -262,16 +290,30 @@ class ConstraintType(Type):
 
 
 @dataclass(frozen=True)
-class UnionType(Type):
-    types: list[Type]
+class FunctionType(Type):
+    params: ParamSpec
+    returns: Type
+
+    @dataclass(frozen=True, kw_only=True)
+    class Parameter:
+        location: Optional[Location] = None
+        name: Optional[Token]
+        type: Type
+        required: bool
 
     def accept(self, visitor: Type.Visitor[T]) -> T:
-        return visitor.visit_union_type(self)
+        return visitor.visit_function_type(self)
 
 
 @dataclass(frozen=True)
-class ComplexType(Type):
-    properties: list[PropertyStmt]
+class FrameType(Type):
+    columns: list[Column]
+
+    @dataclass(frozen=True, kw_only=True)
+    class Column:
+        location: Optional[Location] = None
+        name: Token
+        type: Type
 
     def accept(self, visitor: Type.Visitor[T]) -> T:
-        return visitor.visit_complex_type(self)
+        return visitor.visit_frame_type(self)

midas/ast/printer.py

@@ -1,654 +0,0 @@
-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_stmt_param(param)
-            self._write_line("type", last=True)
-            with self._child_level(single=True):
-                stmt.type.accept(self)
-
-    def _print_type_stmt_param(self, param: m.TypeStmt.Param) -> 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_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(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("type")
-            with self._child_level(single=True):
-                stmt.type.accept(self)
-            self._write_line("operations", last=True)
-            with self._child_level():
-                for i, op in enumerate(stmt.operations):
-                    self._idx = i
-                    if i == len(stmt.operations) - 1:
-                        self._mark_last()
-                    op.accept(self)
-
-    def visit_op_stmt(self, stmt: m.OpStmt) -> None:
-        self._write_line("OpStmt")
-        with self._child_level():
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-
-            self._write_line("operand")
-            with self._child_level(single=True):
-                stmt.operand.accept(self)
-
-            self._write_line("result", last=True)
-            with self._child_level(single=True):
-                stmt.result.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("params", last=True)
-            with self._child_level():
-                for i, param in enumerate(type.params):
-                    self._idx = i
-                    if i == len(type.params) - 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_union_type(self, type: m.UnionType) -> None:
-        self._write_line("UnionType")
-        with self._child_level():
-            self._write_line("types", last=True)
-            with self._child_level():
-                for i, type_ in enumerate(type.types):
-                    self._idx = i
-                    if i == len(type.types) - 1:
-                        self._mark_last()
-                    type_.accept(self)
-
-    def visit_complex_type(self, type: m.ComplexType) -> None:
-        self._write_line("ComplexType")
-        with self._child_level():
-            self._write_line("properties", last=True)
-            with self._child_level():
-                for i, prop in enumerate(type.properties):
-                    self._idx = i
-                    if i == len(type.properties) - 1:
-                        self._mark_last()
-                    prop.accept(self)
-
-
-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):
-        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_template_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_template_param(self, param: m.TypeStmt.Param) -> str:
-        res: str = param.name.lexeme
-        if param.bound is not None:
-            res += "<:" + param.bound.accept(self)
-        return res
-
-    def visit_property_stmt(self, stmt: m.PropertyStmt):
-        res: str = f"{stmt.name.lexeme}: {stmt.type.accept(self)}"
-        return self.indented(res)
-
-    def visit_extend_stmt(self, stmt: m.ExtendStmt):
-        res: str = self.indented(f"extend {stmt.type.accept(self)}")
-        res += " {\n"
-        self.level += 1
-        for op in stmt.operations:
-            res += op.accept(self)
-        self.level -= 1
-        res += "\n" + self.indented("}")
-        return res
-
-    def visit_op_stmt(self, stmt: m.OpStmt):
-        operand: str = stmt.operand.accept(self)
-        result: str = stmt.result.accept(self)
-        return self.indented(f"op {stmt.name.lexeme}({operand}) -> {result}")
-
-    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.params) != 0:
-            params: list[str] = [param.accept(self) for param in type.params]
-            res += f"[{', '.join(params)}]"
-        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_union_type(self, type: m.UnionType) -> str:
-        types: list[str] = [type_.accept(self) for type_ in type.types]
-        return " | ".join(types)
-
-    def visit_complex_type(self, type: m.ComplexType) -> str:
-        res: str = "{\n"
-        self.level += 1
-        for prop in type.properties:
-            res += prop.accept(self)
-            res += "\n"
-        self.level -= 1
-        res += self.indented("}")
-        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_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}")
-
-    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_set_expr(self, expr: p.SetExpr) -> None:
-        self._write_line("SetExpr")
-        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}")
-            self._write_line("value", last=True)
-            with self._child_level(single=True):
-                expr.value.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)

midas/ast/printer/__init__.py

@@ -0,0 +1,3 @@
+from .midas import MidasPrinter as MidasPrinter
+from .midas_ast import MidasAstPrinter as MidasAstPrinter
+from .python_ast import PythonAstPrinter as PythonAstPrinter

midas/ast/printer/base.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import io
+from contextlib import contextmanager
+from enum import Enum, auto
+from typing import Callable, Generator, Generic, Optional, Protocol, Sequence, TypeVar
+
+
+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)
+
+    def _write_sequence(
+        self,
+        label: str,
+        list_: Sequence[T],
+        *,
+        last: bool = False,
+        print_func: Optional[Callable[[T], None]] = None,
+    ):
+        if last:
+            self._mark_last()
+
+        self._write_line(label)
+        with self._child_level():
+            for i, item in enumerate(list_):
+                self._idx = i
+                if i == len(list_) - 1:
+                    self._mark_last()
+                if print_func is not None:
+                    print_func(item)
+                else:
+                    item.accept(self)

midas/ast/printer/midas.py

@@ -0,0 +1,170 @@
+import midas.ast.midas as m
+
+
+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)
+
+    # Statements
+
+    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 visit_alias_stmt(self, stmt: m.AliasStmt) -> str:
+        return self.indented(f"alias {stmt.name.lexeme} = {stmt.type.accept(self)}")
+
+    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
+        sig: str = "".join(self._visit_param_spec(spec) for spec in stmt.params)
+        body: str = stmt.body.accept(self)
+        return self.indented(f"predicate {name}{sig} = {body}")
+
+    # Expressions
+
+    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_call_expr(self, expr: m.CallExpr) -> str:
+        args: list[str] = [arg.accept(self) for arg in expr.arguments] + [
+            f"{name}={arg.accept(self)}" for name, arg in expr.keywords.items()
+        ]
+        return f"{expr.callee.accept(self)}({', '.join(args)})"
+
+    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 "_"
+
+    # Types
+
+    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_function_type(self, type: m.FunctionType) -> str:
+        spec: str = self._visit_param_spec(type.params)
+        return f"fn {spec} -> {type.returns.accept(self)}"
+
+    def _visit_param_spec(self, spec: m.ParamSpec) -> str:
+        pos: list[str] = [self._print_param(param) for param in spec.pos]
+        mixed: list[str] = [self._print_param(param) for param in spec.mixed]
+        kw: list[str] = [self._print_param(param) for param in spec.kw]
+        params: list[str] = pos
+
+        if len(pos) != 0:
+            params.append("/")
+        params += mixed
+        if len(kw) != 0:
+            params.append("*")
+        params += kw
+        return f"({', '.join(params)})"
+
+    def _print_param(self, param: m.FunctionType.Parameter) -> str:
+        res: str = ""
+        if param.name is not None:
+            res += param.name.lexeme
+            res += ": "
+        res += param.type.accept(self)
+        if not param.required:
+            res += "?"
+        return res
+
+    def visit_frame_type(self, type: m.FrameType) -> str:
+        res: str = self.indented("Frame[")
+        if len(type.columns) != 0:
+            res += "\n"
+            self.level += 1
+            columns: list[str] = []
+            for column in type.columns:
+                columns.append(self.indented(self._print_frame_column(column)))
+            res += ",\n".join(columns)
+            self.level -= 1
+            res += "\n"
+        res += "]"
+        return res
+
+    def _print_frame_column(self, column: m.FrameType.Column) -> str:
+        return f"{column.name.lexeme}: {column.type.accept(self)}"

midas/ast/printer/midas_ast.py

@@ -0,0 +1,238 @@
+import midas.ast.midas as m
+from midas.ast.printer.base import AstPrinter
+
+
+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_sequence(
+                "params",
+                stmt.params,
+                print_func=self._print_type_param,
+            )
+            self._write_line("type", last=True)
+            with self._child_level(single=True):
+                stmt.type.accept(self)
+
+    def visit_alias_stmt(self, stmt: m.AliasStmt) -> None:
+        self._write_line("AliasStmt")
+        with self._child_level():
+            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 _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(f'name: "{stmt.name.lexeme}"')
+            self._write_sequence(
+                "params",
+                stmt.params,
+                print_func=self._print_type_param,
+            )
+            self._write_sequence("members", stmt.members, last=True)
+
+    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_sequence(
+                "params",
+                stmt.params,
+                print_func=self._visit_param_spec,
+            )
+            self._write_line("body", last=True)
+            with self._child_level(single=True):
+                stmt.body.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_call_expr(self, expr: m.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_sequence("arguments", expr.arguments)
+            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: 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")
+
+    # Types
+
+    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_sequence("args", type.args, last=True)
+
+    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_function_type(self, type: m.FunctionType) -> None:
+        self._write_line("FunctionType")
+        with self._child_level():
+            self._write_line("params")
+            with self._child_level(single=True):
+                self._visit_param_spec(type.params)
+
+            self._write_line("returns", last=True)
+            with self._child_level(single=True):
+                type.returns.accept(self)
+
+    def _visit_param_spec(self, spec: m.ParamSpec) -> None:
+        self._write_line("ParamSpec")
+        with self._child_level():
+            self._write_sequence(
+                "pos",
+                spec.pos,
+                print_func=self._print_param,
+            )
+            self._write_sequence(
+                "mixed",
+                spec.mixed,
+                print_func=self._print_param,
+            )
+            self._write_sequence(
+                "kw",
+                spec.kw,
+                print_func=self._print_param,
+                last=True,
+            )
+
+    def _print_param(self, param: m.FunctionType.Parameter) -> None:
+        self._write_line("Parameter")
+        with self._child_level():
+            name: str = "None"
+            if param.name is not None:
+                name = f'"{param.name.lexeme}"'
+            self._write_line(f"name: {name}")
+            self._write_line("type")
+            with self._child_level(single=True):
+                param.type.accept(self)
+            self._write_line(f"required: {param.required}", last=True)
+
+    def visit_frame_type(self, type: m.FrameType) -> None:
+        self._write_line("FrameType")
+        with self._child_level(single=True):
+            self._write_sequence(
+                "columns",
+                type.columns,
+                print_func=self._print_frame_column,
+            )
+
+    def _print_frame_column(self, column: m.FrameType.Column) -> None:
+        self._write_line("Column")
+        with self._child_level():
+            self._write_line(f'name: "{column.name.lexeme}"')
+            self._write_line("type")
+            with self._child_level(single=True):
+                column.type.accept(self)

midas/ast/printer/python_ast.py

@@ -0,0 +1,285 @@
+import ast
+
+import midas.ast.python as p
+from midas.ast.printer.base import AstPrinter
+
+
+class PythonAstPrinter(
+    AstPrinter,
+    p.MidasType.Visitor[None],
+    p.Stmt.Visitor[None],
+    p.Expr.Visitor[None],
+):
+    # Types
+
+    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_sequence("args", node.args, 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(single=True):
+            self._write_sequence("columns", node.columns)
+
+    # Statements
+
+    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("params")
+            with self._child_level():
+                self._print_param_spec(stmt.params)
+
+            self._write_optional_child("returns", stmt.returns)
+            self._write_sequence("body", stmt.body, last=True)
+
+    def _print_param_spec(self, spec: p.ParamSpec) -> None:
+        self._write_line("ParamSpec")
+        with self._child_level():
+            self._write_sequence(
+                "pos",
+                spec.pos,
+                print_func=self._print_param,
+            )
+            self._write_sequence(
+                "mixed",
+                spec.mixed,
+                print_func=self._print_param,
+            )
+            self._write_sequence(
+                "kw",
+                spec.kw,
+                print_func=self._print_param,
+                last=True,
+            )
+
+    def _print_param(self, param: p.Function.Parameter) -> None:
+        self._write_line("Parameter")
+        with self._child_level():
+            self._write_line(f"name: {param.name}")
+            self._write_optional_child("type", param.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_sequence("targets", stmt.targets)
+            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_sequence("body", stmt.body)
+            self._write_sequence("orelse", stmt.orelse, last=True)
+
+    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_sequence("body", stmt.body, last=True)
+
+    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)}")
+
+    # Expressions
+
+    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_sequence("arguments", expr.arguments)
+            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")
+            with self._child_level(single=True):
+                expr.expr.accept(self)
+            self._write_line(f"unsafe: {expr.unsafe}", last=True)
+
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> None:
+        self._write_line("TernaryExpr")
+        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_sequence("items", expr.items, last=True)
+
+    def visit_dict_expr(self, expr: p.DictExpr) -> None:
+        self._write_line("DictExpr")
+        with self._child_level():
+            self._write_sequence(
+                "keys",
+                expr.keys,
+                print_func=lambda k: (
+                    self._write_line("None") if k is None else k.accept(self)
+                ),
+            )
+            self._write_sequence("values", expr.values, last=True)
+
+    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_tuple_expr(self, expr: p.TupleExpr) -> None:
+        self._write_line("TupleExpr")
+        with self._child_level():
+            self._write_sequence("items", expr.items, 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

@@ -14,6 +14,28 @@ from midas.ast.location import Location
 
 T = TypeVar("T")
 
+@dataclass(frozen=True, kw_only=True)
+class ParamSpec:
+    pos: list[Function.Parameter]
+    mixed: list[Function.Parameter]
+    kw: list[Function.Parameter]
+
+    @property
+    def all(self) -> list[Function.Parameter]:
+        return self.pos + self.mixed + self.kw
+
+
+@dataclass(frozen=True, kw_only=True)
+class ImportAlias:
+    location: Location
+    name: str
+    alias: Optional[str] = None
+
+    @property
+    def imported_name(self) -> str:
+        return self.alias if self.alias is not None else self.name
+
+
 ####################
 # Type annotations #
 ####################
@@ -43,7 +65,7 @@ class MidasType(ABC):
 @dataclass(frozen=True)
 class BaseType(MidasType):
     base: str
-    param: Optional[MidasType]
+    args: tuple[MidasType, ...]
 
     def accept(self, visitor: MidasType.Visitor[T]) -> T:
         return visitor.visit_base_type(self)
@@ -106,6 +128,21 @@ class Stmt(ABC):
         @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_import_stmt(self, stmt: ImportStmt) -> T: ...
+
+        @abstractmethod
+        def visit_from_import_stmt(self, stmt: FromImportStmt) -> T: ...
+
+        @abstractmethod
+        def visit_raw_stmt(self, stmt: RawStmt) -> T: ...
+
 
 @dataclass(frozen=True)
 class ExpressionStmt(Stmt):
@@ -118,25 +155,17 @@ class ExpressionStmt(Stmt):
 @dataclass(frozen=True)
 class Function(Stmt):
     name: str
-    posonlyargs: list[Argument]
-    args: list[Argument]
-    sink: Optional[Argument]
-    kwonlyargs: list[Argument]
-    kw_sink: Optional[Argument]
+    params: ParamSpec
     returns: Optional[MidasType]
     body: list[Stmt]
 
     @dataclass(frozen=True, kw_only=True)
-    class Argument:
+    class Parameter:
         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)
 
@@ -177,6 +206,50 @@ class IfStmt(Stmt):
         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 ImportStmt(Stmt):
+    imports: list[ImportAlias]
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_import_stmt(self)
+
+
+@dataclass(frozen=True)
+class FromImportStmt(Stmt):
+    module: Optional[str]
+    imports: list[ImportAlias]
+    level: int
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_from_import_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 #
 ###############
@@ -214,15 +287,30 @@ class Expr(ABC):
         @abstractmethod
         def visit_logical_expr(self, expr: LogicalExpr) -> T: ...
 
-        @abstractmethod
-        def visit_set_expr(self, expr: SetExpr) -> 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_tuple_expr(self, expr: TupleExpr) -> T: ...
+
+        @abstractmethod
+        def visit_raw_expr(self, expr: RawExpr) -> T: ...
+
 
 @dataclass(frozen=True)
 class BinaryExpr(Expr):
@@ -298,20 +386,11 @@ class LogicalExpr(Expr):
         return visitor.visit_logical_expr(self)
 
 
-@dataclass(frozen=True)
-class SetExpr(Expr):
-    object: Expr
-    name: str
-    value: Expr
-
-    def accept(self, visitor: Expr.Visitor[T]) -> T:
-        return visitor.visit_set_expr(self)
-
-
 @dataclass(frozen=True)
 class CastExpr(Expr):
     type: MidasType
     expr: Expr
+    unsafe: bool
 
     def accept(self, visitor: Expr.Visitor[T]) -> T:
         return visitor.visit_cast_expr(self)
@@ -325,3 +404,55 @@ class TernaryExpr(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 TupleExpr(Expr):
+    items: tuple[Expr, ...]
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_tuple_expr(self)
+
+
+@dataclass(frozen=True)
+class RawExpr(Expr):
+    expr: ast.expr
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_raw_expr(self)

midas/checker/builtins.midas

@@ -0,0 +1,277 @@
+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]
+
+}
+
+extend str {
+    def capitalize: fn() -> str
+    def casefold: fn() -> str
+    def center: fn(width: int, fillchar: str?, /) -> str
+    def count: fn(sub: str, start: None?, end: None?, /) -> int
+    def count: fn(sub: str, start: int, end: None?, /) -> int
+    def count: fn(sub: str, start: None, end: int, /) -> int
+    def count: fn(sub: str, start: int, end: int, /) -> int
+    def encode: fn(encoding: str?, errors: str?) -> bytes
+    def endswith: fn(suffix: str, start: None?, end: None?, /) -> bool
+    def endswith: fn(suffix: str, start: int, end: None?, /) -> bool
+    def endswith: fn(suffix: str, start: None, end: int, /) -> bool
+    def endswith: fn(suffix: str, start: int, end: int, /) -> bool
+    def expandtabs: fn(tabsize: int?) -> str
+    def find: fn(sub: str, start: None?, end: None?, /) -> int
+    def find: fn(sub: str, start: int, end: None?, /) -> int
+    def find: fn(sub: str, start: None, end: int, /) -> int
+    def find: fn(sub: str, start: int, end: int, /) -> int
+    // def format: fn(*args: object, **kwargs: object) -> str
+    // def format_map: fn(mapping: _FormatMapMapping, /) -> str
+    def index: fn(sub: str, start: None?, end: None?, /) -> int
+    def index: fn(sub: str, start: int, end: None?, /) -> int
+    def index: fn(sub: str, start: None, end: int, /) -> int
+    def index: fn(sub: str, start: int, end: int, /) -> int
+    def isalnum: fn() -> bool
+    def isalpha: fn() -> bool
+    def isascii: fn() -> bool
+    def isdecimal: fn() -> bool
+    def isdigit: fn() -> bool
+    def isidentifier: fn() -> bool
+    def islower: fn() -> bool
+    def isnumeric: fn() -> bool
+    def isprintable: fn() -> bool
+    def isspace: fn() -> bool
+    def istitle: fn() -> bool
+    def isupper: fn() -> bool
+    def join: fn(iterable: list[str], /) -> str  // TODO: use Iterable
+    def ljust: fn(width: int, fillchar: str?, /) -> str
+    def lower: fn() -> str
+    def lstrip: fn(chars: None?, /) -> str
+    def lstrip: fn(chars: str, /) -> str
+    def partition: fn(sep: str, /) -> tuple[str, str, str]
+    
+    def replace: fn(old: str, new: str, count: int?, /) -> str
+
+    def removeprefix: fn(prefix: str, /) -> str
+    def removesuffix: fn(suffix: str, /) -> str
+    def rfind: fn(sub: str, start: None?, end: None?, /) -> int
+    def rfind: fn(sub: str, start: int, end: None?, /) -> int
+    def rfind: fn(sub: str, start: None, end: int, /) -> int
+    def rfind: fn(sub: str, start: int, end: int, /) -> int
+    def rindex: fn(sub: str, start: None?, end: None?, /) -> int
+    def rindex: fn(sub: str, start: int, end: None?, /) -> int
+    def rindex: fn(sub: str, start: None, end: int, /) -> int
+    def rindex: fn(sub: str, start: int, end: int, /) -> int
+    def rjust: fn(width: int, fillchar: str?, /) -> str
+    def rpartition: fn(sep: str, /) -> tuple[str, str, str]
+    def rsplit: fn(sep: None?, maxsplit: int?) -> list[str]
+    def rsplit: fn(sep: str, maxsplit: int?) -> list[str]
+    def rstrip: fn(chars: None?, /) -> str
+    def rstrip: fn(chars: str, /) -> str
+    def split: fn(sep: None?, maxsplit: int?) -> list[str]
+    def split: fn(sep: str, maxsplit: int?) -> list[str]
+    def splitlines: fn(keepends: bool?) -> list[str]
+    def startswith: fn(prefix: str, start: None?, end: None?, /) -> bool
+    def startswith: fn(prefix: str, start: int, end: None?, /) -> bool
+    def startswith: fn(prefix: str, start: None, end: int, /) -> bool
+    def startswith: fn(prefix: str, start: int, end: int, /) -> bool
+    def strip: fn(chars: None?, /) -> str
+    def strip: fn(chars: str, /) -> str
+    def swapcase: fn() -> str
+    def title: fn() -> str
+    // def translate: fn(table: _TranslateTable, /) -> str
+    def upper: fn() -> str
+    def zfill: fn(width: int, /) -> str
+    def __add__: fn(value: str, /) -> str
+    // Incompatible with Sequence.__contains__
+    def __contains__: fn(key: str, /) -> bool
+    def __eq__: fn(value: object, /) -> bool
+    def __ge__: fn(value: str, /) -> bool
+    def __getitem__: fn(key: slice, /) -> str
+    def __getitem__: fn(key: int, /) -> str
+    def __gt__: fn(value: str, /) -> bool
+    def __hash__: fn() -> int
+    // def __iter__: fn() -> Iterator[str]
+    def __le__: fn(value: str, /) -> bool
+    def __len__: fn() -> int
+    def __lt__: fn(value: str, /) -> bool
+    def __mod__: fn(value: Any, /) -> str
+    def __mul__: fn(value: int, /) -> str
+    def __ne__: fn(value: object, /) -> bool
+    def __rmul__: fn(value: int, /) -> str
+    def __getnewargs__: fn() -> tuple[str]
+    def __format__: fn(format_spec: str, /) -> str
+}

midas/checker/builtins.py

@@ -0,0 +1,60 @@
+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]] = {
+    "object": {"float", "list", "dict", "str", "bytes", "tuple"},
+    "float": {"int"},
+}
+"""
+Hard-coded subtype relationships between builtin types
+
+Circular dependencies and diamond inheritance MUST be avoided
+"""
+
+
+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"))
+    bytes = reg.define_type("bytes", BaseType(name="bytes"))
+    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"))
+
+    tuple = reg.define_type("tuple", BaseType(name="tuple"))
+
+    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

@@ -1,546 +1,41 @@
-import logging
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
-import midas.ast.midas as m
-import midas.ast.python as p
-from midas.ast.location import Location
-from midas.checker.diagnostic import Diagnostic, DiagnosticType
-from midas.checker.environment import Environment
-from midas.checker.operators import COMPARATOR_METHODS, OPERATOR_METHODS
-from midas.checker.types import BaseType, Function, SimpleType, Type, UnitType, UnknownType
-from midas.lexer.midas import MidasLexer
-from midas.lexer.token import Token
-from midas.parser.midas import MidasParser
-from midas.resolver.midas import MidasResolver
+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 ReturnException(Exception):
-    pass
+class TypeChecker:
+    """Type checking dispatcher
 
-
-@dataclass(frozen=True, kw_only=True)
-class MappedArgument:
-    expr: p.Expr
-    type: Type
-    argument: Function.Argument
-
-
-class Checker(
-    p.Stmt.Visitor[None],
-    p.Expr.Visitor[Type],
-    p.MidasType.Visitor[Type],
-):
-    """A type checker which can use custom type definitions"""
-
-    def __init__(self, locals: dict[p.Expr, int], file_path: Path):
-        self.logger: logging.Logger = logging.getLogger("Checker")
-        self.file_path: Path = file_path
-        self.ctx: MidasResolver = MidasResolver()
-        self.global_env: Environment = Environment()
-        self.env: Environment = self.global_env
-        self.locals: dict[p.Expr, int] = locals
-        self.diagnostics: list[Diagnostic] = []
-
-    def diagnostic(self, type: DiagnosticType, location: Location, message: str):
-        self.diagnostics.append(
-            Diagnostic(
-                file_path=self.file_path,
-                location=location,
-                type=type,
-                message=message,
-            )
-        )
-
-    def error(self, location: Location, message: str):
-        self.diagnostic(
-            type=DiagnosticType.ERROR,
-            location=location,
-            message=message,
-        )
-
-    def warning(self, location: Location, message: str):
-        self.diagnostic(
-            type=DiagnosticType.WARNING,
-            location=location,
-            message=message,
-        )
-
-    def info(self, location: Location, message: str):
-        self.diagnostic(
-            type=DiagnosticType.INFO,
-            location=location,
-            message=message,
-        )
-
-    def evaluate(self, expr: p.Expr) -> Type:
-        """Evaluate the type of an expression
-
-        Args:
-            expr (p.Expr): the expression to evaluate
-
-        Returns:
-            Type: the type of the given expression
+    Contains a typer for Midas and one for Python, as well as the types registry
     """
-        return expr.accept(self)
 
-    def evaluate_block(self, block: list[p.Stmt], env: Environment) -> bool:
-        """Evaluate a sequence of statements
+    def __init__(self):
+        self.types: TypesRegistry = TypesRegistry()
+        self.reporter: Reporter = Reporter()
 
-        Args:
-            block (list[p.Stmt]): the statements to evaluate
-            env (Environment): the environment in which to evaluate
+        self.midas_typer = MidasTyper(self.types, self.reporter)
+        self.python_typer = PythonTyper(self.types, self.reporter)
 
-        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:
-                stmt.accept(self)
-            except ReturnException:
-                returned = True
-                if i < len(block) - 1:
-                    self.warning(block[i + 1].location, "Unreachable statement")
-                break
-        self.env = previous_env
-        return returned
+    def import_midas(self, path: Path):
+        source: str = path.read_text()
+        return self.import_midas_source(source, path=str(path))
 
-    def check(self, statements: list[p.Stmt]) -> list[Diagnostic]:
-        """Type check a sequence of statements and returns diagnostics
+    def import_midas_source(self, source: str, path: Optional[str] = None):
+        self.midas_typer.process(source, path)
 
-        Args:
-            statements (list[p.Stmt]): the statements to evaluate and check
+    def type_check(self, path: Path) -> TypedAST:
+        source: str = path.read_text()
+        return self.type_check_source(source, path=str(path))
 
-        Returns:
-            list[Diagnostic]: the list of diagnostics (errors, warning, etc.)
-        """
-        self.diagnostics = []
-        for stmt in statements:
-            stmt.accept(self)
+    def type_check_source(self, source: str, path: Optional[str] = None) -> TypedAST:
+        return self.python_typer.process(source, path)
 
-        self.logger.debug(f"Final environment: {self.env.flat_dict()}")
-        return self.diagnostics
-
-    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 parse_midas_import(self, expr: p.CallExpr) -> Optional[Path]:
-        """Parse a Midas import statement
-
-        The statement should be written as `midas.using("path/to/types.midas")`
-
-        Args:
-            expr (p.CallExpr): the import call expression
-
-        Returns:
-            Optional[Path]: the path to the imported file, or None if the expression is malformed
-        """
-        match expr:
-            case p.CallExpr(
-                callee=p.GetExpr(
-                    object=p.VariableExpr(name="midas"),
-                    name="using",
-                ),
-                arguments=[
-                    p.LiteralExpr(value=path),
-                ],
-            ):
-                return Path(path)
-        return None
-
-    def import_midas(self, path: Path) -> None:
-        """Import Midas definitions from a path
-
-        Args:
-            path (Path): the import path
-        """
-        self.logger.debug(f"Importing type definitions from {path}")
-        path = (self.file_path.parent / path).resolve()
-        lexer: MidasLexer = MidasLexer(path.read_text())
-        tokens: list[Token] = lexer.process()
-        parser: MidasParser = MidasParser(tokens)
-        stmts: list[m.Stmt] = parser.parse()
-        self.ctx.resolve(stmts)
-        self.logger.debug(f"Midas types: {self.ctx._types}")
-        self.logger.debug(f"Midas operations: {self.ctx._operations}")
-
-    def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
-        self.evaluate(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 arg.type.accept(self)
-            if arg.default is not None:
-                return arg.default.accept(self)
-            return UnknownType()
-
-        for arg in stmt.posonlyargs:
-            pos_args.append(
-                Function.Argument(
-                    name=arg.name,
-                    type=eval_arg_type(arg),
-                    required=arg.default is None,
-                )
-            )
-        for arg in stmt.args:
-            args.append(
-                Function.Argument(
-                    name=arg.name,
-                    type=eval_arg_type(arg),
-                    required=arg.default is None,
-                )
-            )
-        for arg in stmt.kwonlyargs:
-            kw_args.append(
-                Function.Argument(
-                    name=arg.name,
-                    type=eval_arg_type(arg),
-                    required=arg.default is None,
-                )
-            )
-
-        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 = stmt.returns.accept(self)
-            # Early define to handle simple fully-typed recursion
-            inside_function: Function = Function(
-                name=stmt.name,
-                pos_args=pos_args,
-                args=args,
-                kw_args=kw_args,
-                returns=returns_hint,
-            )
-            self.env.define(stmt.name, inside_function)
-
-        returned: bool = self.evaluate_block(stmt.body, env)
-        inferred_return: Type = UnknownType()
-        if not returned:
-            env.return_types.append(UnitType())
-        return_types: set[Type] = set(env.return_types)
-        if len(return_types) == 1:
-            inferred_return = list(return_types)[0]
-        elif len(return_types) > 1:
-            self.error(
-                stmt.location,
-                f"Mixed return types: {env.return_types}",
-            )
-
-        returns: Type = UnknownType()
-        if returns_hint is not None:
-            assert stmt.returns is not None
-            returns = returns_hint
-            if returns != inferred_return:
-                self.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(
-            name=stmt.name,
-            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 = stmt.type.accept(self)
-        self.env.define(stmt.name, type)
-
-    def visit_assign_stmt(self, stmt: p.AssignStmt) -> None:
-        value: Type = self.evaluate(stmt.value)
-        for target in stmt.targets:
-            if not isinstance(target, p.VariableExpr):
-                self.logger.warning(f"Unsupported assignment to {target}")
-                self.warning(target.location, f"Unsupported assignment to {target}")
-                continue
-            name: str = target.name
-            var_type: Optional[Type] = self.look_up_variable(name, target)
-
-            if var_type is None:
-                self.env.define(name, value)
-            else:
-                # TODO: implement real comparison method
-                if var_type != value:
-                    self.error(
-                        stmt.location,
-                        f"Cannot assign {value} to {name} of type {var_type}",
-                    )
-
-    def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
-        type: Type = stmt.value.accept(self) 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 = stmt.test.accept(self)
-
-        # TODO Allow subtypes or any type
-        if test_type != self.ctx.get_type("bool"):
-            self.error(
-                stmt.test.location, f"If test must be a boolean, got {test_type}"
-            )
-
-        env: Environment = Environment(self.env)
-        body_returned: bool = self.evaluate_block(stmt.body, env)
-        else_returned: bool = self.evaluate_block(stmt.orelse, env)
-        self.env.return_types.extend(env.return_types)
-        if body_returned and else_returned:
-            raise ReturnException()
-
-    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.warning(expr.location, f"Unsupported operator {expr.operator}")
-            return UnknownType()
-        left: Type = self.evaluate(expr.left)
-        right: Type = self.evaluate(expr.right)
-
-        result: Optional[Type] = self.ctx.get_operation_result(left, method, right)
-        if result is None:
-            self.error(
-                expr.location,
-                f"Undefined operation {method} between {left} and {right}",
-            )
-            return UnknownType()
-        return result
-
-    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.warning(expr.location, f"Unsupported operator {expr.operator}")
-            return UnknownType()
-        left: Type = self.evaluate(expr.left)
-        right: Type = self.evaluate(expr.right)
-
-        result: Optional[Type] = self.ctx.get_operation_result(left, method, right)
-        if result is None:
-            self.error(
-                expr.location,
-                f"Undefined operation {method} between {left} and {right}",
-            )
-            return UnknownType()
-        return result
-
-    def visit_unary_expr(self, expr: p.UnaryExpr) -> Type: ...
-
-    def visit_call_expr(self, expr: p.CallExpr) -> Type:
-        if path := self.parse_midas_import(expr):
-            self.import_midas(path)
-            return UnknownType()
-        callee: Type = self.evaluate(expr.callee)
-        if not isinstance(callee, Function):
-            self.error(expr.callee.location, "Callee is not a function")
-            return UnknownType()
-        function: Function = callee
-        mapped: list[MappedArgument] = self.map_call_arguments(function, expr)
-        for arg in mapped:
-            if arg.type != arg.argument.type:
-                self.error(
-                    arg.expr.location,
-                    f"Wrong type for argument '{arg.argument.name}', expected {arg.argument.type}, got {arg.type}",
-                )
-        return function.returns
-
-    def visit_get_expr(self, expr: p.GetExpr) -> Type: ...
-
-    def visit_literal_expr(self, expr: p.LiteralExpr) -> Type:
-        match expr.value:
-            case bool():  # Must be before int
-                return self.ctx.get_type("bool")
-            case int():
-                return self.ctx.get_type("int")
-            case float():
-                return self.ctx.get_type("float")
-            case str():
-                return self.ctx.get_type("str")
-            case _:
-                self.warning(expr.location, f"Unknown literal {expr}")
-                return UnknownType()
-
-    def visit_variable_expr(self, expr: p.VariableExpr) -> Type:
-        return self.look_up_variable(expr.name, expr) or UnknownType()
-
-    def visit_logical_expr(self, expr: p.LogicalExpr) -> Type: ...
-
-    def visit_set_expr(self, expr: p.SetExpr) -> Type: ...
-
-    def visit_cast_expr(self, expr: p.CastExpr) -> Type:
-        return expr.type.accept(self)
-
-    def visit_ternary_expr(self, expr: p.TernaryExpr) -> Type:
-        test_type: Type = expr.test.accept(self)
-
-        # TODO Allow subtypes or any type
-        if test_type != self.ctx.get_type("bool"):
-            self.error(
-                expr.test.location, f"If test must be a boolean, got {test_type}"
-            )
-
-        true_type: Type = expr.if_true.accept(self)
-        false_type: Type = expr.if_false.accept(self)
-        if true_type != false_type:
-            self.error(expr.location, f"Type mismatch in ternary if branches: true={true_type} != false={false_type}")
-            return UnknownType()
-        return true_type
-
-    def visit_base_type(self, node: p.BaseType) -> Type:
-        return self.ctx.get_type(node.base)
-
-    def visit_constraint_type(self, node: p.ConstraintType) -> Type: ...
-
-    def visit_frame_column(self, node: p.FrameColumn) -> Type: ...
-
-    def visit_frame_type(self, node: p.FrameType) -> Type: ...
-
-    def map_call_arguments(
-        self, function: Function, call: p.CallExpr
-    ) -> list[MappedArgument]:
-        """Map call arguments to function 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
-
-        Args:
-            function (Function): the function definition
-            call (p.CallExpr): the call expression
-
-        Returns:
-            list[MappedArgument]: the list of mapped arguments
-        """
-        positional: list[tuple[p.Expr, Type]] = [
-            (arg, self.evaluate(arg)) for arg in call.arguments
-        ]
-        keywords: dict[str, tuple[p.Expr, Type]] = {
-            name: (arg, self.evaluate(arg)) for name, arg in call.keywords.items()
-        }
-        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
-        }
-
-        # 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:
-                self.error(arg[0].location, "Too many positional arguments")
-                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 name in set_args:
-                    self.error(
-                        arg[0].location, f"Multiple values for argument '{name}'"
-                    )
-                else:
-                    self.error(arg[0].location, f"Unknown keyword argument '{name}'")
-                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)
-            self.error(
-                call.location,
-                f"Missing required positional argument{plural}: {args}",
-            )
-
-        if len(required_keyword) != 0:
-            plural: str = "" if len(required_keyword) == 1 else "s"
-            args: str = join_args(required_keyword)
-            self.error(
-                call.location,
-                f"Missing required keyword argument{plural}: {args}",
-            )
-
-        return mapped
+    @property
+    def diagnostics(self) -> list[Diagnostic]:
+        return self.reporter.diagnostics

midas/checker/diagnostic.py

@@ -1,6 +1,5 @@
 from dataclasses import dataclass
 from enum import StrEnum
-from pathlib import Path
 from typing import Optional
 
 from midas.ast.location import Location
@@ -10,16 +9,36 @@ class DiagnosticType(StrEnum):
     ERROR = "Error"
     WARNING = "Warning"
     INFO = "Info"
+    DEBUG = "Debug"
 
 
 @dataclass(frozen=True)
 class Diagnostic:
-    file_path: Path
+    """Information about a diagnostic (warning, errors, etc.)
+
+    Holds a location, a diagnostic type and a message.
+    Optionally bound to a file path
+    """
+
+    file_path: Optional[str]
     location: Location
     type: DiagnosticType
     message: str
 
-    def __str__(self) -> str:
+    @property
+    def location_str(self) -> str:
+        """Get diagnostic type and location as a human readable string
+
+        The location is formatted as "<Type> in <file> from L<start_line>:<start_col> to <end_line>:<end_col>",
+        for example: "Error in /home/user/Desktop/script.py from L12:5 to L12:8"
+
+        If the file is `None`, the "in ..." section is excluded from the result.<br>
+        If the location's end is not specified, the formulation "at L<start_line>:<start_col>" is used.
+
+        Returns:
+            str: the formatted type and location string
+        """
+
         start_loc: str = f"L{self.location.lineno}:{self.location.col_offset+1}"
         end_loc: Optional[str] = ""
         if (
@@ -27,7 +46,16 @@ class Diagnostic:
             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 = (
-            f"at {start_loc}" if end_loc is None else f"from {start_loc} to {end_loc}"
-        )
-        return f"{self.type} in {self.file_path} {loc}: {self.message}"
+
+        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/dispatcher.py

@@ -0,0 +1,532 @@
+import logging
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import Generic, Optional, Protocol, TypeVar, Union
+
+from midas.ast.location import Location
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter
+from midas.checker.types import (
+    AppliedType,
+    DerivedType,
+    Function,
+    GenericType,
+    OverloadedFunction,
+    Type,
+    UnknownType,
+)
+from midas.checker.unifier import Unifier
+
+
+class HasLocation(Protocol):
+    @property
+    def location(self) -> Location: ...
+
+
+E = TypeVar("E", bound=HasLocation)
+
+TypedExpr = tuple[E, Type]
+"""An expression and its type"""
+
+
+@dataclass(frozen=True, kw_only=True)
+class MappedArgument(Generic[E]):
+    """An argument passed in a call and the corresponding parameter"""
+
+    arg_expr: E
+    arg_type: Type
+    parameter: Function.Parameter
+
+
+@dataclass(frozen=True, kw_only=True)
+class OverloadCandidate:
+    """An overloaded function call candidate with its mapped arguments"""
+
+    function: Function
+    mapped: list[MappedArgument]
+
+
+class CallError(StrEnum):
+    """Reason of a call error"""
+
+    INVALID_ARGS = "Invalid arguments"
+    NO_MATCHING_OVERLOAD = "No matching overload"
+    IMPOSSIBLE_UNIFICATION = "Parameters unification failed"
+    NOT_CALLABLE = "Not callable"
+
+
+@dataclass(frozen=True, kw_only=True)
+class CallResult:
+    """The result of a function call
+
+    Holds a return type, an optional error reason and message
+    """
+
+    error: Optional[CallError] = None
+    """The reason of the error, if there is one"""
+
+    result: Type = UnknownType()
+    """The result type. `UnknownType()` if the call is invalid"""
+
+    message: Optional[str] = None
+    """An optional error message"""
+
+    @property
+    def is_valid(self) -> bool:
+        """Whether the call is valid (i.e. no error)"""
+        return self.error is None
+
+    @property
+    def error_message(self) -> str:
+        """A descriptive message for the error, if there is one"""
+        if self.message is not None:
+            return self.message
+        if self.error is not None:
+            return str(self.error)
+        return ""
+
+
+class CallDispatcher(Generic[E]):
+    """Helper class to handle dispatching calls and mapping arguments
+
+    This class is responsible for mapping call-site arguments to function
+    parameters, verifying the validity of calls and computing their
+    return types
+
+    :class:`CallDispatcher` is generic to handle AST nodes from both Midas and Python
+    """
+
+    def __init__(self, types: TypesRegistry, reporter: FileReporter) -> None:
+        self.types: TypesRegistry = types
+        self.reporter: FileReporter = reporter
+        self.logger: logging.Logger = logging.getLogger("CallDispatcher")
+
+    def set_reporter(self, reporter: FileReporter):
+        self.reporter = reporter
+
+    def get_result(
+        self,
+        location: Location,
+        callee: Type,
+        positional: list[TypedExpr[E]],
+        keywords: dict[str, TypedExpr[E]],
+        report_errors: bool = True,
+    ) -> CallResult:
+        """Get the result type of a function call
+
+        If the callee has overloads, this function will try to resolve the
+        appropriate signature.
+        Argument types are matched with the defined parameters.
+        This function doesn't take the raw expression as a parameter to
+        accommodate for desugared calls such as for operators.
+
+        Args:
+            location (Location): the call location
+            callee (Type): the called function
+            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:
+            CallResult: the call result, either a type or an error
+        """
+        match callee:
+            case Function() as function:
+                valid: bool
+                mapped: list[MappedArgument[E]]
+                valid, mapped = self.map_call_arguments(
+                    function, location, positional, keywords
+                )
+                valid = valid and self._are_arguments_valid(mapped, report_errors)
+                if not valid:
+                    return CallResult(error=CallError.INVALID_ARGS)
+                return CallResult(result=function.returns)
+
+            case OverloadedFunction(overloads=overloads):
+                res = self._match_overload(
+                    overloads, location, positional, keywords, report_errors
+                )
+                if res[0] is None:
+                    return CallResult(
+                        error=CallError.NO_MATCHING_OVERLOAD,
+                        message=res[1],
+                    )
+                return CallResult(result=res[0].returns)
+
+            case AppliedType(body=body):
+                return self.get_result(
+                    location, body, positional, keywords, report_errors
+                )
+
+            case UnknownType():
+                return CallResult(result=UnknownType())
+
+            case DerivedType(type=base):
+                return self.get_result(
+                    location, base, positional, keywords, report_errors
+                )
+
+            case GenericType():
+                unifier: Unifier = Unifier(self.types)
+                pos: list[Type] = [a[1] for a in positional]
+                kw: dict[str, Type] = {k: v[1] for k, v in keywords.items()}
+                unified: Optional[Type] = unifier.unify_call(callee, pos, kw)
+                if unified is None:
+                    pos_str: str = ", ".join(str(t) for t in pos)
+                    kw_str: str = ", ".join(f"{k}: {v}" for k, v in kw.items())
+                    message: str = (
+                        f"Could not unify {callee}={callee.body} with pos=[{pos_str}] and kw={{{kw_str}}}"
+                    )
+                    if report_errors:
+                        self.reporter.error(location, message)
+                    return CallResult(
+                        error=CallError.IMPOSSIBLE_UNIFICATION,
+                        message=message,
+                    )
+                return self.get_result(
+                    location,
+                    unified,
+                    positional,
+                    keywords,
+                    report_errors,
+                )
+
+            case _:
+                message: str = f"{callee} ({callee.__class__.__name__}) is not callable"
+                if report_errors:
+                    self.reporter.error(location, message)
+                return CallResult(
+                    error=CallError.NOT_CALLABLE,
+                    message=message,
+                )
+
+    def _unwrap_function(
+        self,
+        callee: Type,
+        positional: list[TypedExpr[E]],
+        keywords: dict[str, TypedExpr[E]],
+    ) -> Union[tuple[Function, None], tuple[None, CallError]]:
+        """Unwrap a type to get a callable `Function`
+
+        Args:
+            callee (Type): the called type
+            positional (list[TypedExpr[E]]): the list of positional arguments
+            keywords (dict[str, TypedExpr[E]]): the map of keyword arguments
+
+        Returns:
+            Union[tuple[Function, None], tuple[None, CallError]]: a tuple
+                containing the callable `Function` type, or `None` if it could
+                not be unwrapped, and an error, or `None` if there was none.
+        """
+        match callee:
+            case DerivedType(type=base):
+                return self._unwrap_function(base, positional, keywords)
+
+            case GenericType():
+                unifier: Unifier = Unifier(self.types)
+                unified: Optional[Type] = unifier.unify_call(
+                    callee,
+                    [a[1] for a in positional],
+                    {k: v[1] for k, v in keywords.items()},
+                )
+                if unified is None:
+                    return None, CallError.IMPOSSIBLE_UNIFICATION
+                return self._unwrap_function(unified, positional, keywords)
+
+            case Function():
+                return callee, None
+
+            case AppliedType(body=body):
+                return self._unwrap_function(body, positional, keywords)
+
+            case _:
+                return None, CallError.NOT_CALLABLE
+
+    def _are_arguments_valid(
+        self,
+        arguments: list[MappedArgument[E]],
+        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 arg.parameter.unsupported:
+                # Always report error
+                self.reporter.error(
+                    arg.arg_expr.location, f"Unsupported argument {arg.parameter.name}"
+                )
+
+            if not self.types.is_subtype(arg.arg_type, arg.parameter.type):
+                if report_errors:
+                    self.reporter.error(
+                        arg.arg_expr.location,
+                        f"Wrong type for argument '{arg.parameter.name}', expected {arg.parameter.type}, got {arg.arg_type}",
+                    )
+                valid = False
+        return valid
+
+    def _match_overload(
+        self,
+        overloads: list[Type],
+        location: Location,
+        positional: list[TypedExpr[E]],
+        keywords: dict[str, TypedExpr[E]],
+        report_errors: bool = True,
+    ) -> Union[tuple[Function, None], tuple[None, str]]:
+        """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:
+            Union[tuple[Function, None], tuple[None, str]]: a tuple containing
+                the resolved function signature if it can be determined
+                unambiguously, or `None`, and an error message, or `None`
+        """
+        candidates: list[OverloadCandidate] = []
+        errors: list[CallError] = []
+        for overload in overloads:
+            function, unwrap_error = self._unwrap_function(
+                overload, positional, keywords
+            )
+            if function is None:
+                errors.append(unwrap_error)  # type: ignore
+                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, None
+
+        # No match -> invalid call
+        if n_candidates == 0:
+            overloads_str: str = ", ".join(map(str, overloads))
+            errors_str: str = ", ".join(errors)
+            message: str = (
+                f"No matching overload in [{overloads_str}] {for_args} (errors: {errors_str})"
+            )
+            if report_errors:
+                self.reporter.error(location, message)
+            return None, message
+
+        # Multiple matches -> see if one <: all others (more specific)
+        for i1, c1 in enumerate(candidates):
+            mapped1: list[MappedArgument[E]] = c1.mapped
+            best_match: bool = True
+            for i2, c2 in enumerate(candidates):
+                if i1 == i2:
+                    continue
+                mapped2: list[MappedArgument[E]] = 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, None
+
+        candidates_str: str = ", ".join(
+            str(candidate.function) for candidate in candidates
+        )
+        message: str = f"Multiple matching overloads {for_args}: {candidates_str}"
+        if report_errors:
+            self.reporter.error(location, message)
+        return None, message
+
+    def map_call_arguments(
+        self,
+        function: Function,
+        location: Location,
+        positional: list[TypedExpr[E]],
+        keywords: dict[str, TypedExpr[E]],
+        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_params: set[str] = set()
+
+        required_positional: list[str] = [
+            param.name
+            for param in function.params.pos + function.params.mixed
+            if param.required
+        ]
+        required_keyword: list[str] = [
+            param.name for param in function.params.kw if param.required
+        ]
+
+        mapped: list[MappedArgument[E]] = []
+
+        pos_params: list[Function.Parameter] = list(function.params.pos)
+        mixed_params: list[Function.Parameter] = list(function.params.mixed)
+        kw_params: dict[str, Function.Parameter] = {
+            param.name: param for param in function.params.kw
+        }
+
+        valid_call: bool = True
+
+        # TODO: handle *args and **kwargs sinks
+        for arg in positional:
+            param: Function.Parameter
+            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_params.add(name)
+            mapped.append(
+                MappedArgument(
+                    arg_expr=arg[0],
+                    arg_type=arg[1],
+                    parameter=param,
+                )
+            )
+
+        kw_params.update({param.name: param for param in mixed_params})
+        for name, arg in keywords.items():
+            param: Function.Parameter
+            if name not in kw_params:
+                if report_errors:
+                    if name in set_params:
+                        self.reporter.error(
+                            arg[0].location, f"Multiple values for parameter '{name}'"
+                        )
+                    else:
+                        self.reporter.error(
+                            arg[0].location, f"Unknown keyword parameter '{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_params.add(name)
+            mapped.append(
+                MappedArgument(
+                    arg_expr=arg[0],
+                    arg_type=arg[1],
+                    parameter=param,
+                )
+            )
+
+        def join_params(params: list[str]) -> str:
+            params = list(map(lambda p: f"'{p}'", params))
+            if len(params) == 0:
+                return ""
+            if len(params) == 1:
+                return params[0]
+            return ", ".join(params[:-1]) + " and " + params[-1]
+
+        if len(required_positional) != 0:
+            plural: str = "" if len(required_positional) == 1 else "s"
+            params: str = join_params(required_positional)
+            if report_errors:
+                self.reporter.error(
+                    location,
+                    f"Missing required positional argument{plural}: {params}",
+                )
+            valid_call = False
+
+        if len(required_keyword) != 0:
+            plural: str = "" if len(required_keyword) == 1 else "s"
+            params: str = join_params(required_keyword)
+            if report_errors:
+                self.reporter.error(
+                    location,
+                    f"Missing required keyword argument{plural}: {params}",
+                )
+            valid_call = False
+
+        return valid_call, mapped
+
+    def _are_mapped_subtypes(
+        self, mapped1: list[MappedArgument[E]], mapped2: list[MappedArgument[E]]
+    ) -> 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[E, Type] = {}
+        for arg in mapped1:
+            by_expr[arg.arg_expr] = arg.parameter.type
+
+        for arg in mapped2:
+            type2: Type = arg.parameter.type
+            type1: Type = by_expr[arg.arg_expr]
+            if not self.types.is_subtype(type1, type2):
+                return False
+        return True

midas/checker/evaluator.py

@@ -0,0 +1,265 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Optional
+
+import midas.ast.midas as m
+from midas.ast.location import Location
+from midas.checker.preamble import Preamble
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter
+from midas.checker.types import Function, Predicate
+from midas.lexer.token import TokenType
+
+
+@dataclass(frozen=True, kw_only=True)
+class PartialPredicate(Predicate):
+    """A partially applied predicate"""
+
+    scope: dict[str, Any]
+    """A dictionary of already applied parameters"""
+
+
+class Evaluator(m.Expr.Visitor[Any]):
+    """Helper class to evaluate an expression
+
+    This class is used to evaluate constraint types on literals at compile-time.
+    """
+
+    def __init__(self, types: TypesRegistry, reporter: Optional[FileReporter] = None):
+        self.types: TypesRegistry = types
+        self.reporter: Optional[FileReporter] = reporter
+        self.preamble: Preamble = Preamble(self.types)
+        self.scopes: list[dict[str, Any]] = [{}]
+
+    def evaluate(self, expr: m.Expr) -> Any:
+        """Evaluate the given expression
+
+        Args:
+            expr (m.Expr): the expression to evaluate
+
+        Returns:
+            Any: the value of the expression
+        """
+        value: Any = expr.accept(self)
+        if self.reporter is not None:
+            self.reporter.debug(expr.location, f"Value: {value}")
+        return value
+
+    def get_value(self, name: str) -> Any:
+        """Get the value of a variable in the current scope
+
+        Args:
+            name (str): the name of the variable
+
+        Raises:
+            KeyError: if the variable is not defined
+
+        Returns:
+            Any: the value of the variable
+        """
+        scope: dict[str, Any] = self.scopes[-1]
+        return scope[name]
+
+    def set_value(self, name: str, value: Any, force_declare: bool = False):
+        """Set the value of a variable
+
+        If `force_declare` is `False`, this function first tries to find the
+        closest scope in which the variable is defined and assign the value in
+        that scope, if it can find one.
+
+        If `force_declare` is `True` or if the variable is not defined in any
+        scope, it is declare and assigned in the current scope
+
+        Args:
+            name (str): the name of the variable
+            value (Any): the value of the variable
+            force_declare (bool, optional): if `False` and the variable is
+                defined in a scope, the value is assigned in that scope (the
+                closest if there are multiple declarations). Defaults to False.
+        """
+        if not force_declare:
+            for scope in reversed(self.scopes):
+                if name in scope:
+                    scope[name] = value
+                    return
+        self.scopes[-1][name] = value
+
+    def visit_logical_expr(self, expr: m.LogicalExpr) -> Any:
+        def left():
+            return self.evaluate(expr.left)
+
+        def right():
+            return self.evaluate(expr.right)
+
+        match expr.operator.type:
+            case TokenType.AND:
+                return left() and right()
+            case _:
+                raise NotImplementedError
+
+    def visit_binary_expr(self, expr: m.BinaryExpr) -> Any:
+        left: Any = self.evaluate(expr.left)
+        right: Any = self.evaluate(expr.right)
+        match expr.operator.type:
+            case TokenType.PLUS:
+                return left + right
+            case TokenType.MINUS:
+                return left - right
+            case TokenType.STAR:
+                return left * right
+            case TokenType.SLASH:
+                return left / right
+            case TokenType.GREATER:
+                return left > right
+            case TokenType.GREATER_EQUAL:
+                return left >= right
+            case TokenType.LESS:
+                return left < right
+            case TokenType.LESS_EQUAL:
+                return left <= right
+            case TokenType.EQUAL_EQUAL:
+                return left == right
+            case TokenType.BANG_EQUAL:
+                return left != right
+            case _:
+                raise NotImplementedError
+
+    def visit_unary_expr(self, expr: m.UnaryExpr) -> Any:
+        right: Any = self.evaluate(expr.right)
+        match expr.operator.type:
+            case TokenType.PLUS:
+                return +right
+            case TokenType.MINUS:
+                return -right
+            case TokenType.BANG:
+                return not right
+            case _:
+                raise NotImplementedError
+
+    def visit_call_expr(self, expr: m.CallExpr) -> Any:
+        callee: Any = self.evaluate(expr.callee)
+        args: list[Any] = [self.evaluate(arg) for arg in expr.arguments]
+        kwargs: dict[str, Any] = {
+            name: self.evaluate(arg) for name, arg in expr.keywords.items()
+        }
+
+        match callee:
+            case Predicate():
+                return self._evaluate_predicate(expr.location, callee, args, kwargs)
+            case _ if callable(callee):
+                return callee(*args, **kwargs)
+            case _:
+                return NotImplementedError
+
+    def visit_get_expr(self, expr: m.GetExpr) -> Any:
+        obj: Any = self.evaluate(expr.expr)
+        return getattr(obj, expr.name.lexeme)
+
+    def visit_variable_expr(self, expr: m.VariableExpr) -> Any:
+        name: str = expr.name.lexeme
+        for scope in reversed(self.scopes):
+            if name in scope:
+                return scope[name]
+
+        predicate: Optional[Predicate] = self.types.lookup_predicate(name)
+        if predicate is not None:
+            if predicate.alias:
+                return self.evaluate(predicate.body)
+            return predicate
+
+        glob: Optional[Callable] = self.preamble.get_py_func(name)
+        if glob is not None:
+            return glob
+        raise NameError(f"Unknown variable '{name}'")
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr) -> Any:
+        return self.evaluate(expr.expr)
+
+    def visit_literal_expr(self, expr: m.LiteralExpr) -> Any:
+        return expr.value
+
+    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> Any:
+        return self.get_value("_")
+
+    def _evaluate_predicate(
+        self,
+        location: Location,
+        predicate: Predicate,
+        args: list[Any],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        """Evaluate a predicate function call
+
+        Args:
+            predicate (Predicate): the predicate to evaluate
+            args (list[Any]): a list of positional arguments
+            kwargs (dict[str, Any]): a map of keyword arguments
+
+        Returns:
+            Any: the value returned by the predicate call
+        """
+        res: Any = None
+        if isinstance(predicate, PartialPredicate):
+            self.scopes.append(predicate.scope)
+        else:
+            self.scopes.append({})
+        match predicate.type:
+            case Function(returns=Function() as inner):
+                self._map_args(location, predicate.type, args, kwargs)
+                res = PartialPredicate(
+                    type=inner,
+                    body=predicate.body,
+                    alias=False,
+                    scope=self.scopes[-1],
+                )
+
+            case Function():
+                self._map_args(location, predicate.type, args, kwargs)
+                res = self.evaluate(predicate.body)
+
+            case _:
+                raise NotImplementedError
+        self.scopes.pop()
+        return res
+
+    def _map_args(
+        self,
+        location: Location,
+        function: Function,
+        args: list[Any],
+        kwargs: dict[str, Any],
+    ):
+        """Map call arguments to a function's parameters and set their values in context
+
+        Each argument is mapped to a parameter of the function, then its value
+        is set in the context using :func:`set_value` with the parameter's name
+
+        Args:
+            function (Function): the called function
+            args (list[Any]): a list of positional arguments
+            kwargs (dict[str, Any]): a map of keyword arguments
+        """
+        positional: list[Function.Parameter] = (
+            function.params.pos + function.params.mixed
+        )
+        keywords: dict[str, Function.Parameter] = {
+            param.name: param for param in function.params.mixed + function.params.kw
+        }
+
+        for i, arg in enumerate(args):
+            if i >= len(positional):
+                if self.reporter is not None:
+                    self.reporter.error(
+                        location,
+                        f"Too many positional arguments, expected at most {len(positional)}, got {len(args)}",
+                    )
+                break
+            param: Function.Parameter = positional[i]
+            self.set_value(param.name, arg)
+
+        for name, arg in kwargs.items():
+            if name not in keywords:
+                if self.reporter is not None:
+                    self.reporter.error(location, f"Unknown keyword argument '{name}'")
+                break
+            param: Function.Parameter = keywords[name]
+            self.set_value(param.name, arg)

midas/checker/frames/column_groupby_methods.py

@@ -0,0 +1,240 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.dispatcher import CallResult
+from midas.checker.frames.utils import MethodRegistry, method
+from midas.checker.types import (
+    ColumnGroupBy,
+    ColumnType,
+    Function,
+    ParamSpec,
+    Type,
+    UnknownType,
+)
+
+if TYPE_CHECKING:
+    from midas.checker.python import TypedExpr
+
+
+@dataclass(frozen=True, kw_only=True)
+class Call:
+    """A column group-by method call, implements :class:`utils.MethodCall`"""
+
+    location: Location
+    call_expr: p.Expr
+    groupby: ColumnGroupBy
+    groupby_expr: p.Expr
+    positional: list[TypedExpr]
+    keywords: dict[str, TypedExpr]
+
+    @property
+    def subject(self) -> TypedExpr:
+        return (self.groupby_expr, self.groupby)
+
+
+class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
+    """The method registry for column group-by types"""
+
+    NAMED_ARGS: dict[str, str] = {
+        "numeric_only": "bool",
+        "skipna": "bool",
+        "engine": "str",
+        "engine_kwargs": "dict",
+    }
+
+    def _aggregate(
+        self,
+        call: Call,
+        method: str,
+        params: list[str | tuple[str, str, bool]] = [],
+    ) -> Type:
+        """Compute the result type of an aggregate method call
+
+        Args:
+            call (Call): the call object
+            method (str): the method name to delegate on :class:`Column`
+            params (list[str | tuple[str, str, bool], optional): a list of extra
+                mixed parameters. The list can contain strings to include
+                parameters predefined in `NAMED_ARGS`, or tuples containing the
+                parameter's name, type and required flag. Defaults to [].
+
+        Returns:
+            Type: the result type
+        """
+        real_params: list[Function.Parameter] = []
+        for i, param in enumerate(params):
+            match param:
+                case str() as name:
+                    param = Function.Parameter(
+                        pos=i,
+                        name=name,
+                        type=self.types.get_type(self.NAMED_ARGS[name]),
+                        required=False,
+                    )
+                case (name, type, required):
+                    param = Function.Parameter(
+                        pos=i,
+                        name=name,
+                        type=self.types.get_type(type),
+                        required=required,
+                    )
+            real_params.append(param)
+
+        # TODO: maybe better to filter arguments and pass some, in case the
+        # return type depends on them
+        returns: Type = self.typer.call_method(
+            location=call.location,
+            call_expr=call.call_expr,
+            obj=(call.groupby_expr, call.groupby.column),
+            method_name=method,
+            positional=[],
+            keywords={},
+        )
+        if not isinstance(returns, ColumnType):
+            returns = ColumnType(type=UnknownType())
+        signature = Function(
+            params=ParamSpec(mixed=real_params),
+            returns=returns,
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    @method()
+    def kurt(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            "kurt",
+            ["skipna", "numeric_only"],
+        )
+
+    @method()
+    def max(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            "max",
+            [
+                "numeric_only",
+                (
+                    "min_count",
+                    "int",
+                    False,
+                ),
+                "skipna",
+                "engine",
+                "engine_kwargs",
+            ],
+        )
+
+    @method()
+    def mean(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            "mean",
+            ["numeric_only", "skipna", "engine", "engine_kwargs"],
+        )
+
+    @method()
+    def median(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            "median",
+            ["numeric_only", "skipna"],
+        )
+
+    @method()
+    def min(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            "min",
+            [
+                "numeric_only",
+                (
+                    "min_count",
+                    "int",
+                    False,
+                ),
+                "skipna",
+                "engine",
+                "engine_kwargs",
+            ],
+        )
+
+    @method()
+    def prod(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            "prod",
+            [
+                "numeric_only",
+                (
+                    "min_count",
+                    "int",
+                    False,
+                ),
+                "skipna",
+            ],
+        )
+
+    @method()
+    def std(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            "std",
+            [
+                (
+                    "ddof",
+                    "int",
+                    False,
+                ),
+                "engine",
+                "engine_kwargs",
+                "numeric_only",
+                "skipna",
+            ],
+        )
+
+    @method()
+    def sum(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            "sum",
+            [
+                "numeric_only",
+                (
+                    "min_count",
+                    "int",
+                    False,
+                ),
+                "skipna",
+                "engine",
+                "engine_kwargs",
+            ],
+        )
+
+    @method()
+    def var(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            "var",
+            [
+                (
+                    "var",
+                    "int",
+                    False,
+                ),
+                "engine",
+                "engine_kwargs",
+                "numeric_only",
+                "skipna",
+            ],
+        )

midas/checker/frames/column_manager.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.dispatcher import CallResult
+from midas.checker.frames.column_groupby_methods import Call as GroupByCall
+from midas.checker.frames.column_groupby_methods import ColumnGroupByMethodRegistry
+from midas.checker.frames.column_methods import Call, ColumnMethodRegistry
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter
+from midas.checker.types import (
+    ColumnGroupBy,
+    ColumnType,
+    Function,
+    OverloadedFunction,
+    ParamSpec,
+    Type,
+)
+
+if TYPE_CHECKING:
+    from midas.checker.python import PythonTyper, TypedExpr
+
+
+class ColumnManager:
+    """Helper class to handle methods and subscripts on column types"""
+
+    def __init__(self, typer: PythonTyper) -> None:
+        self.typer: PythonTyper = typer
+        self.method_resolver: ColumnMethodRegistry = ColumnMethodRegistry(self.typer)
+        self.groupby_method_resolver: ColumnGroupByMethodRegistry = (
+            ColumnGroupByMethodRegistry(self.typer)
+        )
+
+    def get(
+        self,
+        reporter: FileReporter,
+        location: Location,
+        column: ColumnType,
+        index: TypedExpr,
+    ) -> Type:
+        """Compute the type of a subscript access
+
+        Args:
+            reporter (FileReporter): the file reporter to use for diagnostics
+            location (Location): the subscript's location
+            column (DataFrameType): the column type
+            index (TypedExpr): the index
+
+        Returns:
+            Type: the resulting type
+        """
+
+        single = Function(
+            params=ParamSpec(
+                pos=[
+                    Function.Parameter(
+                        pos=0,
+                        name="index",
+                        type=self.typer.types.get_type("int"),
+                        required=True,
+                    )
+                ]
+            ),
+            returns=column.type,
+        )
+        slice = Function(
+            params=ParamSpec(
+                pos=[
+                    Function.Parameter(
+                        pos=0,
+                        name="slice",
+                        type=self.typer.types.get_type("slice"),
+                        required=True,
+                    )
+                ]
+            ),
+            returns=column,
+        )
+
+        overload = OverloadedFunction(overloads=[single, slice])
+
+        result: CallResult = self.typer.dispatcher.get_result(
+            location=location,
+            callee=overload,
+            positional=[index],
+            keywords={},
+        )
+        return result.result
+
+    def call(
+        self,
+        method: str,
+        location: Location,
+        call_expr: p.Expr,
+        column: ColumnType,
+        column_expr: p.Expr,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+    ) -> Type:
+        """Compute the result type of a column's method call
+
+        Args:
+            method (str): the method name
+            location (Location): the call's location
+            call_expr (p.Expr): the call expression
+            column (ColumnType): the column type
+            column_expr (p.Expr): the column expression
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Returns:
+            Type: the result type
+        """
+        call: Call = Call(
+            location=location,
+            call_expr=call_expr,
+            column=column,
+            column_expr=column_expr,
+            positional=positional,
+            keywords=keywords,
+        )
+        return self.method_resolver.call(method, call)
+
+    def groupby_call(
+        self,
+        method: str,
+        location: Location,
+        call_expr: p.Expr,
+        groupby: ColumnGroupBy,
+        groupby_expr: p.Expr,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+    ) -> Type:
+        """Compute the result type of a column group-by's method call
+
+        Args:
+            method (str): the method name
+            location (Location): the call's location
+            call_expr (p.Expr): the call expression
+            groupby (ColumnGroupBy): the column group-by object
+            groupby_expr (p.Expr): the column group-by expression
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Returns:
+            Type: the result type
+        """
+        call: GroupByCall = GroupByCall(
+            location=location,
+            call_expr=call_expr,
+            groupby=groupby,
+            groupby_expr=groupby_expr,
+            positional=positional,
+            keywords=keywords,
+        )
+        return self.groupby_method_resolver.call(method, call)
+
+    def get_attribute(self, column: ColumnType, name: str) -> Optional[Type]:
+        """Get the type of a column's attribute
+
+        Args:
+            column (ColumnType): the column type
+            name (str): the attribute's name
+
+        Returns:
+            Optional[Type]: the attribute's type, or `None` if it doesn't exist
+        """
+        types: TypesRegistry = self.typer.types
+        match name:
+            case "ndim" | "size":
+                return types.get_type("int")
+
+            case "shape":
+                return types.tuple_of("int")
+
+            case "T":
+                return column
+
+            case _:
+                return None

midas/checker/frames/column_methods.py

@@ -0,0 +1,683 @@
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Optional, TypeAlias, Union
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.dispatcher import CallResult
+from midas.checker.frames.utils import MethodRegistry, method
+from midas.checker.types import (
+    ColumnGroupBy,
+    ColumnType,
+    Function,
+    OverloadedFunction,
+    ParamSpec,
+    TopType,
+    Type,
+    UnitType,
+    UnknownType,
+    unfold_type,
+)
+
+if TYPE_CHECKING:
+    from midas.checker.python import TypedExpr
+
+FormulaOperand: TypeAlias = Union["Formula", str, Type]
+"""
+A operand type in a :data:`Formula`
+
+Must be one of the following:
+- a nested formula
+- a type name (a string)
+- a type instance
+"""
+
+Formula: TypeAlias = Union[Type, tuple[FormulaOperand, str, FormulaOperand]]
+"""
+A formula to compute the output type of a function
+
+Must be either a type, or a tuple containing:
+- a left operand
+- an operation / method name (e.g. `"__add__"`)
+- a right operand
+
+For example, to compute the result of a `mean` function, given the input type `T`:
+```python
+mean_formula = ((T, "__add__", T), "__truediv__", "int")
+```
+"""
+
+
+@dataclass(frozen=True, kw_only=True)
+class Call:
+    """A column method call, implements :class:`utils.MethodCall`"""
+
+    location: Location
+    call_expr: p.Expr
+    column: ColumnType
+    column_expr: p.Expr
+    positional: list[TypedExpr]
+    keywords: dict[str, TypedExpr]
+
+    @property
+    def subject(self) -> TypedExpr:
+        return (self.column_expr, self.column)
+
+
+class ColumnMethodRegistry(MethodRegistry[Call]):
+    """The method registry for column types"""
+
+    def _resolve_formula_operand(self, call: Call, operand: FormulaOperand) -> Type:
+        """Resolve the type of a formula operand
+
+        See :data:`FormulaOperand` for more information on the accepted format
+
+        Args:
+            call (Call): the call that triggered this resolution
+            operand (FormulaOperand): the formula operand
+
+        Returns:
+            Type: the type of the operand
+        """
+        match operand:
+            case str():
+                return self.types.get_type(operand)
+            case (_, _, _):
+                return self._resolve_formula_type(call, operand)
+            case _:
+                return operand
+
+    def _resolve_formula_type(self, call: Call, formula: Formula) -> Type:
+        """Resolve the return type of a formula
+
+        See :data:`Formula` for more information on the accepted format
+
+        Args:
+            call (Call): the call that triggered this resolution
+            formula (Formula): the formula to evaluate
+
+        Returns:
+            Type: the return type of the formula
+        """
+        if not isinstance(formula, tuple):
+            return formula
+
+        op1, operator, op2 = formula
+        op1_type: Type = self._resolve_formula_operand(call, op1)
+        op2_type: Type = self._resolve_formula_operand(call, op2)
+        return self.typer.result_of_binary_op(
+            location=call.location,
+            expr=call.call_expr,
+            left=(call.column_expr, op1_type),
+            right=(call.column_expr, op2_type),
+            method=operator,
+        )
+
+    def _simple_call(self, call: Call, function: Type) -> Type:
+        """Get the result of calling a simple method
+
+        This function is a simple wrapper around :func:`dispatcher.CallDispatcher.get_result`
+        Args:
+            call (Call): the call that triggered this resolution
+            function (Type): the function type
+
+        Returns:
+            Type: the return type
+        """
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=function,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    def _element_binary_op(self, call: Call, method: str) -> tuple[Type, bool]:
+        """Compute the result of an element-wise binary operation
+
+        This function delegates to the inner types for computing the resulting
+        type.
+
+        Args:
+            call (Call): the call that triggered this resolution
+            method (str): the method name
+
+        Returns:
+            tuple[Type, bool]: the resulting type and a boolean indicating
+                whether the operand is a column
+        """
+        if len(call.positional) == 0:
+            return UnknownType(), False
+
+        col_type1: Type = call.column.type
+        operand: TypedExpr = call.positional[0]
+        unfolded_operand: Type = unfold_type(operand[1])
+        col_type2: Type
+
+        column_operand: bool = isinstance(unfolded_operand, ColumnType)
+
+        # Operand is a column -> get the inner type
+        if column_operand:
+            col_type2 = unfolded_operand.type
+        # Otherwise use the operand type itself
+        else:
+            col_type2 = operand[1]
+
+        new_inner_type = self.typer.result_of_binary_op(
+            location=call.location,
+            expr=call.call_expr,
+            left=(call.column_expr, col_type1),
+            right=(operand[0], col_type2),
+            method=method,
+        )
+        return ColumnType(type=new_inner_type), column_operand
+
+    def _element_wise(self, call: Call, method: str) -> Type:
+        """Compute the result of an element-wise method call
+
+        If the call is valid, this method also generates an assertion to check
+        that both operands have the same length at runtime
+
+        Args:
+            call (Call): the call object
+            method (str): the method's name
+
+        Returns:
+            Type: the result type
+        """
+
+        # Build signature with new column type and generic operand
+        returns, column_operand = self._element_binary_op(call, method)
+        signature = Function(
+            params=ParamSpec(
+                mixed=[
+                    Function.Parameter(
+                        pos=0,
+                        name="other",
+                        type=TopType(),
+                        required=True,
+                    ),
+                ],
+            ),
+            returns=returns,
+        )
+
+        # Map arguments and compute result type
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        if result.is_valid and column_operand:
+            self._assert_same_length(
+                call.call_expr, call.column_expr, call.positional[0][0]
+            )
+
+        return result.result
+
+    @method()
+    def copy(self, call: Call) -> Type:
+        return self._simple_call(
+            call,
+            Function(
+                params=ParamSpec(
+                    mixed=[
+                        Function.Parameter(
+                            pos=0,
+                            name="deep",
+                            type=self.types.get_type("bool"),
+                            required=False,
+                        )
+                    ]
+                ),
+                returns=call.column,
+            ),
+        )
+
+    @method()
+    def info(self, call: Call) -> Type:
+        def make_overload(memory_usage: Type, required: bool = False) -> Type:
+            return Function(
+                params=ParamSpec(
+                    mixed=[
+                        Function.Parameter(
+                            pos=0,
+                            name="verbose",
+                            type=self.types.get_type("bool"),
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=1,
+                            name="buf",
+                            type=TopType(),
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=2,
+                            name="max_cols",
+                            type=self.types.get_type("int"),
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=3,
+                            name="memory_usage",
+                            type=memory_usage,
+                            required=required,
+                        ),
+                        Function.Parameter(
+                            pos=4,
+                            name="show_counts",
+                            type=self.types.get_type("bool"),
+                            required=False,
+                        ),
+                    ]
+                ),
+                returns=UnitType(),
+            )
+
+        return self._simple_call(
+            call,
+            OverloadedFunction(
+                overloads=[
+                    make_overload(self.types.get_type("bool"), False),
+                    make_overload(self.types.get_type("str"), True),
+                ],
+            ),
+        )
+
+    @method("add", "__add__")
+    def add(self, call: Call) -> Type:
+        return self._element_wise(call, "__add__")
+
+    @method("sub", "__sub__")
+    def sub(self, call: Call) -> Type:
+        return self._element_wise(call, "__sub__")
+
+    @method("mul", "__mul__")
+    def mul(self, call: Call) -> Type:
+        return self._element_wise(call, "__mul__")
+
+    @method("div", "truediv", "__truediv__")
+    def truediv(self, call: Call) -> Type:
+        return self._element_wise(call, "__truediv__")
+
+    @method("floordiv", "__floordiv__")
+    def floordiv(self, call: Call) -> Type:
+        return self._element_wise(call, "__floordiv__")
+
+    @method("mod", "__mod__")
+    def mod(self, call: Call) -> Type:
+        return self._element_wise(call, "__mod__")
+
+    @method("pow", "__pow__")
+    def pow(self, call: Call) -> Type:
+        return self._element_wise(call, "__pow__")
+
+    @method("lt", "__lt__")
+    def lt(self, call: Call) -> Type:
+        return self._element_wise(call, "__lt__")
+
+    @method("gt", "__gt__")
+    def gt(self, call: Call) -> Type:
+        return self._element_wise(call, "__gt__")
+
+    @method("le", "__le__")
+    def le(self, call: Call) -> Type:
+        return self._element_wise(call, "__le__")
+
+    @method("ge", "__ge__")
+    def ge(self, call: Call) -> Type:
+        return self._element_wise(call, "__ge__")
+
+    @method("ne", "__ne__")
+    def ne(self, call: Call) -> Type:
+        return self._element_wise(call, "__ne__")
+
+    @method("eq", "__eq__")
+    def eq(self, call: Call) -> Type:
+        return self._element_wise(call, "__eq__")
+
+    def _aggregate(
+        self,
+        call: Call,
+        kwargs: list[Function.Parameter] = [],
+        *,
+        formula: Optional[Callable[[Type], Formula]] = None,
+    ) -> Type:
+        """Compute the result type of an aggregate method call
+
+        Args:
+            call (Call): the call object
+            kwargs (list[Function.Parameter], optional): a list of extra
+                keyword-only parameters. Defaults to [].
+            formula (Callable[[Type], Formula], optional): optional formula
+                builder function to compute the return type. If set, the function
+                should accept the inner column type and return a formula.
+                If `None`, the result is typed as `Column[Any]`. Defaults to None.
+
+        Returns:
+            Type: the result type
+        """
+
+        returns: Type = ColumnType(type=TopType())
+        if formula:
+            returns = ColumnType(
+                type=self._resolve_formula_type(
+                    call,
+                    formula(call.column.type),
+                )
+            )
+        signature = Function(
+            params=ParamSpec(
+                kw=[
+                    Function.Parameter(
+                        pos=0,
+                        name="axis",
+                        type=TopType(),
+                        required=False,
+                    ),
+                    *kwargs,
+                ],
+            ),
+            returns=returns,
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    @method("kurtosis", "kurt")
+    def kurtosis(self, call: Call) -> Type:
+        return self._aggregate(call)
+
+    @method()
+    def max(self, call: Call) -> Type:
+        return self._aggregate(call, formula=lambda t: t)
+
+    @method()
+    def mean(self, call: Call) -> Type:
+        return self._aggregate(
+            call, formula=lambda t: ((t, "__add__", t), "__truediv__", "int")
+        )
+
+    @method()
+    def median(self, call: Call) -> Type:
+        return self._aggregate(call, formula=lambda t: t)
+
+    @method()
+    def min(self, call: Call) -> Type:
+        return self._aggregate(call, formula=lambda t: t)
+
+    @method()
+    def mode(self, call: Call) -> Type:
+        return self._aggregate(call, formula=lambda t: t)
+
+    @method("product", "prod")
+    def product(self, call: Call) -> Type:
+        return self._aggregate(call, formula=lambda t: (t, "__mul__", t))
+
+    @method()
+    def std(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            [
+                Function.Parameter(
+                    pos=1,
+                    name="ddof",
+                    type=self.types.get_type("int"),
+                    required=False,
+                )
+            ],
+        )
+
+    @method()
+    def sum(self, call: Call) -> Type:
+        return self._aggregate(call, formula=lambda t: (t, "__add__", t))
+
+    @method()
+    def var(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            [
+                Function.Parameter(
+                    pos=1,
+                    name="var",
+                    type=self.types.get_type("int"),
+                    required=False,
+                )
+            ],
+        )
+
+    @method()
+    def head(self, call: Call) -> Type:
+        signature = Function(
+            params=ParamSpec(
+                mixed=[
+                    Function.Parameter(
+                        pos=0,
+                        name="n",
+                        type=self.types.get_type("int"),
+                        required=False,
+                    ),
+                ],
+            ),
+            returns=call.column,
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    @method()
+    def tail(self, call: Call) -> Type:
+        signature = Function(
+            params=ParamSpec(
+                mixed=[
+                    Function.Parameter(
+                        pos=0,
+                        name="n",
+                        type=self.types.get_type("int"),
+                        required=False,
+                    ),
+                ],
+            ),
+            returns=call.column,
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    @method()
+    def sort_values(self, call: Call) -> Type:
+        str_ = self.types.get_type("str")
+        bool_ = self.types.get_type("bool")
+
+        def make_overload(ascending: Type) -> Function:
+            return Function(
+                params=ParamSpec(
+                    kw=[
+                        Function.Parameter(
+                            pos=0,
+                            name="axis",
+                            type=TopType(),
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=1,
+                            name="ascending",
+                            type=ascending,
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=2,
+                            name="inplace",
+                            type=bool_,
+                            required=False,
+                            unsupported=True,
+                        ),
+                        Function.Parameter(
+                            pos=3,
+                            name="kind",
+                            type=str_,
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=4,
+                            name="na_position",
+                            type=str_,
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=5,
+                            name="ignore_index",
+                            type=bool_,
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=6,
+                            name="key",
+                            type=TopType(),
+                            required=False,
+                        ),
+                    ],
+                ),
+                returns=call.column,
+            )
+
+        list_of = self.types.list_of
+        overloads: list[Type] = [
+            make_overload(bool_),
+            make_overload(bool_),
+            make_overload(list_of(bool_)),
+            make_overload(list_of(bool_)),
+        ]
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=OverloadedFunction(overloads=overloads),
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    @method()
+    def groupby(self, call: Call) -> Type:
+        bool_: Type = self.types.get_type("bool")
+        function: Function = Function(
+            params=ParamSpec(
+                mixed=[
+                    Function.Parameter(
+                        pos=0,
+                        name="by",
+                        type=TopType(),
+                        required=False,
+                    ),
+                    Function.Parameter(
+                        pos=1,
+                        name="level",
+                        type=TopType(),
+                        required=False,
+                    ),
+                ],
+                kw=[
+                    Function.Parameter(
+                        pos=i + 2,
+                        name=name,
+                        type=bool_,
+                        required=False,
+                    )
+                    for i, name in enumerate(
+                        ["as_index", "sort", "group_keys", "observed", "dropna"]
+                    )
+                ],
+            ),
+            returns=ColumnGroupBy(column=call.column),
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=function,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    def _assert_same_length(self, call_expr: p.Expr, column1: p.Expr, column2: p.Expr):
+        """Generate an assertion to check that two columns have the same length
+
+        Args:
+            call_expr (p.Expr): the call expression, to insert the assertion
+                at the right place
+            column1 (p.Expr): the first column expression
+            column2 (p.Expr): the second column expression
+        """
+        func_name: str = "__midas_column_same_length__"
+
+        # Efficiently compute length
+        # https://stackoverflow.com/a/15943975/11109181
+        def len_of_col(col: ast.expr) -> ast.expr:
+            return ast.Call(
+                func=ast.Name(id="len"),
+                args=[
+                    ast.Attribute(
+                        value=col,
+                        attr="index",
+                    )
+                ],
+                keywords=[],
+            )
+
+        self.assertions.define(
+            func_name,
+            ast.FunctionDef(
+                name=func_name,
+                args=ast.arguments(
+                    posonlyargs=[],
+                    args=[
+                        ast.arg(arg="column1"),
+                        ast.arg(arg="column2"),
+                    ],
+                    kwonlyargs=[],
+                    defaults=[],
+                    kw_defaults=[],
+                ),
+                body=[
+                    ast.Return(
+                        value=ast.Compare(
+                            left=len_of_col(ast.Name(id="column1")),
+                            ops=[ast.Eq()],
+                            comparators=[
+                                len_of_col(ast.Name(id="column2")),
+                            ],
+                        )
+                    )
+                ],
+                decorator_list=[],
+            ),
+        )
+        self.assertions.add(
+            bound_expr=call_expr,
+            inputs=[column1, column2],
+            builder=lambda c1, c2: ast.Call(
+                func=ast.Name(id=func_name),
+                args=[c1, c2],
+                keywords=[],
+            ),
+            message="Columns must have the same length",
+        )

midas/checker/frames/frame_groupby_methods.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.frames.utils import MethodRegistry, method
+from midas.checker.types import (
+    ColumnGroupBy,
+    ColumnType,
+    DataFrameType,
+    FrameGroupBy,
+    Type,
+    UnknownType,
+)
+
+if TYPE_CHECKING:
+    from midas.checker.python import TypedExpr
+
+
+@dataclass(frozen=True, kw_only=True)
+class Call:
+    """A frame group-by method call, implements :class:`utils.MethodCall`"""
+
+    location: Location
+    call_expr: p.Expr
+    groupby: FrameGroupBy
+    groupby_expr: p.Expr
+    positional: list[TypedExpr]
+    keywords: dict[str, TypedExpr]
+
+    @property
+    def subject(self) -> TypedExpr:
+        return (self.groupby_expr, self.groupby)
+
+
+class FrameGroupByMethodRegistry(MethodRegistry[Call]):
+    """The method registry for frame group-by types"""
+
+    def _aggregate(self, call: Call, method: str) -> Type:
+        """Compute the result type of an aggregate method call
+
+        Args:
+            call (Call): the call object
+            method (str): the method's name
+
+        Returns:
+            Type: the result type
+        """
+        new_columns: list[DataFrameType.Column] = []
+
+        for column in call.groupby.frame.columns:
+            with self.reporter.with_context(f"in column '{column.name}'"):
+                column_groupby: ColumnGroupBy = ColumnGroupBy(column=column.type)
+                result_type: Type = self.typer.call_method(
+                    location=call.location,
+                    call_expr=call.call_expr,
+                    obj=(call.groupby_expr, column_groupby),
+                    method_name=method,
+                    positional=call.positional,
+                    keywords=call.keywords,
+                )
+                if not isinstance(result_type, ColumnType):
+                    result_type = ColumnType(type=UnknownType())
+                new_columns.append(
+                    DataFrameType.Column(
+                        index=column.index,
+                        name=column.name,
+                        type=result_type,
+                    )
+                )
+
+        return DataFrameType(columns=new_columns)
+
+    @method()
+    def kurt(self, call: Call) -> Type:
+        return self._aggregate(call, "kurt")
+
+    @method()
+    def max(self, call: Call) -> Type:
+        return self._aggregate(call, "max")
+
+    @method()
+    def mean(self, call: Call) -> Type:
+        return self._aggregate(call, "mean")
+
+    @method()
+    def median(self, call: Call) -> Type:
+        return self._aggregate(call, "median")
+
+    @method()
+    def min(self, call: Call) -> Type:
+        return self._aggregate(call, "min")
+
+    @method()
+    def prod(self, call: Call) -> Type:
+        return self._aggregate(call, "prod")
+
+    @method()
+    def std(self, call: Call) -> Type:
+        return self._aggregate(call, "std")
+
+    @method()
+    def sum(self, call: Call) -> Type:
+        return self._aggregate(call, "sum")
+
+    @method()
+    def var(self, call: Call) -> Type:
+        return self._aggregate(call, "var")

midas/checker/frames/frame_manager.py

@@ -0,0 +1,386 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, TypeGuard, cast
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.frames.frame_groupby_methods import Call as GroupByCall
+from midas.checker.frames.frame_groupby_methods import FrameGroupByMethodRegistry
+from midas.checker.frames.frame_methods import Call, FrameMethodRegistry
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter
+from midas.checker.types import (
+    ColumnGroupBy,
+    ColumnType,
+    DataFrameType,
+    FrameGroupBy,
+    TupleType,
+    Type,
+    UnknownType,
+)
+
+if TYPE_CHECKING:
+    from midas.checker.python import PythonTyper, TypedExpr
+
+
+def is_list_of_literals(exprs: list[p.Expr]) -> TypeGuard[list[p.LiteralExpr]]:
+    """Check whether the given list only contains literal expressions
+
+    Args:
+        exprs (list[p.Expr]): the list to check
+
+    Returns:
+        TypeGuard[list[p.LiteralExpr]]: whether `exprs` only contains literal expressions
+    """
+    return all(isinstance(expr, p.LiteralExpr) for expr in exprs)
+
+
+class FrameManager:
+    """Helper class to handle methods and subscripts on frame types"""
+
+    def __init__(self, typer: PythonTyper) -> None:
+        self.typer: PythonTyper = typer
+        self.method_resolver: FrameMethodRegistry = FrameMethodRegistry(self.typer)
+        self.groupby_method_resolver: FrameGroupByMethodRegistry = (
+            FrameGroupByMethodRegistry(self.typer)
+        )
+
+    def assign(
+        self,
+        reporter: FileReporter,
+        location: Location,
+        frame: DataFrameType,
+        index: p.Expr,
+        value_type: Type,
+    ) -> Type:
+        """Compute the new frame type after assigning a value to an index
+
+        Args:
+            reporter (FileReporter): the file reporter to use for diagnostics
+            location (Location): the assignment's location
+            frame (DataFrameType): the frame type
+            index (p.Expr): the index expression
+            value_type (Type): the assigned value
+
+        Returns:
+            Type: the resulting frame type
+        """
+        match index:
+            case p.LiteralExpr(value=str() as name):
+                return self.assign_column(reporter, location, frame, name, value_type)
+
+            case p.ListExpr(items=indices) if is_list_of_literals(indices) and all(
+                isinstance(index.value, str) for index in indices
+            ):
+                names: list[str] = [cast(str, index.value) for index in indices]
+
+                if not isinstance(value_type, TupleType):
+                    reporter.error(
+                        location,
+                        f"Cannot assign {type} to dataframe columns. Must be a tuple of columns",
+                    )
+                    return UnknownType()
+
+                if len(names) != len(value_type.items):
+                    reporter.error(
+                        location,
+                        f"Wrong number of columns. Cannot assign {len(value_type.items)} to {len(names)} targets",
+                    )
+                    return UnknownType()
+
+                new_frame: Type = frame
+                for name, value in zip(names, value_type.items):
+                    new_frame = self.assign_column(
+                        reporter,
+                        location,
+                        new_frame,
+                        name,
+                        value,
+                    )
+                    if not isinstance(new_frame, DataFrameType):
+                        return new_frame
+                return new_frame
+
+            case _:
+                reporter.error(
+                    location, f"Invalid index type {index} on {frame} (assignment)"
+                )
+                return UnknownType()
+
+    def assign_column(
+        self,
+        reporter: FileReporter,
+        location: Location,
+        frame: DataFrameType,
+        name: str,
+        type: Type,
+    ) -> Type:
+        """Compute the new frame type after assigning a single value to a column
+
+        Args:
+            reporter (FileReporter): the file reporter to use for diagnostics
+            location (Location): the assignment's location
+            frame (DataFrameType): the frame type
+            name (str): the column name
+            type (Type): the assigned value type
+
+        Returns:
+            Type: the resulting frame type
+        """
+        if not isinstance(type, ColumnType):
+            reporter.error(
+                location,
+                f"Cannot assign {type} to dataframe column. Must be a ColumnType",
+            )
+            return self._set_column(frame, name, ColumnType(type=UnknownType()))
+        return self._set_column(frame, name, type)
+
+    def get(
+        self,
+        reporter: FileReporter,
+        location: Location,
+        frame: DataFrameType,
+        index: p.Expr,
+    ) -> Type:
+        """Compute the type of a subscript access
+
+        Args:
+            reporter (FileReporter): the file reporter to use for diagnostics
+            location (Location): the subscript's location
+            frame (DataFrameType): the frame type
+            index (p.Expr): the index expression
+
+        Returns:
+            Type: the resulting type
+        """
+        match index:
+            case p.LiteralExpr(value=str() as name):
+                column: Optional[ColumnType] = FrameManager._get_column(frame, name)
+                if column is None:
+                    reporter.error(location, f"Unknown column '{name}' on {frame}")
+                    return UnknownType()
+                return column
+
+            case p.ListExpr(items=indices) if is_list_of_literals(indices) and all(
+                isinstance(index.value, str) for index in indices
+            ):
+                names: list[str] = [cast(str, index.value) for index in indices]
+                columns: list[ColumnType] = []
+                for name in names:
+                    column: Optional[ColumnType] = FrameManager._get_column(frame, name)
+                    if column is None:
+                        reporter.error(location, f"Unknown column '{name}' on {frame}")
+                        return UnknownType()
+                    columns.append(column)
+                return TupleType(items=tuple(columns))
+
+            case _:
+                reporter.error(
+                    location, f"Invalid index type {index} on {frame} (access)"
+                )
+                return UnknownType()
+
+    def groupby_get(
+        self,
+        reporter: FileReporter,
+        location: Location,
+        groupby: FrameGroupBy,
+        index: p.Expr,
+    ) -> Type:
+        """Compute the type of a subscript access on a frame group-by object
+
+        Args:
+            reporter (FileReporter): the file reporter to use for diagnostics
+            location (Location): the subscript's location
+            groupby (FrameGroupBy): the group-by object
+            index (p.Expr): the index expression
+
+        Returns:
+            Type: the resulting type
+        """
+        result: Type = self.get(reporter, location, groupby.frame, index)
+        match result:
+            case ColumnType():
+                result = ColumnGroupBy(column=result)
+            case TupleType(items=columns):
+                result = TupleType(
+                    items=tuple(
+                        ColumnGroupBy(column=cast(ColumnType, column))
+                        for column in columns
+                    )
+                )
+        return result
+
+    @classmethod
+    def _set_column(
+        cls, frame: DataFrameType, name: str, column: ColumnType
+    ) -> DataFrameType:
+        """Set a frame's column to the given type
+
+        Args:
+            frame (DataFrameType): the frame type
+            name (str): the column's name
+            column (ColumnType): the new column's type
+
+        Returns:
+            DataFrameType: the new frame type
+        """
+        new_columns: list[DataFrameType.Column] = []
+        index: int = len(frame.columns)
+        replace: bool = False
+        for i, col in enumerate(frame.columns):
+            if col.name == name:
+                index = i
+                replace = True
+                # TODO: check column type here to prevent changing it
+            new_columns.append(col)
+
+        new_col: DataFrameType.Column = DataFrameType.Column(
+            index=index,
+            name=name,
+            type=column,
+        )
+        if replace:
+            new_columns[index] = new_col
+        else:
+            new_columns.append(new_col)
+
+        return DataFrameType(columns=new_columns)
+
+    @classmethod
+    def _set_columns(
+        cls, frame: DataFrameType, names: list[str], columns: list[ColumnType]
+    ) -> DataFrameType:
+        """Set multiple columns of a frame to the given types
+
+        Args:
+            frame (DataFrameType): the frame type
+            names (list[str]): the column names
+            columns (list[ColumnType]): the new column types
+
+        Returns:
+            DataFrameType: the new frame type
+        """
+        for name, col in zip(names, columns):
+            frame = cls._set_column(frame, name, col)
+        return frame
+
+    @classmethod
+    def _get_column(cls, frame: DataFrameType, name: str) -> Optional[ColumnType]:
+        """Get a column's type by name
+
+        Args:
+            frame (DataFrameType): the frame type
+            name (str): the column's name
+
+        Returns:
+            Optional[ColumnType]: the column's type, or `None` if it doesn't exist
+        """
+        for col in frame.columns:
+            if col.name == name:
+                return col.type
+        return None
+
+    @classmethod
+    def _get_columns(
+        cls, frame: DataFrameType, names: list[str]
+    ) -> list[Optional[ColumnType]]:
+        """Get multiple column types by name
+
+        Args:
+            frame (DataFrameType): the frame type
+            names (list[str]): the column names
+
+        Returns:
+            list[Optional[ColumnType]]: the column types (see :func:`_get_column`)
+        """
+        return [cls._get_column(frame, name) for name in names]
+
+    def call(
+        self,
+        method: str,
+        location: Location,
+        call_expr: p.Expr,
+        frame: DataFrameType,
+        frame_expr: p.Expr,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+    ) -> Type:
+        """Compute the result type of a frame's method call
+
+        Args:
+            method (str): the method name
+            location (Location): the call's location
+            call_expr (p.Expr): the call expression
+            frame (DataFrameType): the frame type
+            frame_expr (p.Expr): the frame expression
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Returns:
+            Type: the result type
+        """
+        call: Call = Call(
+            location=location,
+            call_expr=call_expr,
+            frame=frame,
+            frame_expr=frame_expr,
+            positional=positional,
+            keywords=keywords,
+        )
+        return self.method_resolver.call(method, call)
+
+    def groupby_call(
+        self,
+        method: str,
+        location: Location,
+        call_expr: p.Expr,
+        groupby: FrameGroupBy,
+        groupby_expr: p.Expr,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+    ) -> Type:
+        """Compute the result type of a frame group-by's method call
+
+        Args:
+            method (str): the method name
+            location (Location): the call's location
+            call_expr (p.Expr): the call expression
+            groupby (FrameGroupBy): the frame group-by object
+            groupby_expr (p.Expr): the frame group-by expression
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Returns:
+            Type: the result type
+        """
+        call: GroupByCall = GroupByCall(
+            location=location,
+            call_expr=call_expr,
+            groupby=groupby,
+            groupby_expr=groupby_expr,
+            positional=positional,
+            keywords=keywords,
+        )
+        return self.groupby_method_resolver.call(method, call)
+
+    def get_attribute(self, frame: DataFrameType, name: str) -> Optional[Type]:
+        """Get the type of a frame's attribute
+
+        Args:
+            frame (DataFrameType): the frame type
+            name (str): the attribute's name
+
+        Returns:
+            Optional[Type]: the attribute's type, or `None` if it doesn't exist
+        """
+        types: TypesRegistry = self.typer.types
+        match name:
+            case "ndim" | "size":
+                return types.get_type("int")
+
+            case "shape":
+                return types.tuple_of("int", "int")
+
+            case _:
+                return None

midas/checker/frames/frame_methods.py

@@ -0,0 +1,757 @@
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.dispatcher import CallResult
+from midas.checker.frames.utils import MethodRegistry, method
+from midas.checker.types import (
+    ColumnType,
+    DataFrameType,
+    FrameGroupBy,
+    Function,
+    OverloadedFunction,
+    ParamSpec,
+    TopType,
+    Type,
+    UnitType,
+    UnknownType,
+    unfold_type,
+)
+
+if TYPE_CHECKING:
+    from midas.checker.python import TypedExpr
+
+
+@dataclass(frozen=True, kw_only=True)
+class Call:
+    """A frame method call, implements :class:`utils.MethodCall`"""
+
+    location: Location
+    call_expr: p.Expr
+    frame: DataFrameType
+    frame_expr: p.Expr
+    positional: list[TypedExpr]
+    keywords: dict[str, TypedExpr]
+
+    @property
+    def subject(self) -> TypedExpr:
+        return (self.frame_expr, self.frame)
+
+
+class FrameMethodRegistry(MethodRegistry[Call]):
+    """The method registry for frame types"""
+
+    def _simple_call(self, call: Call, function: Type) -> Type:
+        """Get the result of calling a simple method
+
+        This function is a simple wrapper around :func:`dispatcher.CallDispatcher.get_result`
+        Args:
+            call (Call): the call that triggered this resolution
+            function (Type): the function type
+
+        Returns:
+            Type: the return type
+        """
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=function,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    def _get_method_result(
+        self,
+        call: Call,
+        column1: ColumnType,
+        column2: ColumnType,
+        method: str,
+    ) -> ColumnType:
+        """Get the result of calling a method on a column, passing a second
+
+        This function delegates to the main typer the resolution of the method
+        member, as well as computing the result type. Because we don't have any
+        AST expression for the individual columns, the frame expressions are
+        used instead.
+
+        Args:
+            call (Call): the call that triggered this resolution
+            column1 (ColumnType): the first column, i.e. left operand
+            column2 (ColumnType): the second column, i.e. right operand
+            method (str): the method name
+
+        Returns:
+            ColumnType: the resulting column.
+                If the operation is invalid / doesn't exist,
+                `ColumnType(type=UnknownType())` is returned
+        """
+
+        result: Type = self.typer.result_of_binary_op(
+            location=call.location,
+            expr=call.call_expr,
+            left=(call.frame_expr, column1),
+            right=(call.positional[0][0], column2),
+            method=method,
+        )
+
+        if not isinstance(result, ColumnType):
+            return ColumnType(type=UnknownType())
+        return result
+
+    def _element_binary_op(self, call: Call, method: str) -> tuple[Type, bool]:
+        """Compute the result of an element-wise binary operation
+
+        This function delegates to the matching columns for computing resulting
+        types. Any column only present in one of the frames is forwarded as a
+        generic `ColumnType(type=UnknownType())`. Columns only in the second
+        frame are append at the end of the schema.
+
+        Args:
+            call (Call): the call that triggered this resolution
+            method (str): the method name
+
+        Returns:
+            tuple[Type, bool]: the resulting type and a boolean indicating
+                whether the operand is a frame
+        """
+
+        if len(call.positional) == 0:
+            return UnknownType(), False
+
+        operand: TypedExpr = call.positional[0]
+        new_columns: list[DataFrameType.Column] = []
+
+        by_name: dict[str, DataFrameType.Column] = {}
+        frame2: Optional[DataFrameType] = None
+        # Get map of operand's columns by name, if the operand is a dataframe
+        unfolded_other: Type = unfold_type(operand[1])
+        frame_operand: bool = isinstance(unfolded_other, DataFrameType)
+        if frame_operand:
+            frame2 = unfolded_other
+            by_name = {col.name: col for col in frame2.columns if col.name is not None}
+
+        # Compute new schema:
+        # Step 1: for all columns in frame1:
+        #  - if present in frame2 -> delegate operation to columns
+        #  - if not -> add to schema as unknown
+        in_frame1: set[str] = set()
+        for column in call.frame.columns:
+            if column.name is not None:
+                in_frame1.add(column.name)
+
+            col_type1: ColumnType = column.type
+            col_type: ColumnType = ColumnType(type=UnknownType())
+
+            col_type2: Optional[ColumnType] = None
+
+            # Operand is a frame -> lookup column with the same name
+            if frame2 is not None:
+                if column.name in by_name:
+                    column2 = by_name[column.name]
+                    col_type2 = column2.type
+
+            # Operand is not a frame -> scalar operation -> ad-hoc column
+            else:
+                col_type2 = ColumnType(type=operand[1])
+
+            if col_type2 is not None:
+                with self.reporter.with_context(f"in column '{column.name}'"):
+                    col_type = self._get_method_result(
+                        call, col_type1, col_type2, method
+                    )
+
+            new_column = DataFrameType.Column(
+                index=column.index,
+                name=column.name,
+                type=col_type,
+            )
+            new_columns.append(new_column)
+
+        # Step 2: for all columns in frame2
+        #  - if not in frame1 -> add to schema as unknown
+        if frame2 is not None:
+            for column in frame2.columns:
+                if column.name in in_frame1:
+                    continue
+                new_columns.append(
+                    DataFrameType.Column(
+                        index=len(new_columns),
+                        name=column.name,
+                        type=ColumnType(type=UnknownType()),
+                    )
+                )
+
+        return DataFrameType(columns=new_columns), frame_operand
+
+    def _element_wise(self, call: Call, method: str) -> Type:
+        """Compute the result of an element-wise method call
+
+        If the call is valid, this method also generates an assertion to check
+        that both operands have the same length at runtime
+
+        Args:
+            call (Call): the call object
+            method (str): the method's name
+
+        Returns:
+            Type: the result type
+        """
+        # TODO: support sequence, Series, dict operand
+        returns, frame_operand = self._element_binary_op(call, method)
+
+        # Build signature with new schema and generic operand
+        signature = Function(
+            params=ParamSpec(
+                mixed=[
+                    Function.Parameter(
+                        pos=0,
+                        name="other",
+                        type=TopType(),
+                        required=True,
+                    ),
+                ],
+            ),
+            returns=returns,
+        )
+
+        # Map arguments and compute result type
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        if result.is_valid and frame_operand:
+            self._assert_same_length(
+                call.call_expr, call.frame_expr, call.positional[0][0]
+            )
+
+        return result.result
+
+    @method()
+    def copy(self, call: Call) -> Type:
+        return self._simple_call(
+            call,
+            Function(
+                params=ParamSpec(
+                    mixed=[
+                        Function.Parameter(
+                            pos=0,
+                            name="deep",
+                            type=self.types.get_type("bool"),
+                            required=False,
+                        )
+                    ]
+                ),
+                returns=call.frame,
+            ),
+        )
+
+    @method()
+    def info(self, call: Call) -> Type:
+        def make_overload(memory_usage: Type, required: bool = False) -> Type:
+            return Function(
+                params=ParamSpec(
+                    mixed=[
+                        Function.Parameter(
+                            pos=0,
+                            name="verbose",
+                            type=self.types.get_type("bool"),
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=1,
+                            name="buf",
+                            type=TopType(),
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=2,
+                            name="max_cols",
+                            type=self.types.get_type("int"),
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=3,
+                            name="memory_usage",
+                            type=memory_usage,
+                            required=required,
+                        ),
+                        Function.Parameter(
+                            pos=4,
+                            name="show_counts",
+                            type=self.types.get_type("bool"),
+                            required=False,
+                        ),
+                    ]
+                ),
+                returns=UnitType(),
+            )
+
+        return self._simple_call(
+            call,
+            OverloadedFunction(
+                overloads=[
+                    make_overload(self.types.get_type("bool"), False),
+                    make_overload(self.types.get_type("str"), True),
+                ],
+            ),
+        )
+
+    @method("add", "__add__")
+    def add(self, call: Call) -> Type:
+        return self._element_wise(call, "__add__")
+
+    @method("sub", "__sub__")
+    def sub(self, call: Call) -> Type:
+        return self._element_wise(call, "__sub__")
+
+    @method("mul", "__mul__")
+    def mul(self, call: Call) -> Type:
+        return self._element_wise(call, "__mul__")
+
+    @method("div", "truediv", "__truediv__")
+    def truediv(self, call: Call) -> Type:
+        return self._element_wise(call, "__truediv__")
+
+    @method("floordiv", "__floordiv__")
+    def floordiv(self, call: Call) -> Type:
+        return self._element_wise(call, "__floordiv__")
+
+    @method("mod", "__mod__")
+    def mod(self, call: Call) -> Type:
+        return self._element_wise(call, "__mod__")
+
+    @method("pow", "__pow__")
+    def pow(self, call: Call) -> Type:
+        return self._element_wise(call, "__pow__")
+
+    @method("lt", "__lt__")
+    def lt(self, call: Call) -> Type:
+        return self._element_wise(call, "__lt__")
+
+    @method("gt", "__gt__")
+    def gt(self, call: Call) -> Type:
+        return self._element_wise(call, "__gt__")
+
+    @method("le", "__le__")
+    def le(self, call: Call) -> Type:
+        return self._element_wise(call, "__le__")
+
+    @method("ge", "__ge__")
+    def ge(self, call: Call) -> Type:
+        return self._element_wise(call, "__ge__")
+
+    @method("ne", "__ne__")
+    def ne(self, call: Call) -> Type:
+        return self._element_wise(call, "__ne__")
+
+    @method("eq", "__eq__")
+    def eq(self, call: Call) -> Type:
+        return self._element_wise(call, "__eq__")
+
+    def _aggregate(self, call: Call, kwargs: list[Function.Parameter] = []) -> Type:
+        """Compute the result type of an aggregate method call
+
+        Args:
+            call (Call): the call object
+            kwargs (list[Function.Parameter], optional): a list of extra
+                keyword-only parameters. Defaults to [].
+
+        Returns:
+            Type: the result type
+        """
+        with_axis = Function(
+            params=ParamSpec(
+                kw=[
+                    Function.Parameter(
+                        pos=0,
+                        name="axis",
+                        type=self.types.get_type("int"),
+                        required=False,
+                    ),
+                    *kwargs,
+                ],
+            ),
+            returns=ColumnType(type=TopType()),
+        )
+        without_axis = Function(
+            params=ParamSpec(
+                kw=[
+                    Function.Parameter(
+                        pos=0,
+                        name="axis",
+                        type=self.types.get_type("None"),
+                        required=True,
+                    ),
+                    *kwargs,
+                ],
+            ),
+            returns=TopType(),
+        )
+        overload = OverloadedFunction(
+            overloads=[
+                with_axis,
+                without_axis,
+            ]
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=overload,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    @method("kurtosis", "kurt")
+    def kurtosis(self, call: Call) -> Type:
+        return self._aggregate(call)
+
+    @method()
+    def max(self, call: Call) -> Type:
+        return self._aggregate(call)
+
+    @method()
+    def mean(self, call: Call) -> Type:
+        return self._aggregate(call)
+
+    @method()
+    def median(self, call: Call) -> Type:
+        return self._aggregate(call)
+
+    @method()
+    def min(self, call: Call) -> Type:
+        return self._aggregate(call)
+
+    @method()
+    def mode(self, call: Call) -> Type:
+        return self._aggregate(call)
+
+    @method("product", "prod")
+    def product(self, call: Call) -> Type:
+        return self._aggregate(call)
+
+    @method()
+    def std(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            [
+                Function.Parameter(
+                    pos=1,
+                    name="ddof",
+                    type=self.types.get_type("int"),
+                    required=False,
+                )
+            ],
+        )
+
+    @method()
+    def sum(self, call: Call) -> Type:
+        return self._aggregate(call)
+
+    @method()
+    def var(self, call: Call) -> Type:
+        return self._aggregate(
+            call,
+            [
+                Function.Parameter(
+                    pos=1,
+                    name="var",
+                    type=self.types.get_type("int"),
+                    required=False,
+                )
+            ],
+        )
+
+    @method()
+    def head(self, call: Call) -> Type:
+        signature = Function(
+            params=ParamSpec(
+                mixed=[
+                    Function.Parameter(
+                        pos=0,
+                        name="n",
+                        type=self.types.get_type("int"),
+                        required=False,
+                    ),
+                ],
+            ),
+            returns=call.frame,
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    @method()
+    def tail(self, call: Call) -> Type:
+        signature = Function(
+            params=ParamSpec(
+                mixed=[
+                    Function.Parameter(
+                        pos=0,
+                        name="n",
+                        type=self.types.get_type("int"),
+                        required=False,
+                    ),
+                ],
+            ),
+            returns=call.frame,
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    @method()
+    def sort_values(self, call: Call) -> Type:
+        str_ = self.types.get_type("str")
+        bool_ = self.types.get_type("bool")
+
+        def make_overload(by: Type, ascending: Type) -> Function:
+            return Function(
+                params=ParamSpec(
+                    mixed=[
+                        Function.Parameter(
+                            pos=0,
+                            name="by",
+                            type=by,
+                            required=True,
+                        ),
+                    ],
+                    kw=[
+                        Function.Parameter(
+                            pos=1,
+                            name="axis",
+                            type=TopType(),
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=2,
+                            name="ascending",
+                            type=ascending,
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=3,
+                            name="inplace",
+                            type=bool_,
+                            required=False,
+                            unsupported=True,
+                        ),
+                        Function.Parameter(
+                            pos=4,
+                            name="kind",
+                            type=str_,
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=5,
+                            name="na_position",
+                            type=str_,
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=6,
+                            name="ignore_index",
+                            type=bool_,
+                            required=False,
+                        ),
+                        Function.Parameter(
+                            pos=7,
+                            name="key",
+                            type=TopType(),
+                            required=False,
+                        ),
+                    ],
+                ),
+                returns=call.frame,
+            )
+
+        list_of = self.types.list_of
+        overloads: list[Type] = [
+            make_overload(by=str_, ascending=bool_),
+            make_overload(by=list_of(str_), ascending=bool_),
+            make_overload(by=str_, ascending=list_of(bool_)),
+            make_overload(by=list_of(str_), ascending=list_of(bool_)),
+        ]
+
+        # TODO: check that literal strings in `by` are valid columns
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=OverloadedFunction(overloads=overloads),
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    def _filter_groupby_columns(
+        self, frame: DataFrameType, by: TypedExpr
+    ) -> DataFrameType:
+        """Remove columns passed as string literals in groupby's `by` argument
+
+        Args:
+            frame (DataFrameType): the original dataframe
+            by (TypedExpr): the by argument
+
+        Returns:
+            DataFrameType: the filtered dataframe
+        """
+        by_columns: list[str] = []
+
+        by_expr, _ = by
+
+        match by_expr:
+            case p.ListExpr(items=items):
+                for item in items:
+                    match item:
+                        case p.LiteralExpr(value=str() as name):
+                            by_columns.append(name)
+
+            case p.LiteralExpr(value=str() as name):
+                by_columns.append(name)
+
+        if len(by_columns) == 0:
+            return frame
+
+        new_columns: list[DataFrameType.Column] = []
+        for column in frame.columns:
+            if column.name in by_columns:
+                continue
+            new_columns.append(
+                DataFrameType.Column(
+                    index=len(new_columns),
+                    name=column.name,
+                    type=column.type,
+                )
+            )
+
+        return DataFrameType(columns=new_columns)
+
+    @method()
+    def groupby(self, call: Call) -> Type:
+        new_frame: DataFrameType = call.frame
+
+        by: Optional[TypedExpr] = None
+        if len(call.positional) != 0:
+            by = call.positional[0]
+        elif "by" in call.keywords:
+            by = call.keywords["by"]
+
+        if by is not None:
+            new_frame = self._filter_groupby_columns(call.frame, by)
+
+        bool_: Type = self.types.get_type("bool")
+        function: Function = Function(
+            params=ParamSpec(
+                mixed=[
+                    Function.Parameter(
+                        pos=0,
+                        name="by",
+                        type=TopType(),
+                        required=False,
+                    ),
+                    Function.Parameter(
+                        pos=1,
+                        name="level",
+                        type=TopType(),
+                        required=False,
+                    ),
+                ],
+                kw=[
+                    Function.Parameter(
+                        pos=i + 2,
+                        name=name,
+                        type=bool_,
+                        required=False,
+                    )
+                    for i, name in enumerate(
+                        ["as_index", "sort", "group_keys", "observed", "dropna"]
+                    )
+                ],
+            ),
+            returns=FrameGroupBy(frame=new_frame),
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=function,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    def _assert_same_length(self, call_expr: p.Expr, frame1: p.Expr, frame2: p.Expr):
+        """Generate an assertion to check that two frames have the same length
+
+        Args:
+            call_expr (p.Expr): the call expression, to insert the assertion
+                at the right place
+            frame1 (p.Expr): the first frame expression
+            frame2 (p.Expr): the second frame expression
+        """
+        func_name: str = "__midas_frame_same_length__"
+
+        # Efficiently compute length
+        # https://stackoverflow.com/a/15943975/11109181
+        def len_of_df(df: ast.expr) -> ast.expr:
+            return ast.Call(
+                func=ast.Name(id="len"),
+                args=[
+                    ast.Attribute(
+                        value=df,
+                        attr="index",
+                    )
+                ],
+                keywords=[],
+            )
+
+        self.assertions.define(
+            func_name,
+            ast.FunctionDef(
+                name=func_name,
+                args=ast.arguments(
+                    posonlyargs=[],
+                    args=[
+                        ast.arg(arg="frame1"),
+                        ast.arg(arg="frame2"),
+                    ],
+                    kwonlyargs=[],
+                    defaults=[],
+                    kw_defaults=[],
+                ),
+                body=[
+                    ast.Return(
+                        value=ast.Compare(
+                            left=len_of_df(ast.Name(id="frame1")),
+                            ops=[ast.Eq()],
+                            comparators=[len_of_df(ast.Name(id="frame2"))],
+                        )
+                    )
+                ],
+                decorator_list=[],
+            ),
+        )
+        self.assertions.add(
+            bound_expr=call_expr,
+            inputs=[frame1, frame2],
+            builder=lambda f1, f2: ast.Call(
+                func=ast.Name(id=func_name),
+                args=[f1, f2],
+                keywords=[],
+            ),
+            message="DataFrames must have the same length",
+        )

midas/checker/frames/utils.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generic,
+    Optional,
+    Protocol,
+    Self,
+    TypeVar,
+)
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.dispatcher import CallDispatcher
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter
+from midas.checker.types import Type, UnknownType
+from midas.generator.collector import AssertionCollector
+
+if TYPE_CHECKING:
+    from midas.checker.python import PythonTyper, TypedExpr
+
+
+class _MethodRegistryMeta(type):
+    """Meta-class for :class:`MethodRegistry`
+
+    Collects methods marked with the :func:`method` decorator into a dictionary
+    named `_methods` on the class itself
+    """
+
+    _methods: dict[str, Callable[..., Type]] = {}
+
+    def __new__(
+        cls,
+        name: str,
+        bases: tuple[type, ...],
+        namespace: dict[str, Any],
+    ):
+        new_class = super().__new__(cls, name, bases, namespace)
+        new_class._methods = {}
+        for attr in namespace.values():
+            if callable(attr) and hasattr(attr, "__method_names__"):
+                for name in attr.__method_names__:  # type: ignore
+                    new_class._methods[name] = attr  # type: ignore
+        return new_class
+
+
+class MethodCall(Protocol):
+    """A method call object
+
+    Must have at least `location`, `call_expr` and `subject` properties
+    """
+
+    @property
+    def location(self) -> Location: ...
+
+    @property
+    def call_expr(self) -> p.Expr: ...
+
+    @property
+    def subject(self) -> TypedExpr: ...
+
+
+T = TypeVar("T", bound=MethodCall)
+
+
+class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta):
+    """A registry of methods"""
+
+    def __init__(self, typer: PythonTyper) -> None:
+        self.typer: PythonTyper = typer
+
+    @property
+    def reporter(self) -> FileReporter:
+        return self.typer.reporter
+
+    @property
+    def types(self) -> TypesRegistry:
+        return self.typer.types
+
+    @property
+    def dispatcher(self) -> CallDispatcher[p.Expr]:
+        return self.typer.dispatcher
+
+    @property
+    def assertions(self) -> AssertionCollector:
+        return self.typer.assertions
+
+    def call(self, method: str, call: T) -> Type:
+        """Compute the result type of a call to the given method
+
+        Args:
+            method (str): the method's name
+            call (T): the call
+
+        Returns:
+            Type: the result type
+        """
+        func: Optional[Callable[[Self, T], Type]] = self._methods.get(method)
+        if func is None:
+            self.reporter.warning(
+                call.location, f"Unknown method {method} on {call.subject[1]}"
+            )
+            return UnknownType()
+        return func(self, call)
+
+
+_Self = TypeVar("_Self", bound=MethodRegistry[Any])
+Method = Callable[[_Self, T], Type]
+
+
+def method(*names: str) -> Callable[[Method[_Self, T]], Method[_Self, T]]:
+    """Simple decorator to mark a method as part of the registry
+
+    Args:
+        names (str): names by which the method can be called. If left empty, the
+            Python method's name will be used
+    """
+
+    def wrapper(func: Method[_Self, T]) -> Method[_Self, T]:
+        names_: tuple[str, ...] = names
+        if len(names_) == 0:
+            names_ = (func.__name__,)
+        setattr(func, "__method_names__", names_)
+        return func
+
+    return wrapper

midas/checker/midas.py

@@ -0,0 +1,476 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+import midas.ast.midas as m
+from midas.ast.location import Location
+from midas.checker.builtins import define_builtins
+from midas.checker.dispatcher import CallDispatcher, CallResult
+from midas.checker.environment import Environment
+from midas.checker.operators import MIDAS_BINARY_METHODS, MIDAS_UNARY_METHODS
+from midas.checker.preamble import Preamble
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter, Reporter
+from midas.checker.types import (
+    ColumnType,
+    ConstraintType,
+    DataFrameType,
+    DerivedType,
+    Function,
+    GenericType,
+    ParamSpec,
+    Predicate,
+    Type,
+    TypeVar,
+    UnknownType,
+)
+from midas.checker.variance import VarianceInferrer
+from midas.lexer.midas import MidasLexer
+from midas.lexer.token import Token, TokenType
+from midas.parser.midas import MidasParser
+
+
+class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], 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.dispatcher: CallDispatcher[m.Expr] = CallDispatcher[m.Expr](
+            self.types, self.reporter
+        )
+
+        self._local_variables: dict[str, TypeVar] = {}
+
+        self._predicate_params: dict[str, Type] = {}
+
+        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))
+
+        self._bool: Type = self.get_type("bool")
+
+        self._preamble: Environment = Preamble(self.types)
+
+    def set_reporter(self, reporter: FileReporter):
+        """Set the file reporter to use for diagnostics
+
+        Args:
+            reporter (FileReporter): the file reporter
+        """
+        self.reporter = reporter
+        self.dispatcher.set_reporter(reporter)
+
+    def process(self, source: str, path: Optional[str]):
+        """Process some Midas source code
+
+        Args:
+            source (str): the Midas source code
+            path (Optional[str]): the path of the source file, if known
+        """
+        reporter: FileReporter = self.reporter.for_file(path)
+        self.set_reporter(reporter)
+
+        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 type_of(self, expr: m.Expr) -> Type:
+        """Compute the type of the given expression
+
+        Args:
+            expr (m.Expr): the expression to type
+
+        Returns:
+            Type: the type of the expression
+        """
+        type: Type = expr.accept(self)
+        return 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._local_variables:
+            return self._local_variables[name]
+        return self.types.get_type(name)
+
+    def get_variable(self, name: str) -> Type:
+        """Get the type of a variable
+
+        This function will first look into the current predicate's parameters if
+        we are in a predicate definition.
+        The the variable is looked up in the preamble (i.e. global environment)
+
+        Args:
+            name (str): the name of the variable
+
+        Raises:
+            NameError: if the variable cannot be found
+
+        Returns:
+            Type: the type of the variable
+        """
+        if name in self._predicate_params:
+            return self._predicate_params[name]
+        predicate: Optional[Predicate] = self.types.lookup_predicate(name)
+        if predicate is not None:
+            return predicate.type
+
+        global_: Optional[Type] = self._preamble.get(name)
+        if global_ is not None:
+            return global_
+
+        raise NameError(f"Unknown variable '{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)
+
+        for name, type in self.types._types.items():
+            if isinstance(type, GenericType):
+                inferrer = VarianceInferrer(self.types)
+                self.types._types[name] = inferrer.infer(type)
+
+    def assert_bool(self, expr: m.Expr):
+        """Check that the given expression is a subtype of `bool` or report an error
+
+        Args:
+            expr (m.Expr): the expression to check
+        """
+        type: Type = self.type_of(expr)
+        if not self.types.is_subtype(type, self._bool):
+            self.reporter.error(expr.location, f"Must be a boolean but is {type}")
+
+    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 = DerivedType(name=name, type=type)
+        self.types.define_type(name, type)
+        self._local_variables.clear()
+        self._current_name = None
+
+    def visit_alias_stmt(self, stmt: m.AliasStmt) -> None:
+        name: str = stmt.name.lexeme
+        self._current_name = name
+        type: Type = stmt.type.accept(self)
+        self.types.define_type(name, type)
+        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,
+            )
+
+    def visit_predicate_stmt(self, stmt: m.PredicateStmt) -> None:
+        for spec in stmt.params:
+            for param in spec.mixed:
+                assert param.name is not None
+                self._predicate_params[param.name.lexeme] = param.type.accept(self)
+
+        type: Type = self.type_of(stmt.body)
+        params: list[ParamSpec] = [self._visit_param_spec(spec) for spec in stmt.params]
+
+        if not self._is_valid_predicate(type):
+            self.reporter.error(
+                stmt.body.location,
+                f"Predicate function body must evaluate to a boolean, got {type}",
+            )
+        if len(params) != 0:
+            type = self._bool
+            for spec in reversed(params):
+                type = Function(
+                    params=spec,
+                    returns=type,
+                )
+        self._predicate_params = {}
+        self.types.define_predicate(
+            stmt.name.lexeme,
+            Predicate(
+                type=type,
+                body=stmt.body,
+                alias=len(params) == 0,
+            ),
+        )
+
+    def _is_valid_predicate(self, body: Type) -> bool:
+        """Check whether the given type is valid as a predicate's body
+
+        Accepted types are either subtypes of `bool` or valid predicates
+
+        Args:
+            body (Type): the potential predicate body
+
+        Returns:
+            bool: `True` if `body` can be a predicate body, `False` otherwise
+        """
+        match body:
+            case Function(returns=returns):
+                return self._is_valid_predicate(returns)
+            case _ if self.types.is_subtype(body, self._bool):
+                return True
+            case _:
+                return False
+
+    def visit_logical_expr(self, expr: m.LogicalExpr) -> Type:
+        self.assert_bool(expr.left)
+        self.assert_bool(expr.right)
+        return self._bool
+
+    def visit_binary_expr(self, expr: m.BinaryExpr) -> Type:
+        method: Optional[str] = MIDAS_BINARY_METHODS.get(expr.operator.type)
+        if method is None:
+            self.logger.warning(f"Unsupported operator {expr.operator.lexeme}")
+            self.reporter.warning(
+                expr.location, f"Unsupported operator {expr.operator.lexeme}"
+            )
+            return UnknownType()
+
+        return self._visit_binary_expr(expr.location, expr.left, expr.right, method)
+
+    def _visit_binary_expr(
+        self,
+        location: Location,
+        left_expr: m.Expr,
+        right_expr: m.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: CallResult = self.dispatcher.get_result(
+            location=location,
+            callee=operation,
+            positional=[(right_expr, right)],
+            keywords={},
+        )
+        return result.result
+
+    def visit_unary_expr(self, expr: m.UnaryExpr) -> Type:
+        # Special case because there is no __not__ dunder method
+        match expr.operator:
+            case Token(type=TokenType.BANG):
+                return self.types.get_type("bool")
+
+        method: Optional[str] = MIDAS_UNARY_METHODS.get(expr.operator.type)
+        if method is None:
+            self.logger.warning(f"Unsupported operator {expr.operator.lexeme}")
+            self.reporter.warning(
+                expr.location, f"Unsupported operator {expr.operator.lexeme}"
+            )
+            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: CallResult = self.dispatcher.get_result(
+            location=expr.location,
+            callee=operation,
+            positional=[],
+            keywords={},
+        )
+        return result.result
+
+    def visit_call_expr(self, expr: m.CallExpr) -> Type:
+        callee: Type = expr.callee.accept(self)
+        positional: list[tuple[m.Expr, Type]] = [
+            (arg, self.type_of(arg)) for arg in expr.arguments
+        ]
+        keywords: dict[str, tuple[m.Expr, Type]] = {
+            name: (arg, self.type_of(arg)) for name, arg in expr.keywords.items()
+        }
+        result: CallResult = self.dispatcher.get_result(
+            location=expr.location,
+            callee=callee,
+            positional=positional,
+            keywords=keywords,
+        )
+        return result.result
+
+    def visit_get_expr(self, expr: m.GetExpr) -> Type:
+        object: Type = expr.expr.accept(self)
+        member: Optional[Type] = self.types.lookup_member(object, expr.name.lexeme)
+        if member is None:
+            self.reporter.error(
+                expr.location, f"Unknown member '{expr.name.lexeme}' of {object}"
+            )
+            return UnknownType()
+        return member
+
+    def visit_variable_expr(self, expr: m.VariableExpr) -> Type:
+        return self.get_variable(expr.name.lexeme)
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr) -> Type:
+        return expr.expr.accept(self)
+
+    def visit_literal_expr(self, expr: m.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_wildcard_expr(self, expr: m.WildcardExpr) -> Type:
+        return self.get_variable("_")
+
+    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:
+        match type.type:
+            case m.NamedType(name=Token(lexeme="Column")):
+                if len(type.args) != 1:
+                    self.reporter.error(
+                        type.location,
+                        f"Column requires 1 type argument, {len(type.args)} provided",
+                    )
+                    return ColumnType(type=UnknownType())
+                return ColumnType(type=type.args[0].accept(self))
+
+        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:
+        base_type: Type = type.type.accept(self)
+        self._predicate_params["_"] = base_type
+        constraint_type: Type = self.type_of(type.constraint)
+        self._predicate_params = {}
+        if not self.types.is_subtype(constraint_type, self._bool):
+            self.reporter.error(
+                type.location,
+                f"Constraint must evaluate to a boolean, got {constraint_type}",
+            )
+
+        return ConstraintType(
+            type=base_type,
+            constraint=type.constraint,
+        )
+
+    def visit_function_type(self, type: m.FunctionType) -> Type:
+        return Function(
+            params=self._visit_param_spec(type.params),
+            returns=type.returns.accept(self),
+        )
+
+    def _visit_param_spec(self, spec: m.ParamSpec) -> ParamSpec:
+        n_pos: int = len(spec.pos)
+        n_mixed: int = len(spec.mixed)
+
+        def process_param(
+            param: m.FunctionType.Parameter, i: int
+        ) -> Function.Parameter:
+            return Function.Parameter(
+                pos=i,
+                name=param.name.lexeme if param.name is not None else str(i),
+                type=param.type.accept(self),
+                required=param.required,
+            )
+
+        return ParamSpec(
+            pos=[process_param(param, i) for i, param in enumerate(spec.pos)],
+            mixed=[
+                process_param(param, i + n_pos) for i, param in enumerate(spec.mixed)
+            ],
+            kw=[
+                process_param(param, i + n_pos + n_mixed)
+                for i, param in enumerate(spec.kw)
+            ],
+        )
+
+    def visit_frame_type(self, type: m.FrameType) -> Type:
+        def process_column(i: int, col: m.FrameType.Column) -> DataFrameType.Column:
+            return DataFrameType.Column(
+                index=i,
+                name=col.name.lexeme,
+                type=ColumnType(type=col.type.accept(self)),
+            )
+
+        return DataFrameType(
+            columns=[process_column(i, col) for i, col in enumerate(type.columns)]
+        )
+
+    def _resolve_type_params(self, params: list[m.TypeParam]):
+        vars: list[TypeVar] = []
+        for param in params:
+            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

@@ -1,7 +1,9 @@
 import ast
 from typing import Type
 
-OPERATOR_METHODS: dict[Type[ast.operator], str] = {
+from midas.lexer.token import TokenType
+
+PY_OPERATOR_METHODS: dict[Type[ast.operator], str] = {
     ast.Add: "__add__",
     ast.Sub: "__sub__",
     ast.Mult: "__mul__",
@@ -17,9 +19,9 @@ OPERATOR_METHODS: dict[Type[ast.operator], str] = {
     ast.FloorDiv: "__floordiv__",
 }
 
-COMPARATOR_METHODS: dict[Type[ast.cmpop], str] = {
+PY_COMPARATOR_METHODS: dict[Type[ast.cmpop], str] = {
     ast.Eq: "__eq__",
-    # ast.NotEq: "__noteq__",
+    ast.NotEq: "__eq__",
     ast.Lt: "__lt__",
     ast.LtE: "__le__",
     ast.Gt: "__gt__",
@@ -29,3 +31,41 @@ COMPARATOR_METHODS: dict[Type[ast.cmpop], str] = {
     # ast.In: "__in__",
     # ast.NotIn: "__notin__",
 }
+
+PY_UNARY_METHODS: dict[Type[ast.unaryop], str] = {
+    ast.Invert: "__invert__",
+    # ast.Not: "",
+    ast.UAdd: "__pos__",
+    ast.USub: "__neg__",
+}
+
+
+MIDAS_BINARY_METHODS: dict[TokenType, str] = {
+    TokenType.PLUS: "__add__",
+    TokenType.MINUS: "__sub__",
+    TokenType.STAR: "__mul__",
+    TokenType.SLASH: "__truediv__",
+    # TokenType.MODULO: "__mod__",
+    # TokenType.POW: "__pow__",
+    # ast.BitOr: "__or__",
+    # ast.BitXor: "__xor__",
+    # ast.BitAnd: "__and__",
+    # ast.FloorDiv: "__floordiv__",
+    TokenType.EQUAL_EQUAL: "__eq__",
+    TokenType.BANG_EQUAL: "__eq__",
+    TokenType.LESS: "__lt__",
+    TokenType.LESS_EQUAL: "__le__",
+    TokenType.GREATER: "__gt__",
+    TokenType.GREATER_EQUAL: "__ge__",
+    # ast.Is: "__is__",
+    # ast.IsNot: "__isnot__",
+    # ast.In: "__in__",
+    # ast.NotIn: "__notin__",
+}
+
+MIDAS_UNARY_METHODS: dict[TokenType, str] = {
+    # ast.Invert: "__invert__",
+    # ast.Not: "",
+    # TokenType.PLUS: "__pos__",
+    TokenType.MINUS: "__neg__",
+}

midas/checker/preamble.py

@@ -0,0 +1,216 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Optional
+
+from midas.checker.environment import Environment
+from midas.checker.registry import TypesRegistry
+from midas.checker.types import (
+    Function,
+    GenericType,
+    OverloadedFunction,
+    ParamSpec,
+    TopType,
+    Type,
+    TypeVar,
+    UnitType,
+)
+
+
+@dataclass(frozen=True)
+class Param:
+    name: str
+    type: Type
+    required: bool = True
+
+
+class Preamble(Environment):
+    """The initial environment containing some of Python's builtin functions"""
+
+    def __init__(self, types: TypesRegistry) -> None:
+        super().__init__()
+        self._types: TypesRegistry = types
+        self._python_funcs: dict[str, Callable[..., Any]] = {}
+
+        self.define("__name__", self._types.get_type("str"))
+        self._def_type_constructor("object", object)
+        self._def_type_constructor("float", float)
+        self._def_type_constructor("int", int)
+        self._def_type_constructor("bool", bool)
+        self._def_type_constructor("str", str)
+        self._def_function(
+            name="list",
+            pos=[Param("object", TopType())],
+            returns=self._list_of(TopType()),
+            py_function=list,
+        )
+
+        # TODO: use sink
+        self._def_function(
+            name="print",
+            pos=[Param("object", TopType(), required=False)],
+            returns=UnitType(),
+            py_function=print,
+        )
+
+        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]
+            type_vars=[map_in, map_out],
+            py_function=map,
+        )
+        self._def_function(
+            name="input",
+            pos=[Param("prompt", TopType(), required=False)],
+            returns=self._types.get_type("str"),
+        )
+        self._def_function(
+            name="len",
+            pos=[Param("object", TopType())],
+            returns=self._types.get_type("int"),
+        )
+
+        T = TypeVar(name="T", bound=None)
+        self._def_overloads(
+            name="max",
+            py_function=max,
+            signatures=[
+                (
+                    [Param("arg1", T), Param("arg2", T)],
+                    [],
+                    [],
+                    T,
+                    [T],
+                ),
+                ([Param("iterable", self._list_of(T))], [], [], T, [T]),
+            ],
+        )
+        self._def_overloads(
+            name="min",
+            py_function=min,
+            signatures=[
+                (
+                    [Param("arg1", T), Param("arg2", T)],
+                    [],
+                    [],
+                    T,
+                    [T],
+                ),
+                ([Param("iterable", self._list_of(T))], [], [], T, [T]),
+            ],
+        )
+
+    def _list_of(self, item_type: str | Type) -> Type:
+        return self._types.list_of(item_type)
+
+    def _def_type_constructor(
+        self, name: str, py_function: Optional[Callable[..., Any]] = None
+    ):
+        # TODO: more specific arg types
+        self._def_function(
+            name=name,
+            pos=[Param("object", TopType(), required=False)],
+            returns=self._types.get_type(name),
+            py_function=py_function,
+        )
+
+    def _make_function(
+        self,
+        *,
+        name: str,
+        pos: list[Param] = [],
+        mixed: list[Param] = [],
+        kw: list[Param] = [],
+        returns: Type = UnitType(),
+        type_vars: list[TypeVar] = [],
+    ) -> Type:
+        def map_params(params: list[Param], offset: int) -> list[Function.Parameter]:
+            return [
+                Function.Parameter(
+                    pos=i + offset,
+                    name=param.name,
+                    type=param.type,
+                    required=param.required,
+                )
+                for i, param in enumerate(params)
+            ]
+
+        function = Function(
+            params=ParamSpec(
+                pos=map_params(pos, 0),
+                mixed=map_params(mixed, len(pos)),
+                kw=map_params(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] = [],
+        py_function: Optional[Callable[..., Any]] = None,
+    ):
+        function: Type = self._make_function(
+            name=name,
+            pos=pos,
+            mixed=mixed,
+            kw=kw,
+            returns=returns,
+            type_vars=type_vars,
+        )
+        self.define(name, function)
+        if py_function is not None:
+            self._python_funcs[name] = py_function
+
+    def _def_overloads(
+        self,
+        *,
+        name: str,
+        signatures: list[
+            tuple[list[Param], list[Param], list[Param], Type, list[TypeVar]]
+        ],
+        py_function: Optional[Callable[..., Any]] = None,
+    ):
+        overloads: list[Type] = []
+        for pos, mixed, kw, returns, type_vars in signatures:
+            overloads.append(
+                self._make_function(
+                    name=name,
+                    pos=pos,
+                    mixed=mixed,
+                    kw=kw,
+                    returns=returns,
+                    type_vars=type_vars,
+                )
+            )
+        function: Type = OverloadedFunction(overloads=overloads)
+        self.define(name, function)
+        if py_function is not None:
+            self._python_funcs[name] = py_function
+
+    def get_py_func(self, name: str) -> Optional[Callable[..., Any]]:
+        return self._python_funcs.get(name)

midas/checker/registry.py

@@ -0,0 +1,584 @@
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+from midas.ast.midas import MemberKind
+from midas.checker.builtins import BUILTIN_SUBTYPES
+from midas.checker.types import (
+    AppliedType,
+    BaseType,
+    ColumnType,
+    ConstraintType,
+    DataFrameType,
+    DerivedType,
+    Function,
+    GenericType,
+    OverloadedFunction,
+    Predicate,
+    TopType,
+    TupleType,
+    Type,
+    TypeVar,
+    UnknownType,
+    Variance,
+    substitute_typevars,
+)
+
+
+@dataclass
+class Member:
+    """A member of a type (property or method)"""
+
+    kind: MemberKind
+    type: Type
+
+
+class TypesRegistry:
+    """A registry of types, type members and predicates"""
+
+    def __init__(self) -> None:
+        self.logger: logging.Logger = logging.getLogger("TypesRegistry")
+        self._types: dict[str, Type] = {}
+        self._members: dict[str, dict[str, Member]] = {}
+        self._predicates: dict[str, Predicate] = {}
+
+    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,
+        kind: MemberKind,
+    ):
+        """Define a member on a type
+
+        If the member is a method and a member with the same name is already
+        defined on the given type, the two are combined into an :class:`OverloadedFunction`.
+
+        If the member is a property and a member with the same name is already
+        defined on the given type, the new definition is dropped and an error
+        is reported.
+
+        In any case, if a member with the same name but a different kind is
+        already defined on the given type, the new definition is dropped and
+        an error is reported.
+
+        Args:
+            type_name (str): the name of the type on which the member is defined
+            member_name (str): the name of the new member
+            member_type (Type): the type of the new member
+            kind (MemberKind): the kind of member to define (property or method)
+        """
+        members: dict[str, Member] = self._members.setdefault(type_name, {})
+        if member_name in members:
+            current: Member = members[member_name]
+            if current.kind != kind:
+                self.logger.error(
+                    f"Member '{member_name}' is already defined as a {current.kind},"
+                    + f" cannot define a {kind} with the same name"
+                )
+                return
+            if kind != MemberKind.METHOD:
+                self.logger.error(
+                    f"Member '{member_name}' already defined for type {type_name},"
+                    + " only methods can be overloaded"
+                )
+                return
+
+            combined: Type
+            match current.type:
+                case OverloadedFunction(overloads=overloads):
+                    combined = OverloadedFunction(overloads=overloads + [member_type])
+                case _:
+                    combined = OverloadedFunction(overloads=[current.type, member_type])
+            members[member_name] = Member(kind=current.kind, type=combined)
+
+        else:
+            members[member_name] = Member(kind=kind, type=member_type)
+
+    def define_predicate(self, name: str, predicate: Predicate):
+        """Define a predicate
+
+        Args:
+            name (str): the name of the new predicate
+            predicate (Predicate): the predicate to define
+
+        Raises:
+            ValueError: if a predicate with the same name is already defined
+        """
+        if name in self._predicates:
+            raise ValueError(f"Predicate {name} already defined")
+        self._predicates[name] = predicate
+
+    def is_builtin_subtype(self, name1: str, name2: str) -> bool:
+        """Check whether a type is a subtype of another base on builtin subtype rules
+
+        Args:
+            name1 (str): the name of the potential subtype
+            name2 (str): the name of the potential supertype
+
+        Returns:
+            bool: _description_
+        """
+        subtypes: set[str] = BUILTIN_SUBTYPES.get(name2, set())
+        if name1 in subtypes:
+            return True
+        for subtype in subtypes:
+            if self.is_builtin_subtype(name1, subtype):
+                return True
+        return False
+
+    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 (_, UnknownType()):
+                return True
+
+            case (TypeVar(bound=bound), _):
+                if bound is None:
+                    return False
+                return self.is_subtype(bound, type2)
+
+            case (_, TypeVar(bound=bound)):
+                if bound is None:
+                    return True
+                return self.is_subtype(type1, bound)
+
+            case (DerivedType(type=base1), _):
+                return self.is_subtype(base1, type2)
+
+            case (BaseType(name=name1), BaseType(name=name2)):
+                return self.is_builtin_subtype(name1, name2)
+
+            case (DataFrameType(columns=columns1), DataFrameType(columns=columns2)):
+                # TODO: check order?
+                by_name1: dict[str, DataFrameType.Column] = {
+                    col.name: col for col in columns1 if col.name is not None
+                }
+                for col2 in columns2:
+                    if col2.name not in by_name1:
+                        return False
+                    if not self.is_subtype(by_name1[col2.name].type, col2.type):
+                        return False
+                return True
+
+            case (ColumnType(type=inner1), ColumnType(type=inner2)):
+                # TODO: invariant, replace ColumnType with simple GenericType
+                if not self.are_equivalent(inner1, inner2):
+                    return False
+                return True
+
+            case (Function(), Function()):
+                return self.is_func_subtype(type1, type2)
+
+            case (ConstraintType(type=base1), _):
+                return self.is_subtype(base1, type2)
+
+            case (
+                AppliedType(name=name1, args=args1),
+                AppliedType(name=name2, args=args2),
+            ) if (
+                name1 == name2
+            ):
+                generic: Type = self.get_type(name1)
+                assert isinstance(generic, GenericType)
+                for param, arg1, arg2 in zip(generic.params, args1, args2):
+                    variance: Variance = param.variance
+                    if variance in {Variance.INVARIANT, Variance.COVARIANT}:
+                        if not self.is_subtype(arg1, arg2):
+                            return False
+                    if variance in {Variance.INVARIANT, Variance.CONTRAVARIANT}:
+                        if not self.is_subtype(arg2, arg1):
+                            return False
+                return True
+
+            # TODO: verify legitimacy
+            case (AppliedType(body=body), _):
+                return self.is_subtype(body, type2)
+
+        return False
+
+    def are_equivalent(self, type1: Type, type2: Type) -> bool:
+        """Check whether two types are equivalent (T <: S and S <: T)
+
+        Args:
+            type1 (Type): the first type
+            type2 (Type): the second type
+
+        Returns:
+            bool: whether `type1` is a subtype and a supertype of `type2`
+        """
+        return self.is_subtype(type1, type2) and self.is_subtype(type2, type1)
+
+    # 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.Parameter] = func1.params.pos
+        mixed1: list[Function.Parameter] = func1.params.mixed
+        kw1: dict[str, Function.Parameter] = {
+            param.name: param for param in func1.params.kw
+        }
+        pos2: list[Function.Parameter] = func2.params.pos
+        mixed2: list[Function.Parameter] = func2.params.mixed
+        kw2: dict[str, Function.Parameter] = {
+            param.name: param for param in func2.params.kw
+        }
+
+        mixed_by_pos: dict[int, Function.Parameter] = {
+            param.pos: param for param in mixed2
+        }
+        mixed_by_name: dict[str, Function.Parameter] = {
+            param.name: param for param in mixed2
+        }
+
+        def is_arg_subtype(sub: Function.Parameter, sup: Function.Parameter) -> bool:
+            if not self.is_subtype(sub.type, sup.type):
+                return False
+            if not sup.required and sub.required:
+                return False
+            return True
+
+        for param1 in pos1:
+            param2: Function.Parameter
+            if param1.pos < len(pos2):
+                param2 = pos2[param1.pos]
+            elif param1.pos in mixed_by_pos:
+                param2 = mixed_by_pos[param1.pos]
+            elif not param1.required:
+                continue
+            else:
+                return False
+            if not is_arg_subtype(param2, param1):
+                return False
+
+        for name, param1 in kw1.items():
+            param2: Function.Parameter
+            if name in kw2:
+                param2 = kw2[name]
+            elif name in mixed_by_name:
+                param2 = mixed_by_name[name]
+            elif not param1.required:
+                continue
+            else:
+                return False
+            if not is_arg_subtype(param2, param1):
+                return False
+
+        for param1 in mixed1:
+            pos_param2: Optional[Function.Parameter] = None
+            kw_param2: Optional[Function.Parameter] = None
+            if param1.name in kw2:
+                kw_param2 = kw2[param1.name]
+            elif param1.name in mixed_by_name:
+                kw_param2 = mixed_by_name[param1.name]
+            if param1.pos < len(pos2):
+                pos_param2 = pos2[param1.pos]
+            elif param1.pos in mixed_by_pos:
+                pos_param2 = mixed_by_pos[param1.pos]
+
+            # No match in func2 and arg is required
+            if pos_param2 is None and kw_param2 is None and param1.required:
+                return False
+
+            # Matching keyword argument
+            if kw_param2 is not None and not is_arg_subtype(kw_param2, param1):
+                return False
+
+            # Matching positional argument
+            if pos_param2 is not None and not is_arg_subtype(pos_param2, param1):
+                return False
+
+        mixed_positions: set[int] = {param.pos for param in mixed1}
+        mixed_names: set[str] = {param.name for param in mixed1}
+        for param2 in pos2:
+            if not param2.required:
+                continue
+            if param2.pos >= len(pos1) and param2.pos not in mixed_positions:
+                return False
+
+        for name, param2 in kw2.items():
+            if not param2.required:
+                continue
+            if name not in kw1 and name not in mixed_names:
+                return False
+
+        for param2 in mixed2:
+            if param2.required:
+                continue
+            pos_match: bool = param2.pos < len(pos1) or param2.pos in mixed_positions
+            kw_match: bool = param2.name in kw1 or param2.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:
+        """Instantiate a generic type with the given type arguments
+
+        Args:
+            type (Type): the generic
+            args (list[Type]): the type arguments
+
+        Raises:
+            ValueError: if the arguments are invalid (wrong number, bound violation)
+
+        Returns:
+            Type: the applied generic type
+        """
+        match type:
+            case DerivedType(name=name, type=base):
+                return DerivedType(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 BaseType(name="tuple"):
+                return TupleType(items=tuple(args))
+
+            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]:
+        """Lookup a member by name on a given type
+
+        This function first looks up directly on the specified type, then
+        recurse through supertypes until it finds the member or reaches
+        the root type
+
+        Args:
+            type (Type): the type on which to lookup the member
+            member_name (str): the member's name
+
+        Returns:
+            Optional[Type]: the member's type, or `None` if it is not defined
+        """
+        match type:
+            case BaseType(name=name):
+                if name in self._members:
+                    if member_name in self._members[name]:
+                        return self._members[name][member_name].type
+                return None
+
+            case DerivedType(name=name, type=base):
+                if name in self._members:
+                    if member_name in self._members[name]:
+                        return self._members[name][member_name].type
+                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].type
+                        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 ConstraintType(type=base):
+                return self.lookup_member(base, member_name)
+
+            case TypeVar(bound=bound) if bound is not None:
+                return self.lookup_member(bound, member_name)
+
+            case UnknownType():
+                return UnknownType()
+
+            case _:
+                self.logger.debug(f"Can't get member on {type}")
+                return None
+
+    def lookup_predicate(self, name: str) -> Optional[Predicate]:
+        """Lookup a predicate by name
+
+        Args:
+            name (str): the name of the predicate
+
+        Returns:
+            Optional[Predicate]: the predicate, or `None` if is not defined
+        """
+        return self._predicates.get(name)
+
+    def _by_name_or_type(self, name_or_type: str | Type) -> Type:
+        """Get a type by name or return it as is
+
+        If `name_or_type` is a string, the associated type is looked up and returned.
+        Otherwise, the type is returned as is.
+
+        Args:
+            name_or_type (str | Type): the type or type's name
+
+        Returns:
+            Type: the type
+        """
+        if isinstance(name_or_type, str):
+            return self.get_type(name_or_type)
+        return name_or_type
+
+    def list_of(self, item_type: str | Type) -> Type:
+        """Helper method to type a list of a given item type
+
+        Args:
+            item_type (str | Type): the item type
+
+        Returns:
+            Type: the list type
+        """
+        list_ = self.get_type("list")
+        return self.apply_generic(list_, [self._by_name_or_type(item_type)])
+
+    def tuple_of(self, *item_types: str | Type) -> Type:
+        """Helper method to type a tuple of given item types
+
+        Args:
+            item_type (str | Type): the item types
+
+        Returns:
+            Type: the tuple type
+        """
+
+        tuple_ = self.get_type("tuple")
+        return self.apply_generic(
+            tuple_,
+            [self._by_name_or_type(item_type) for item_type in item_types],
+        )
+
+    def dict_of(self, key_type: str | Type, value_type: str | Type) -> Type:
+        """Helper method to type a dict of given key and value types
+
+        Args:
+            key_type (str | Type): the key type
+            value_type (str | Type): the value type
+
+        Returns:
+            Type: the dict type
+        """
+        dict_ = self.get_type("dict")
+        return self.apply_generic(
+            dict_,
+            [
+                self._by_name_or_type(key_type),
+                self._by_name_or_type(value_type),
+            ],
+        )

midas/checker/reporter.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Optional
+
+from midas.ast.location import Location
+from midas.checker.diagnostic import Diagnostic, DiagnosticType
+
+
+class Reporter:
+    """Helper class to store diagnostics"""
+
+    def __init__(self):
+        self.diagnostics: list[Diagnostic] = []
+
+    def report(
+        self,
+        path: Optional[str],
+        type: DiagnosticType,
+        location: Location,
+        message: str,
+    ):
+        """Create and record a diagnostic
+
+        Args:
+            path (Optional[str]): the path linked to this diagnostic
+            type (DiagnosticType): the type of diagnostic
+            location (Location): the location if the diagnostic in the file
+            message (str): the diagnostic's message
+        """
+        self.diagnostics.append(
+            Diagnostic(
+                file_path=path,
+                location=location,
+                type=type,
+                message=message,
+            )
+        )
+
+    def for_file(self, path: Optional[str]) -> FileReporter:
+        """Create a new file reporter for the given path using this reporter
+
+        Args:
+            path (Optional[str]): the path for the new file reporter
+
+        Returns:
+            FileReporter: the new file reporter, linked to this reporter
+        """
+        return FileReporter(self, path)
+
+
+class FileReporter:
+    """Helper class to manage diagnostics for a file"""
+
+    def __init__(self, base_reporter: Reporter, path: Optional[str]) -> None:
+        self.base_reporter: Reporter = base_reporter
+        self.path: Optional[str] = path
+        self._context: list[str] = []
+
+    def for_file(self, path: Optional[str]) -> FileReporter:
+        """Create a new file reporter for the given path with the same base reporter
+
+        Args:
+            path (Optional[str]): the path for the new file reporter
+
+        Returns:
+            FileReporter: the file reporter
+        """
+        return FileReporter(self.base_reporter, path)
+
+    @contextmanager
+    def with_context(self, ctx: str):
+        self._context.append(ctx)
+        try:
+            yield
+        finally:
+            self._context.pop()
+
+    def report(self, type: DiagnosticType, location: Location, message: str):
+        """Report a diagnostic to the base reporter
+
+        Args:
+            type (DiagnosticType): the type of diagnostic
+            location (Location): the location of the diagnostic in the file
+            message (str): the diagnostic's message
+        """
+        for ctx in self._context:
+            message = message + ", " + ctx
+        self.base_reporter.report(self.path, type, location, message)
+
+    def error(self, location: Location, message: str):
+        """Report an error diagnostic
+
+        Args:
+            location (Location): the location of the diagnostic in the file
+            message (str): the diagnostic's message
+        """
+        self.report(
+            type=DiagnosticType.ERROR,
+            location=location,
+            message=message,
+        )
+
+    def warning(self, location: Location, message: str):
+        """Report a warning diagnostic
+
+        Args:
+            location (Location): the location of the diagnostic in the file
+            message (str): the diagnostic's message
+        """
+        self.report(
+            type=DiagnosticType.WARNING,
+            location=location,
+            message=message,
+        )
+
+    def info(self, location: Location, message: str):
+        """Report an info diagnostic
+
+        Args:
+            location (Location): the location of the diagnostic in the file
+            message (str): the diagnostic's message
+        """
+        self.report(
+            type=DiagnosticType.INFO,
+            location=location,
+            message=message,
+        )
+
+    def debug(self, location: Location, message: str):
+        """Report a debug diagnostic
+
+        Args:
+            location (Location): the location of the diagnostic in the file
+            message (str): the diagnostic's message
+        """
+        self.report(
+            type=DiagnosticType.DEBUG,
+            location=location,
+            message=message,
+        )

midas/checker/resolver.py

@@ -1,4 +1,6 @@
 import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.reporter import FileReporter
 
 
 class ResolverError(Exception): ...
@@ -11,9 +13,10 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
     scope is referred to when a variable is referenced
     """
 
-    def __init__(self):
+    def __init__(self, reporter: FileReporter):
         self.locals: dict[p.Expr, int] = {}
-        self.scopes: list[dict[str, bool]] = []
+        self.scopes: list[dict[str, bool]] = [{}]
+        self.reporter: FileReporter = reporter
 
     def resolve(self, *objects: p.Stmt | p.Expr) -> None:
         """Resolve the given statements or expressions"""
@@ -25,28 +28,28 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
         """Begin a new scope inside the current one"""
         self.scopes.append({})
 
-    def end_scope(self):
-        """Close the current scope"""
-        self.scopes.pop()
+    def end_scope(self) -> dict[str, bool]:
+        """Close and return the current scope"""
+        return self.scopes.pop()
 
-    def declare(self, name: str) -> None:
+    def declare(self, location: Location, name: str) -> None:
         """Declare a variable in the current scope
 
         This method must be called *before* evaluating the variable initializer
 
         Args:
+            location (Location): the location where the name is declared
             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"
+            self.reporter.error(
+                location,
+                f"A variable with the name '{name}' is already declared in this scope",
             )
+        else:
             scope[name] = False
 
     def define(self, name: str) -> None:
@@ -77,6 +80,20 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
                 self.locals[expr] = i
                 return
 
+    def is_declared(self, name: str) -> bool:
+        """Check whether the given variable is defined in any scope
+
+        Args:
+            name (str): the name of the variable
+
+        Returns:
+            bool: `True` if the variable is defined in a scope, `False` otherwise
+        """
+        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
 
@@ -87,8 +104,12 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
             function (p.Function): the function to resolve
         """
         self.begin_scope()
-        for param in function.all_args:
-            self.declare(param.name)
+        for param in function.params.all:
+            if param.default is not None:
+                self.resolve(param.default)
+
+        for param in function.params.all:
+            self.declare(function.location, param.name)
             self.define(param.name)
         self.resolve(*function.body)
         self.end_scope()
@@ -98,24 +119,36 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
 
     def visit_function(self, stmt: p.Function) -> None:
         # Declare before resolving body to allow recursion
-        self.declare(stmt.name)
+        self.declare(stmt.location, 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)
+        self.declare(stmt.location, 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):
-                    self.resolve_local(target, name)
-                    # TODO: declare if not found
+                if not self.is_declared(name):
+                    self.declare(target.location, name)
+                self.define(name)
+                target.accept(self)
+
+            case p.GetExpr():
+                target.accept(self)
+
+            case p.SubscriptExpr():
+                target.accept(self)
+
             case _:
-                    raise Exception(f"Unsupported assignment to {target}")
+                self.reporter.error(
+                    target.location, f"Unsupported assignment to {target}"
+                )
 
     def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
         if stmt.value is not None:
@@ -132,13 +165,43 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
         # Body
         self.begin_scope()
         self.resolve(*stmt.body)
-        self.end_scope()
+        body: dict[str, bool] = self.end_scope()
 
         # Else
         self.begin_scope()
         self.resolve(*stmt.orelse)
+        else_: dict[str, bool] = self.end_scope()
+
+        # Define variables in this scope if it was defined in both body and else blocks
+        for name, is_defined in body.items():
+            if is_defined and else_.get(name, False):
+                self.define(name)
+
+    def visit_pass(self, stmt: p.Pass) -> None:
+        pass
+
+    def visit_for_stmt(self, stmt: p.ForStmt) -> None:
+        self.begin_scope()
+        self.resolve(stmt.iterator)
+        self._visit_assign(stmt.target)
+        self.resolve(*stmt.body)
         self.end_scope()
 
+    def visit_import_stmt(self, stmt: p.ImportStmt) -> None:
+        self._resolve_imports(stmt.imports)
+
+    def visit_from_import_stmt(self, stmt: p.FromImportStmt) -> None:
+        self._resolve_imports(stmt.imports)
+
+    def _resolve_imports(self, imports: list[p.ImportAlias]) -> None:
+        for import_ in imports:
+            name: str = import_.imported_name
+            self.declare(import_.location, name)
+            self.define(name)
+
+    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)
@@ -165,8 +228,9 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
 
     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"
+            self.reporter.error(
+                expr.location,
+                f"Variable '{expr.name}' is declared but may not be defined",
             )  # aka. UnboundLocalError
         self.resolve_local(expr, expr.name)
 
@@ -174,10 +238,6 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
         self.resolve(expr.left)
         self.resolve(expr.right)
 
-    def visit_set_expr(self, expr: p.SetExpr) -> None:
-        self.resolve(expr.value)
-        self.resolve(expr.object)
-
     def visit_cast_expr(self, expr: p.CastExpr) -> None:
         self.resolve(expr.expr)
 
@@ -185,3 +245,33 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[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_tuple_expr(self, expr: p.TupleExpr) -> None:
+        for item in expr.items:
+            self.resolve(item)
+
+    def visit_raw_expr(self, expr: p.RawExpr) -> None:
+        pass

midas/checker/types.py

@@ -1,54 +1,488 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import Optional, assert_never, cast
+
+import midas.ast.midas as m
+from midas.ast.printer import MidasPrinter
+
+
+@dataclass(frozen=True, kw_only=True)
+class TopType:
+    """The top type (`Any`)"""
+
+    def __str__(self) -> str:
+        return "Any"
 
 
 @dataclass(frozen=True, kw_only=True)
 class BaseType:
+    """A base / builtin type"""
+
     name: str
 
+    def __str__(self) -> str:
+        return self.name
+
 
 @dataclass(frozen=True, kw_only=True)
-class AliasType:
+class DerivedType:
+    """A derived type, i.e. a named subtype of another type"""
+
     name: str
     type: Type
 
+    def __str__(self) -> str:
+        return self.name
+
 
 @dataclass(frozen=True, kw_only=True)
 class UnknownType:
-    pass
+    """An unknown type"""
+
+    def __str__(self) -> str:
+        return "<Unknown>"
 
 
 @dataclass(frozen=True, kw_only=True)
 class UnitType:
-    pass
+    """The unit type (`None`)"""
+
+    def __str__(self) -> str:
+        return "None"
 
 
 @dataclass(frozen=True, kw_only=True)
 class Function:
-    name: str
-    pos_args: list[Argument]
-    args: list[Argument]
-    kw_args: list[Argument]
+    """A function type"""
+
+    params: ParamSpec
     returns: Type
 
+    def __str__(self) -> str:
+        return f"{self.params} -> {self.returns}"
+
     @dataclass(frozen=True, kw_only=True)
-    class Argument:
+    class Parameter:
+        pos: int
         name: str
         type: Type
         required: bool
+        unsupported: bool = False
+
+        def __str__(self) -> str:
+            opt: str = "" if self.required else "?"
+            param: str = f"{self.name}: {self.type}{opt}"
+            if self.unsupported:
+                param = f"({param})"
+            return param
 
 
 @dataclass(frozen=True, kw_only=True)
-class ComplexType:
-    properties: dict[str, Type]
+class ParamSpec:
+    """A function's parameter spec"""
+
+    pos: list[Function.Parameter] = field(default_factory=list)
+    mixed: list[Function.Parameter] = field(default_factory=list)
+    kw: list[Function.Parameter] = field(default_factory=list)
+
+    def __str__(self) -> str:
+        params: list[str] = []
+        if len(self.pos) != 0:
+            params += list(map(str, self.pos))
+            params.append("/")
+
+        if len(self.mixed) != 0:
+            params += list(map(str, self.mixed))
+
+        if len(self.kw) != 0:
+            params.append("*")
+            params += list(map(str, self.kw))
+
+        return f"({', '.join(params)})"
 
 
 @dataclass(frozen=True, kw_only=True)
-class UnionType:
-    alternatives: list[Type]
+class OverloadedFunction:
+    """A list of method overloads"""
+
+    overloads: list[Type]
+
+    def __str__(self) -> str:
+        return "<overloaded function>"
+
+
+class Variance(StrEnum):
+    """The variance of a :class:`TypeVar`"""
+
+    INVARIANT = "INVARIANT"
+    COVARIANT = "COVARIANT"
+    CONTRAVARIANT = "CONTRAVARIANT"
+
+
+@dataclass(frozen=True, kw_only=True)
+class TypeVar:
+    """A type variable, often used as type parameters for a generic type"""
+
+    name: str
+    bound: Optional[Type]
+    variance: Variance = Variance.INVARIANT
+
+    def __str__(self) -> str:
+        variance: str = {
+            Variance.COVARIANT: "+",
+            Variance.CONTRAVARIANT: "-",
+        }.get(self.variance, "")
+        res: str = f"{variance}{self.name}"
+        if self.bound is not None:
+            res = f"{res} <: {self.bound}"
+        return res
+
+
+@dataclass(frozen=True, kw_only=True)
+class GenericType:
+    """A generic type, with type parameters and a generic body type"""
+
+    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:
+    """An instance of a :class:`GenericType`, with concrete type arguments substituted in its body"""
+
+    name: str
+    args: list[Type]
+    body: Type
+
+    def __str__(self) -> str:
+        return f"{self.name}[{', '.join(map(str, self.args))}]"
+
+
+@dataclass(frozen=True, kw_only=True)
+class ConstraintType:
+    """A type with a constraint expression"""
+
+    type: Type
+    constraint: m.Expr
+
+    def __str__(self) -> str:
+        printer = MidasPrinter()
+        return f"{self.type} where {printer.print(self.constraint)}"
+
+
+@dataclass(frozen=True, kw_only=True)
+class TupleType:
+    """A tuple type, containing any number of ordered item types"""
+
+    items: tuple[Type, ...]
+
+    def __str__(self) -> str:
+        return f"({', '.join(map(str, self.items))})"
+
+
+@dataclass(frozen=True, kw_only=True)
+class ColumnType:
+    """A column type containing items of a given unique type"""
+
+    type: Type
+
+    def __str__(self) -> str:
+        return f"Column[{self.type}]"
+
+
+@dataclass(frozen=True, kw_only=True)
+class DataFrameType:
+    """A data-frame type, containing named columns of specific :class:`ColumnType`"""
+
+    columns: list[Column]
+
+    def __str__(self) -> str:
+        schema: list[str] = [f"{col.name}: {col.type}" for col in self.columns]
+        return f"Frame[{', '.join(schema)}]"
+
+    @dataclass(frozen=True, kw_only=True)
+    class Column:
+        index: int
+        name: Optional[str]
+        type: ColumnType
+
+
+@dataclass(frozen=True, kw_only=True)
+class FrameGroupBy:
+    """A frame group-by object"""
+
+    frame: DataFrameType
+
+    def __str__(self) -> str:
+        return f"FrameGroupBy[{self.frame}]"
+
+
+@dataclass(frozen=True, kw_only=True)
+class ColumnGroupBy:
+    """A column group-by object"""
+
+    column: ColumnType
+
+    def __str__(self) -> str:
+        return f"ColumnGroupBy[{self.column}]"
+
+
+def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
+    """Substitute type variables in the given type
+
+    This function is called recursively on inner type structures
+
+    Args:
+        type (Type): the type in which to substitute type variables
+        substitutions (dict[str, Type]): a mapping of type variable names to
+            concrete types
+
+    Returns:
+        Type: the resulting type with substitutions applied
+    """
+
+    def sub_parameter(param: Function.Parameter):
+        return Function.Parameter(
+            pos=param.pos,
+            name=param.name,
+            type=substitute_typevars(param.type, substitutions),
+            required=param.required,
+        )
+
+    def sub_param_spec(spec: ParamSpec):
+        return ParamSpec(
+            pos=list(map(sub_parameter, spec.pos)),
+            mixed=list(map(sub_parameter, spec.mixed)),
+            kw=list(map(sub_parameter, spec.kw)),
+        )
+
+    def sub_column(col: DataFrameType.Column):
+        return DataFrameType.Column(
+            index=col.index,
+            name=col.name,
+            type=cast(ColumnType, substitute_typevars(col.type, substitutions)),
+        )
+
+    match type:
+        case TopType():
+            return type
+
+        case BaseType(name=name) if name in substitutions:
+            return substitutions[name]
+
+        case BaseType():
+            return type
+
+        case DerivedType(name=name, type=type2):
+            return DerivedType(
+                name=name, type=substitute_typevars(type2, substitutions)
+            )
+
+        case Function(
+            params=params,
+            returns=returns,
+        ):
+            return Function(
+                params=sub_param_spec(params),
+                returns=substitute_typevars(returns, substitutions),
+            )
+
+        case OverloadedFunction(overloads=overloads):
+            return OverloadedFunction(
+                overloads=[
+                    substitute_typevars(overload, substitutions)
+                    for overload in overloads
+                ]
+            )
+
+        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 ConstraintType():
+            return ConstraintType(
+                type=substitute_typevars(type.type, substitutions),
+                constraint=type.constraint,
+            )
+
+        case TypeVar(name=name):
+            if name in substitutions:
+                return substitutions[name]
+            raise ValueError(f"Missing TypeVar substitution for {name}")
+
+        case GenericType(name=name, params=params, body=body):
+            params2: list[TypeVar] = []
+            for param in params:
+                param2: Type = substitute_typevars(param, substitutions)
+                if not isinstance(param2, TypeVar):
+                    raise ValueError(
+                        f"Invalid type parameter substitution, expected TypeVar, got {param2}"
+                    )
+                params2.append(param2)
+            return GenericType(
+                name=name,
+                params=params2,
+                body=substitute_typevars(body, substitutions),
+            )
+
+        case TupleType(items=items):
+            return TupleType(
+                items=tuple(substitute_typevars(item, substitutions) for item in items),
+            )
+
+        case ColumnType(type=items_type):
+            return ColumnType(
+                type=substitute_typevars(items_type, substitutions),
+            )
+
+        case DataFrameType(columns=columns):
+            return DataFrameType(
+                columns=list(map(sub_column, columns)),
+            )
+
+        case FrameGroupBy(frame=frame):
+            return FrameGroupBy(
+                frame=cast(DataFrameType, substitute_typevars(frame, substitutions))
+            )
+
+        case ColumnGroupBy(column=column):
+            return ColumnGroupBy(
+                column=cast(ColumnType, substitute_typevars(column, substitutions))
+            )
+
+        case UnknownType() | UnitType():
+            return type
+
+        case TopType() | GenericType():
+            raise NotImplementedError(f"Unsupported type {type}")
+
+        # Ensure exhaustiveness
+        case _:
+            assert_never(type)
+
+
+def unfold_type(type: Type) -> Type:
+    """Unfold a chain of :class:`DerivedType` to get the root supertype
+
+    Args:
+        type (Type): the type to unfold
+
+    Returns:
+        Type: the root supertype
+    """
+    match type:
+        case DerivedType(type=ref_type):
+            return unfold_type(ref_type)
+        case _:
+            return type
+
+
+def to_annotation(type: Type) -> str:
+    """Convert the given type to a Python annotation string
+
+    Args:
+        type (Type): the type to convert
+
+    Returns:
+        str: the annotation string
+    """
+
+    def _params_annotation(spec: ParamSpec) -> str:
+        if len(spec.kw) != 0:
+            return "..."
+
+        params: str = ", ".join(
+            to_annotation(param.type) for param in spec.pos + spec.mixed
+        )
+        return f"[{params}]"
+
+    match type:
+        case TopType():
+            return "Any"
+
+        case BaseType(name=name):
+            return name
+
+        case DerivedType(name=name):
+            return name
+
+        case UnknownType():
+            return "Any"
+
+        case UnitType():
+            return "None"
+
+        case Function(params=params, returns=returns):
+            params_annot: str = _params_annotation(params)
+            return f"Callable[{params_annot}, {to_annotation(returns)}]"
+
+        case OverloadedFunction():
+            return "Callable"
+
+        case TypeVar(name=name):
+            return name
+
+        case GenericType(name=name, params=params):
+            return f"{name}[{', '.join(map(to_annotation, params))}]"
+
+        case AppliedType(name=name, args=args):
+            return f"{name}[{', '.join(map(to_annotation, args))}]"
+
+        case ConstraintType():
+            return str(type)
+
+        case TupleType(items=items):
+            return f"Tuple[{', '.join(map(to_annotation, items))}]"
+
+        case ColumnType():
+            return "pd.Series"
+
+        case DataFrameType():
+            return "pd.DataFrame"
+
+        case FrameGroupBy():
+            return "pd.api.typing.DataFrameGroupBy"
+
+        case ColumnGroupBy():
+            return "pd.api.typing.SeriesGroupBy"
+
+        case _:
+            assert_never(type)
+
+
+@dataclass(frozen=True, kw_only=True)
+class Predicate:
+    """A predicate"""
+
+    type: Type
+    body: m.Expr
+    alias: bool
 
 
 Type = (
-    BaseType | AliasType | UnknownType | UnitType | Function | ComplexType | UnionType
+    TopType
+    | BaseType
+    | DerivedType
+    | UnknownType
+    | UnitType
+    | Function
+    | OverloadedFunction
+    | TypeVar
+    | GenericType
+    | AppliedType
+    | ConstraintType
+    | TupleType
+    | ColumnType
+    | DataFrameType
+    | FrameGroupBy
+    | ColumnGroupBy
 )

midas/checker/unifier.py

@@ -0,0 +1,268 @@
+import logging
+from typing import Optional
+
+from midas.checker.registry import TypesRegistry
+from midas.checker.types import (
+    AppliedType,
+    ColumnType,
+    DataFrameType,
+    Function,
+    GenericType,
+    ParamSpec,
+    TopType,
+    Type,
+    TypeVar,
+)
+
+
+class UnificationError(Exception): ...
+
+
+class Unifier:
+    """
+    Helper class to unify generic types in concrete usages
+
+    This can be used for example when a generic function is called with concrete
+    arguments, at which point the type parameters of the function signature
+    should be resolvable
+    """
+
+    def __init__(self, types: TypesRegistry) -> None:
+        self.types: TypesRegistry = types
+        self.logger: logging.Logger = logging.getLogger("Unifier")
+
+    def unify_call(
+        self,
+        type: GenericType,
+        positional: list[Type],
+        keywords: dict[str, Type],
+    ) -> Optional[Type]:
+        """Try and unify a generic function call given concrete arguments
+
+        Args:
+            type (GenericType): the generic function type
+            positional (list[Type]): the list of positional arguments
+            keywords (dict[str, Type]): the map of keyword arguments
+
+        Returns:
+            Optional[Type]: the concrete function type if unifiable, or `None`
+        """
+        concrete_func: Function = Function(
+            params=ParamSpec(
+                pos=[
+                    Function.Parameter(
+                        pos=i,
+                        name=str(i),
+                        type=arg,
+                        required=True,
+                    )
+                    for i, arg in enumerate(positional)
+                ],
+                kw=[
+                    Function.Parameter(
+                        pos=len(positional) + i,
+                        name=name,
+                        type=arg,
+                        required=True,
+                    )
+                    for i, (name, arg) in enumerate(keywords.items())
+                ],
+            ),
+            returns=TopType(),  # TODO: use expected type
+        )
+        return self.unify_generic(type, concrete_func, match_return=False)
+
+    def unify_generic(
+        self,
+        template: GenericType,
+        concrete: Type,
+        match_return: bool = True,
+    ) -> Optional[Type]:
+        """Unify a generic type's parameters given a concrete usage
+
+        Args:
+            template (GenericType): the generic type
+            concrete (Type): a concrete usage
+            match_return (bool, optional): if `template` is a function type,
+                whether its return type must be matched (see :func:`match`).
+                Defaults to True.
+
+        Returns:
+            Optional[Type]: the concrete type if unifiable, or `None`
+        """
+        substitutions: dict[str, Type]
+        try:
+            substitutions = self.match(template.body, concrete, match_return)
+        except UnificationError:
+            return None
+
+        args: list[Type] = []
+        for param in template.params:
+            if param.name not in substitutions:
+                return None
+            args.append(substitutions[param.name])
+
+        applied: Type = self.types.apply_generic(template, args)
+        return applied
+
+    def match(
+        self,
+        template: Type,
+        concrete: Type,
+        match_return: bool = True,
+    ) -> dict[str, Type]:
+        """Match a generic type with a concrete usage, recording parameter substitutions
+
+        Args:
+            template (Type): the generic type
+            concrete (Type): a concrete usage
+            match_return (bool, optional): if `template` and `concrete` are both
+                :class:`Function`, whether their return types are also matched.
+                Defaults to True.
+
+        Raises:
+            UnificationError: if there is a conflict in parameter substitutions
+
+        Returns:
+            dict[str, Type]: the parameter substitutions which,
+                applied to `template`, yield `concrete`
+        """
+        # TODO: if concrete is Generic, record bound TypeVar. Then when merging
+        # substitutions, check that the constraint is respected
+        match (template, concrete):
+            case (TypeVar(name=name), _):
+                return {name: concrete}
+
+            case (
+                AppliedType(name=template_name, args=template_args),
+                AppliedType(name=concrete_name, args=concrete_args),
+            ) if template_name == concrete_name and len(template_args) == len(
+                concrete_args
+            ):
+                substitutions: dict[str, Type] = {}
+                for template_arg, concrete_arg in zip(template_args, concrete_args):
+                    new_substistutions: dict[str, Type] = self.match(
+                        template_arg, concrete_arg
+                    )
+                    substitutions = self.merge(substitutions, new_substistutions)
+
+                return substitutions
+
+            case (
+                DataFrameType(columns=template_columns),
+                DataFrameType(columns=concrete_columns),
+            ) if len(template_columns) == len(concrete_columns):
+                substitutions: dict[str, Type] = {}
+                for template_column, concrete_column in zip(
+                    template_columns, concrete_columns
+                ):
+                    if template_column.index != concrete_column or (
+                        template_column.name != concrete_column.name
+                    ):
+                        self.logger.debug(
+                            f"Column mismatch: template={template_column}, concrete={concrete_column}"
+                        )
+                        raise UnificationError
+                    new_substistutions: dict[str, Type] = self.match(
+                        template_column.type, concrete_column.type
+                    )
+                    substitutions = self.merge(substitutions, new_substistutions)
+                return substitutions
+
+            case (ColumnType(type=template_column), ColumnType(type=concrete_column)):
+                return self.match(template_column, concrete_column)
+
+            case (Function(), Function()):
+                mapped: list[tuple[Function.Parameter, Function.Parameter]] = (
+                    self.map_params(template, concrete)
+                )
+                substitutions: dict[str, Type] = {}
+                for template_arg, concrete_arg in mapped:
+                    arg_subs: dict[str, Type] = self.match(
+                        template_arg.type, concrete_arg.type
+                    )
+                    substitutions = self.merge(substitutions, arg_subs)
+
+                if match_return:
+                    return_subs: dict[str, Type] = self.match(
+                        template.returns, concrete.returns
+                    )
+                    substitutions = self.merge(substitutions, return_subs)
+
+                return substitutions
+
+            case _:
+                self.logger.debug(f"Can't match {concrete!r} with {template!r}")
+                return {}
+
+    def merge(self, subs1: dict[str, Type], subs2: dict[str, Type]) -> dict[str, Type]:
+        """Merge two maps of substitutions and raise an error if incompatible
+
+        Args:
+            subs1 (dict[str, Type]): the first substitutions
+            subs2 (dict[str, Type]): the second substitutions
+
+        Raises:
+            UnificationError: if there is a conflict between the two maps
+
+        Returns:
+            dict[str, Type]: the merged map of substitutions
+        """
+        merged: dict[str, Type] = subs1.copy()
+
+        for k, v in subs2.items():
+            if k in merged and merged[k] != v:
+                self.logger.debug(
+                    f"Substitution already defined for {k} with type {merged[k]}, got {v}"
+                )
+                raise UnificationError
+            merged[k] = v
+        return merged
+
+    def map_params(
+        self, func1: Function, func2: Function
+    ) -> list[tuple[Function.Parameter, Function.Parameter]]:
+        """Map parameters of two functions
+
+        Args:
+            func1 (Function): the first function
+            func2 (Function): the second function
+
+        Returns:
+            list[tuple[Function.Parameter, Function.Parameter]]: the list of parameter pairs
+        """
+        pos1: list[Function.Parameter] = func1.params.pos
+        mixed1: list[Function.Parameter] = func1.params.mixed
+        kw1: list[Function.Parameter] = func1.params.kw
+
+        pos2: list[Function.Parameter] = func2.params.pos
+        mixed2: list[Function.Parameter] = func2.params.mixed
+        kw2: list[Function.Parameter] = func2.params.kw
+
+        mapped: list[tuple[Function.Parameter, Function.Parameter]] = []
+
+        by_pos2: dict[int, Function.Parameter] = {
+            param.pos: param for param in pos2 + mixed2
+        }
+        by_name2: dict[str, Function.Parameter] = {
+            param.name: param for param in mixed2 + kw2
+        }
+
+        for arg1 in pos1:
+            if (arg2 := by_pos2.get(arg1.pos)) is not None:
+                mapped.append((arg1, arg2))
+
+        for arg1 in mixed1:
+            # Match both positionally and by name, conflicts are caught
+            # when merging substitutions
+            if (arg2 := by_pos2.get(arg1.pos)) is not None:
+                mapped.append((arg1, arg2))
+
+            if (arg2 := by_name2.get(arg1.name)) is not None:
+                mapped.append((arg1, arg2))
+
+        for arg1 in kw1:
+            if (arg2 := by_name2.get(arg1.name)) is not None:
+                mapped.append((arg1, arg2))
+
+        return mapped

midas/checker/variance.py

@@ -0,0 +1,181 @@
+from typing import Literal, Optional, cast
+
+from midas.checker.registry import Member, TypesRegistry
+from midas.checker.types import (
+    AppliedType,
+    ConstraintType,
+    Function,
+    GenericType,
+    OverloadedFunction,
+    Type,
+    TypeVar,
+    Variance,
+)
+
+Polarity = Literal[-1, 0, 1]
+
+
+class Tracker:
+    """Helper class to track the polarity of type parameter references and computer their variance"""
+
+    def __init__(self, vars: list[TypeVar]) -> None:
+        self.vars: list[TypeVar] = vars
+        self.refs: dict[str, set[Polarity]] = {var.name: set() for var in self.vars}
+
+    def record(self, var: TypeVar, polarity: Polarity):
+        """Record a polarity of the given type parameter
+
+        Args:
+            var (TypeVar): the type parameter
+            polarity (Polarity): the polarity
+        """
+        self.refs[var.name].add(polarity)
+
+    def get_updated_vars(self) -> list[TypeVar]:
+        """Get a list of the tracked type variables with their recorded variance
+
+        Returns:
+            list[TypeVar]: the list of update type parameters
+        """
+        return [
+            TypeVar(
+                name=var.name, bound=var.bound, variance=self.get_variance(var.name)
+            )
+            for var in self.vars
+        ]
+
+    def get_variance(self, name: str) -> Variance:
+        """Get the variance of a type parameter
+
+        If the type parameter is only referenced in positive positions, it is
+        covariant. If it is only referenced in negative positions, it is
+        contravariant. Otherwise, it is invariant
+
+        Args:
+            name (str): the name of the type parameter
+
+        Returns:
+            Variance: the variance of the type parameter
+        """
+        refs: set[Polarity] = self.refs[name]
+        if refs == {-1}:
+            return Variance.CONTRAVARIANT
+        if refs == {1}:
+            return Variance.COVARIANT
+        return Variance.INVARIANT
+
+    def __contains__(self, item: TypeVar | str):
+        if isinstance(item, TypeVar):
+            return item.name in self
+        return item in self.refs
+
+
+class VarianceInferrer:
+    """Helper class to compute type parameter variance"""
+
+    def __init__(self, types: TypesRegistry) -> None:
+        self.types: TypesRegistry = types
+        self.tracker: Tracker = Tracker([])
+
+    def infer(self, type: GenericType) -> GenericType:
+        """Infer the variance of a generic type's parameters
+
+        Args:
+            type (GenericType): the generic type
+
+        Returns:
+            GenericType: a new generic type with its parameters updated with
+                their inferred variance
+        """
+        self.tracker = Tracker(type.params)
+
+        self.walk(type.body, 1, type.name)
+        members: dict[str, Member] = self.types._members.get(type.name, {})
+        for name, member in members.items():
+            self.walk(member.type, 1, type.name, [f"member:'{name}'"])
+
+        return GenericType(
+            name=type.name,
+            params=self.tracker.get_updated_vars(),
+            body=type.body,
+        )
+
+    def walk(
+        self,
+        type: Type,
+        polarity: Polarity,
+        base_name: str,
+        path: Optional[list[str]] = None,
+    ):
+        """Walk the type nodes and record variance
+
+        This function recurses into type substructures (e.g. function parameters,
+        overloads, constraint type bases, etc.)
+
+        When recursing, the polarity is flipped for consumer positions (e.g. function
+        parameters) or kept the same for producer positions (e.g. return type)
+
+        Args:
+            type (Type): the type to visit
+            polarity (Polarity): the current polarity
+            base_name (str): the root generic type name (used to detect and
+                handle cyclic references)
+            path (Optional[list[str]], optional): the path to reach the current
+                type from the root generic type (used for debugging). Defaults to None.
+        """
+        if path is None:
+            path = []
+
+        match type:
+            # Arguments are negative positions -> flip polarity
+            # Return is positive position -> keep polarity
+            case Function(params=spec):
+                all_params: list[Function.Parameter] = spec.pos + spec.mixed + spec.kw
+                for param in all_params:
+                    self.walk(
+                        param.type,
+                        -polarity,
+                        base_name,
+                        path + [f"param:'{param.name}'"],
+                    )
+
+                self.walk(type.returns, polarity, base_name, path + ["return"])
+
+            # Walk all overloads
+            case OverloadedFunction(overloads=overloads):
+                for overload in overloads:
+                    self.walk(overload, polarity, base_name, path)
+
+            # If same name as root generic -> skip
+            # Get inferred variance of parameters and multiply with current
+            # polarity to recurse through arguments
+            case AppliedType(name=name, args=args):
+                # TODO: handle mutually recursive types
+                if name == base_name:
+                    return
+                generic: Type = self.types.get_type(name)
+                assert isinstance(generic, GenericType)
+                params: list[TypeVar] = generic.params
+                polarities: dict[Variance, Polarity] = {
+                    Variance.INVARIANT: 0,
+                    Variance.COVARIANT: 1,
+                    Variance.CONTRAVARIANT: -1,
+                }
+                for arg, param in zip(args, params):
+                    param_polarity: Polarity = polarities[param.variance]
+                    self.walk(
+                        arg,
+                        cast(Polarity, polarity * param_polarity),
+                        base_name,
+                        path + [f"applied:'{name}'"],
+                    )
+
+            # Walk base type
+            case ConstraintType(type=base):
+                self.walk(base, polarity, base_name, path + ["constraint"])
+
+            # Reached end
+            # If tracked, record polarity
+            case TypeVar():
+                if type in self.tracker:
+                    self.tracker.record(type, polarity)

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,9 @@
+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 .stubs import stubs as stubs
+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,51 @@
+# **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 Optional, 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)
+@click.option("-s", "--stubs", type=str, multiple=True)
+@click.option("--ignore-errors", is_flag=True)
+def compile(
+    file: TextIO,
+    types: tuple[TextIO],
+    stubs: tuple[str],
+    ignore_errors: bool,
+):
+    source: str = file.read()
+    source_path: Path = Path(file.name).resolve()
+
+    checker = TypeChecker()
+    type_files: list[tuple[Path, Optional[str]]] = []
+    for i, types_file in enumerate(types):
+        in_path: Path = Path(types_file.name).resolve()
+        checker.import_midas(in_path)
+        type_files.append((in_path, stubs[i] if i < len(stubs) else None))
+
+    typed_ast: TypedAST = checker.type_check_source(source, str(source_path))
+    diagnostics: list[Diagnostic] = checker.diagnostics.copy()
+    printer = DiagnosticPrinter()
+    printer.print_all(diagnostics)
+
+    if not ignore_errors and any(
+        map(lambda d: d.type == DiagnosticType.ERROR, diagnostics)
+    ):
+        sys.exit(1)
+
+    generator = Generator(workdir=source_path.parent, types=checker.types)
+    generator.generate(typed_ast, source_path, type_files=type_files)

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,66 @@
+# **Dump types registry**
+# ```shell
+# midas dump-registry [--types <file.midas>]
+# ```
+
+from pathlib import Path
+from typing import TextIO
+
+import click
+
+from midas.ast.printer import MidasPrinter
+from midas.checker.checker import TypeChecker
+from midas.checker.registry import Member
+from midas.checker.types import AppliedType, BaseType, DerivedType, GenericType, Type
+
+
+def base_type(type: Type) -> Type:
+    match type:
+        case BaseType():
+            return type
+        case DerivedType(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())
+
+    print("##### Types #####")
+    for name, type in checker.types._types.items():
+        members: dict[str, Member] = checker.types._members.get(name, {})
+        params: str = ""
+        if isinstance(type, GenericType):
+            params = ", ".join(map(str, type.params))
+            params = f"[{params}]"
+        print(f"{name}{params} = {base_type(type)}")
+        if len(members) != 0:
+            print(" " * 4 + "Members:")
+            for member_name, member in members.items():
+                kind: str = member.kind.name
+                print(" " * 8 + f"({kind:8}) {member_name}: {member.type}")
+
+    print("##### Predicates #####")
+    printer = MidasPrinter()
+    for name, predicate in checker.types._predicates.items():
+        body: str = printer.print(predicate.body)
+        if predicate.alias:
+            print(f"{name}: {predicate.type} = {body}")
+        else:
+            print(f"{name}{predicate.type}:")
+            body = "\n".join(
+                "    " + ("return " if i == 0 else "") + line
+                for i, line in enumerate(body.split("\n"))
+            )
+            print(body)

midas/cli/commands/stubs.py

@@ -0,0 +1,66 @@
+import ast
+import time
+from pathlib import Path
+from typing import Optional, TextIO
+
+import black
+import click
+from watchdog.events import DirModifiedEvent, FileModifiedEvent, FileSystemEventHandler
+from watchdog.observers import Observer
+
+from midas.checker.checker import TypeChecker
+from midas.generator.stubs import StubsGenerator
+
+
+def generate_stubs(in_path: Path, out_path: Path):
+    checker = TypeChecker()
+    checker.import_midas(in_path)
+
+    generator = StubsGenerator(checker.types)
+    module: ast.Module = generator.generate_stubs()
+    module = ast.fix_missing_locations(module)
+
+    output: str = ast.unparse(module)
+    output = black.format_str(output, mode=black.Mode(is_pyi=True))
+
+    out_path.write_text(output)
+
+
+class Handler(FileSystemEventHandler):
+    def __init__(self, in_path: Path, out_path: Path) -> None:
+        super().__init__()
+        self.in_path: Path = in_path
+        self.out_path: Path = out_path
+
+    def on_modified(self, event: DirModifiedEvent | FileModifiedEvent) -> None:
+        generate_stubs(self.in_path, self.out_path)
+
+
+@click.command(help="Generate stubs from Midas definitions")
+@click.argument("file", type=click.File("r"))
+@click.option("-o", "--output", type=click.File("w"))
+@click.option("-w", "--watch", is_flag=True)
+def stubs(
+    file: TextIO,
+    output: Optional[TextIO],
+    watch: bool,
+):
+    source_path: Path = Path(file.name).resolve()
+    out_path: Path = source_path.with_suffix(".pyi")
+    if output is not None:
+        out_path = Path(output.name).resolve()
+    generate_stubs(source_path, out_path)
+
+    if watch:
+        print(f"Watching {source_path}...")
+        print("Press CTRL+C to stop")
+        handler = Handler(source_path, out_path)
+        observer = Observer()
+        observer.schedule(handler, str(source_path))
+        observer.start()
+        try:
+            while True:
+                time.sleep(1)
+        except KeyboardInterrupt:
+            observer.stop()
+        observer.join()

midas/cli/commands/types.py

@@ -0,0 +1,52 @@
+# **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}",
+            )
+        )
+    diagnostics.extend(checker.diagnostics)
+    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

@@ -53,5 +53,6 @@ span {
 
     &.keyword {
         color: rgb(211, 72, 9);
+        pointer-events: none;
     }
 }

midas/cli/highlighter.py

@@ -1,6 +1,7 @@
 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
 
@@ -8,6 +9,7 @@ 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)
 
@@ -22,6 +24,15 @@ class Locatable(Protocol):
     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
@@ -123,9 +134,9 @@ class PythonHighlighter(
 
     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)
+        for arg in node.args:
+            self.wrap(arg, "arg")
+            arg.accept(self)
 
     def visit_constraint_type(self, node: p.ConstraintType) -> None:
         self.wrap(node, "constraint-type")
@@ -146,15 +157,18 @@ class PythonHighlighter(
 
     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)
+        self._highlight_param_spec(stmt.params)
         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 _highlight_param_spec(self, spec: p.ParamSpec) -> None:
+        for param in spec.all:
+            self._highlight_function_param(param)
+
+    def _highlight_function_param(self, param: p.Function.Parameter) -> None:
+        self.wrap(param, "parameter")
+        if param.type is not None:
+            param.type.accept(self)
 
     def visit_type_assign(self, stmt: p.TypeAssign) -> None:
         stmt.type.accept(self)
@@ -177,6 +191,16 @@ class PythonHighlighter(
         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: ...
@@ -199,60 +223,70 @@ class PythonHighlighter(
 
     def visit_logical_expr(self, expr: p.LogicalExpr) -> None: ...
 
-    def visit_set_expr(self, expr: p.SetExpr) -> 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)
 
-class MidasHighlighter(Highlighter, m.Stmt.Visitor[None], m.Expr.Visitor[None]):
+    def visit_dict_expr(self, expr: p.DictExpr) -> None:
+        for key in expr.keys:
+            if key is not None:
+                key.accept(self)
+        for value in expr.values:
+            value.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)
+
+    def visit_tuple_expr(self, expr: p.TupleExpr) -> None:
+        for item in expr.items:
+            item.accept(self)
+
+    def visit_raw_expr(self, expr: p.RawExpr) -> None: ...
+
+    def visit_raw_stmt(self, stmt: p.RawStmt) -> None: ...
+
+
+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_simple_type_stmt(self, stmt: m.SimpleTypeStmt) -> None:
-        self.wrap(stmt, "simple-type")
-        if stmt.template is not None:
-            stmt.template.accept(self)
-        stmt.base.accept(self)
-        if stmt.constraint is not None:
-            self.wrap(stmt.constraint, "constraint")
-            stmt.constraint.accept(self)
-
-    def visit_complex_type_stmt(self, stmt: m.ComplexTypeStmt) -> None:
-        self.wrap(stmt, "complex-type")
-        if stmt.template is not None:
-            stmt.template.accept(self)
-        for prop in stmt.properties:
-            prop.accept(self)
-
-    def visit_property_stmt(self, stmt: m.PropertyStmt) -> None:
-        self.wrap(stmt, "property")
+    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)
-        if stmt.constraint is not None:
-            self.wrap(stmt.constraint, "constraint")
-            stmt.constraint.accept(self)
 
     def visit_extend_stmt(self, stmt: m.ExtendStmt) -> None:
         self.wrap(stmt, "extend")
-        stmt.type.accept(self)
-        for op in stmt.operations:
-            op.accept(self)
-
-    def visit_op_stmt(self, stmt: m.OpStmt) -> None:
-        self.wrap(stmt, "op")
-        stmt.operand.accept(self)
-        stmt.result.accept(self)
+        for member in stmt.members:
+            member.accept(self)
 
     def visit_predicate_stmt(self, stmt: m.PredicateStmt) -> None:
         self.wrap(stmt, "predicate")
-        stmt.type.accept(self)
-        stmt.condition.accept(self)
-
-    def visit_simple_type_expr(self, expr: m.SimpleTypeExpr) -> None:
-        self.wrap(expr, "simple-type-expr")
+        self.wrap(LocatableToken(stmt.name), "predicate-name")
+        for spec in stmt.params:
+            self._visit_param_spec(spec)
+        stmt.body.accept(self)
 
     def visit_logical_expr(self, expr: m.LogicalExpr) -> None:
         self.wrap(expr, "logical-expr")
@@ -268,6 +302,14 @@ class MidasHighlighter(Highlighter, m.Stmt.Visitor[None], m.Expr.Visitor[None]):
         self.wrap(expr, "unary-expr")
         expr.right.accept(self)
 
+    def visit_call_expr(self, expr: m.CallExpr) -> None:
+        self.wrap(expr, "call-expr")
+        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: m.GetExpr) -> None:
         self.wrap(expr, "get-expr")
         expr.expr.accept(self)
@@ -282,14 +324,36 @@ class MidasHighlighter(Highlighter, m.Stmt.Visitor[None], m.Expr.Visitor[None]):
 
     def visit_wildcard_expr(self, expr: m.WildcardExpr) -> None: ...
 
-    def visit_template_expr(self, expr: m.TemplateExpr) -> None:
-        self.wrap(expr, "template")
-        expr.type.accept(self)
+    def visit_named_type(self, type: m.NamedType) -> None:
+        self.wrap(type, "named-type")
 
-    def visit_type_expr(self, expr: m.TypeExpr) -> None:
-        self.wrap(expr, "type")
-        if expr.template is not None:
-            expr.template.accept(self)
+    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_function_type(self, type: m.FunctionType) -> None:
+        self.wrap(type, "function")
+        self._visit_param_spec(type.params)
+        type.returns.accept(self)
+
+    def _visit_param_spec(self, spec: m.ParamSpec) -> None:
+        for param in spec.pos + spec.mixed + spec.kw:
+            param.type.accept(self)
+
+    def visit_frame_type(self, type: m.FrameType) -> None:
+        self.wrap(type, "frame")
+        for column in type.columns:
+            self._visit_frame_column(column)
+
+    def _visit_frame_column(self, column: m.FrameType.Column) -> None:
+        self.wrap(column, "column")
 
 
 class DiagnosticsHighlighter(Highlighter):

midas/cli/hl_diagnostic.css

@@ -4,9 +4,11 @@ span {
     &.error {
         --col: 255, 0, 0;
     }
+
     &.warning {
         --col: 250, 160, 0;
     }
+
     &.info {
         --col: 150, 190, 250;
     }
@@ -19,12 +21,12 @@ span {
         }
 
         &:hover:not(:has(.with-msg:hover)) {
-            .message {
+            &>.message {
                 display: inline-block;
             }
         }
 
-        .message {
+        &>.message {
             position: absolute;
             top: calc(100% + 0.2em);
             left: -.2em;
@@ -33,7 +35,8 @@ span {
             padding: 0.2em 0.4em;
             border-radius: .2em;
             z-index: 10;
-            width: 300%;
+            width: max-content;
+            max-width: 60vw;
         }
     }
 }

midas/cli/hl_midas.css

@@ -5,12 +5,10 @@ span {
         font-style: italic;
     }
 
-    &.simple-type {
-        --col: 108, 233, 108;
-    }
-
-    &.complex-type {
-        --col: 233, 206, 108;
+    &.named-type,
+    &.generic-type,
+    &.constraint-type {
+        --col: 150, 150, 150;
     }
 
     &.constraint {
@@ -33,10 +31,6 @@ span {
         --col: 193, 108, 233;
     }
 
-    &.simple-type-expr {
-        --col: 150, 150, 150;
-    }
-
     &.logical-expr,
     &.binary-expr,
     &.unary-expr,
@@ -48,7 +42,9 @@ span {
         --col: 163, 117, 71;
     }
 
-    &.type {
+    &.type-name,
+    &.op-name,
+    &.predicate-name {
         --col: 200, 200, 200;
         font-weight: bold;
     }

midas/cli/hl_python.css

@@ -3,7 +3,7 @@ span {
         --col: 108, 233, 108;
     }
 
-    &.param {
+    &.arg {
         --col: 103, 192, 224;
     }
 
@@ -23,7 +23,7 @@ span {
         --col: 215, 103, 224;
     }
 
-    &.argument {
+    &.parameter {
         --col: 103, 192, 224;
     }
 }

midas/cli/main.py

@@ -1,179 +1,25 @@
-import ast
-import json
 import logging
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Optional, TextIO, get_args
 
 import click
 
-import midas.ast.midas as m
-import midas.ast.python as p
-from midas.ast.location import Location
-from midas.ast.printer import MidasAstPrinter, PythonAstPrinter
-from midas.checker.checker import Checker
-from midas.checker.diagnostic import Diagnostic
-from midas.checker.types import Type
-from midas.cli.highlighter import (
-    DiagnosticsHighlighter,
-    Highlighter,
-    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
-from midas.resolver.resolver import Resolver
-from midas.utils import UniversalJSONDumper
+from midas.cli import commands
 
 
 @click.group()
-def midas():
-    click.echo("Welcome to Midas!")
+@click.option("-v", "--verbose", is_flag=True)
+def midas(verbose: bool):
+    logging.basicConfig(level=logging.DEBUG if verbose else logging.WARN)
 
 
-@midas.command()
-@click.option("-l", "--highlight", type=click.File("w"))
-@click.argument("file", type=click.File("r"))
-def compile(highlight: Optional[TextIO], file: TextIO):
-    logging.basicConfig(level=logging.DEBUG)
-    source: str = file.read()
-    tree: ast.Module = ast.parse(source, filename=file.name)
-    parser = PythonParser()
-    stmts: list[p.Stmt] = parser.parse_module(tree)
-    resolver = Resolver()
-    resolver.resolve(*stmts)
-    checker = Checker(resolver.locals, file_path=Path(file.name).resolve())
-    diagnostics: list[Diagnostic] = checker.check(stmts)
-    for diagnostic in diagnostics:
-        print(diagnostic)
-
-    print(
-        json.dumps(
-            UniversalJSONDumper.dump(
-                checker.global_env,
-                [("Environment", "_children")],
-                lambda obj: isinstance(obj, get_args(Type)),
-            ),
-            indent=4,
-        )
-    )
-    if highlight is not None:
-        highlighter = DiagnosticsHighlighter(source)
-        highlighter.highlight(diagnostics)
-        highlighter.dump(highlight)
-
-
-@midas.group()
-def utils():
-    pass
-
-
-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
-
-
-@utils.command()
-@click.option("-o", "--output", type=click.File("w"))
-@click.option("-p", "--parse", is_flag=True)
-@click.argument("file", type=click.File("r"))
-def dump_ast(output: Optional[TextIO], parse: bool, file: TextIO):
-    source: str = file.read()
-
-    dump: str
-    if file.name.endswith(".py"):
-        tree: ast.Module = ast.parse(source, filename=file.name)
-        if parse:
-            dump = dump_python_ast(tree)
-        else:
-            dump = ast.dump(tree, indent=4)
-    elif file.name.endswith(".midas"):
-        dump = dump_midas_ast(source, file.name)
-    else:
-        raise ValueError("Unsupported file type")
-
-    if output is None:
-        click.echo(dump)
-    else:
-        output.write(dump)
-
-
-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())
-
-    @dataclass(frozen=True)
-    class LocatableToken:
-        token: Token
-
-        @property
-        def location(self) -> Location:
-            return self.token.get_location()
-
-    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
-
-
-@utils.command()
-@click.option("-o", "--output", type=click.File("w"), default="-")
-@click.argument("file", type=click.File("r"))
-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.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.stubs)
+midas.add_command(commands.validate)
 
 
 if __name__ == "__main__":

midas/cli/utils.py

@@ -0,0 +1,121 @@
+from collections import defaultdict
+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:
+    COLORS: dict[DiagnosticType, int] = {
+        DiagnosticType.ERROR: Ansi.RED,
+        DiagnosticType.WARNING: Ansi.YELLOW,
+        DiagnosticType.INFO: Ansi.CYAN,
+        DiagnosticType.DEBUG: Ansi.MAGENTA,
+    }
+
+    def __init__(self) -> None:
+        self.files: dict[Optional[str], list[str]] = {}
+
+    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):
+        by_type: dict[DiagnosticType, int] = defaultdict(int)
+        for diagnostic in diagnostics:
+            filename: Optional[str] = diagnostic.file_path
+            lines = self.get_lines(filename)
+            self.print(lines, diagnostic, indent=indent)
+            by_type[diagnostic.type] += 1
+
+        if len(diagnostics) == 0:
+            return
+
+        counts: list[str] = []
+        for type in DiagnosticType:
+            if type not in by_type:
+                continue
+            count: int = by_type[type]
+            color: int = self.COLORS.get(type, Ansi.WHITE)
+            counts.append(f"{Ansi.FG(color)}{type.value}s{Ansi.RESET}: {count}")
+
+        print("    ".join(counts))
+
+    def print(self, lines: list[str], diagnostic: Diagnostic, indent: int = 4):
+        """Pretty-print a diagnostic, showing some context if possible
+
+        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:
+            self.print_multiline(lines, diagnostic, indent)
+            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 = self.COLORS.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()
+
+    def print_multiline(
+        self, all_lines: list[str], diagnostic: Diagnostic, indent: int = 4
+    ):
+        loc: Location = diagnostic.location
+        lines: list[str] = all_lines[loc.lineno - 1 : loc.end_lineno]
+
+        start_offset: int = loc.col_offset
+        end_offset: int = loc.end_col_offset or (start_offset + 1)
+
+        indent_str: str = " " * indent
+        color: int = self.COLORS.get(diagnostic.type, Ansi.WHITE)
+        res: str = indent_str + lines[0][:start_offset]
+        res += Ansi.FG(color) + lines[0][start_offset:]
+        for line in lines[1:-1]:
+            res += "\n" + indent_str + line
+        res += "\n" + indent_str + lines[-1][:end_offset]
+        res += Ansi.RESET + lines[-1][end_offset:]
+
+        print(diagnostic.location_str + ":")
+        print(res)
+        print()
+        print(Ansi.FG(color) + diagnostic.message + Ansi.RESET)
+        print()

midas/generator/collector.py

@@ -0,0 +1,124 @@
+import ast
+from dataclasses import dataclass
+from typing import Callable
+
+import midas.ast.python as p
+
+AssertionBuilder = Callable[..., ast.expr]
+"""A callback function which builds an assertion test given some input expressions"""
+
+
+@dataclass
+class Assertion:
+    """Runtime assertion to generate, bound to an expression"""
+
+    bound_expr: p.Expr
+    """The expression the assertion is bound to"""
+
+    inputs: list[p.Expr]
+    """
+    Expressions needed for the assertion
+    
+    Each expression will be converted by the generator and passed as individual
+    arguments to `builder`
+    """
+
+    builder: AssertionBuilder
+    """The callback to build the assertion test given converted expression from `inputs`"""
+
+    message: str
+    """The assertion message"""
+
+    def is_bound_to(self, expr: p.Expr) -> bool:
+        """Check whether this assertion is bound to the given expression
+
+        Args:
+            expr (p.Expr): the expression
+
+        Returns:
+            bool: whether this assertion is bound to `expr`
+        """
+        return expr == self.bound_expr
+
+
+class AssertionCollector:
+    """Helper class to collect assertions from outside the generator"""
+
+    def __init__(self):
+        self.assertions: list[Assertion] = []
+        self.definitions: dict[str, ast.stmt] = {}
+
+    def add(
+        self,
+        bound_expr: p.Expr,
+        inputs: list[p.Expr],
+        builder: AssertionBuilder,
+        message: str,
+    ):
+        """Add an assertion bound to the given expression
+
+        Args:
+            bound_expr (p.Expr): the expression before which the assertion
+                must be generated
+            inputs (list[p.Expr]): the list of input expressions (see :class:`Assertion`)
+            builder (AssertionBuilder): the builder callback (see :class:`Assertion`)
+            message (str): the assertion message
+        """
+        self.assertions.append(
+            Assertion(
+                bound_expr=bound_expr,
+                inputs=inputs,
+                builder=builder,
+                message=message,
+            )
+        )
+
+    def remove(self, assertion: Assertion):
+        """Remove the given assertion from the collection
+
+        Args:
+            assertion (Assertion): the assertion to remove
+        """
+        try:
+            self.assertions.remove(assertion)
+        except ValueError:
+            pass
+
+    def define(self, name: str, stmt: ast.stmt):
+        """Register a statement definition
+
+        This method will only register the first definition of any given name
+
+        Args:
+            name (str): the name of the definition
+            stmt (ast.stmt): the definition statement, like a function def
+        """
+        if name not in self.definitions:
+            self.definitions[name] = stmt
+
+    def get_definitions(self) -> list[ast.stmt]:
+        """Get the list of definitions
+
+        Returns:
+            list[ast.stmt]: the list of definitions
+        """
+        return list(self.definitions.values())
+
+    def get_assertions(self) -> list[Assertion]:
+        """Get the list of assertions
+
+        Returns:
+            list[Assertion]: the list of assertions
+        """
+        return self.assertions
+
+    def get_assertions_for(self, expr: p.Expr) -> list[Assertion]:
+        """Get the list of assertions bound to a given expression
+
+        Args:
+            expr (p.Expr): the expression
+
+        Returns:
+            list[Assertion]: the list of assertions bound to `expr`
+        """
+        return list(filter(lambda a: a.is_bound_to(expr), self.assertions))

midas/generator/constraints.py

@@ -0,0 +1,304 @@
+import ast
+from typing import Optional
+
+import midas.ast.midas as m
+from midas.checker.registry import TypesRegistry
+from midas.checker.types import (
+    Function,
+    ParamSpec,
+    Predicate,
+    Type,
+    to_annotation,
+)
+from midas.lexer.token import TokenType
+
+LOGICAL_OPERATORS: dict[TokenType, type[ast.boolop]] = {
+    TokenType.AND: ast.And,
+    # TokenType.OR: ast.Or,
+}
+
+BINARY_OPERATORS: dict[TokenType, type[ast.operator]] = {
+    TokenType.PLUS: ast.Add,
+    TokenType.MINUS: ast.Sub,
+    TokenType.STAR: ast.Mult,
+    TokenType.SLASH: ast.Div,
+}
+
+UNARY_OPERATORS: dict[TokenType, type[ast.unaryop]] = {
+    TokenType.PLUS: ast.UAdd,
+    TokenType.MINUS: ast.USub,
+}
+
+COMPARISON_OPERATORS: dict[TokenType, type[ast.cmpop]] = {
+    TokenType.GREATER: ast.Gt,
+    TokenType.GREATER_EQUAL: ast.GtE,
+    TokenType.LESS: ast.Lt,
+    TokenType.LESS_EQUAL: ast.LtE,
+    TokenType.EQUAL_EQUAL: ast.Eq,
+    TokenType.BANG_EQUAL: ast.NotEq,
+}
+
+
+class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
+    """Class to generate Python code for constraint expressions"""
+
+    def __init__(self, types: TypesRegistry):
+        self.types: TypesRegistry = types
+        self._id: int = 0
+        self._definitions: list[ast.stmt] = []
+        self._aliases: dict[str, str] = {}
+
+    def get_definitions(self) -> list[ast.stmt]:
+        """Get the list of definitions
+
+        Returns:
+            list[ast.stmt]: the list of definitions
+        """
+        return self._definitions
+
+    def generate(self, expr: m.Expr) -> ast.expr:
+        """Translate the given Midas expression to a Python expression
+
+        Args:
+            expr (m.Expr): the expression to translate
+
+        Returns:
+            ast.expr: the equivalent Python expression
+        """
+        match expr:
+            case m.VariableExpr():
+                return expr.accept(self)
+            case _:
+                func = Function(
+                    params=ParamSpec(
+                        mixed=[
+                            Function.Parameter(
+                                pos=0,
+                                name="_",
+                                type=self.types.get_type("Any"),
+                                required=True,
+                            )
+                        ],
+                    ),
+                    returns=self.types.get_type("bool"),
+                )
+                alias: str = self.make_alias(None)
+                definition: ast.stmt = self.make_definition(
+                    alias, Predicate(type=func, body=expr, alias=False)
+                )
+                self._definitions.append(definition)
+                return ast.Name(id=alias)
+
+    def make_alias(self, name: Optional[str]) -> str:
+        """Get a unique alias for a predicate
+
+        Args:
+            name (Optional[str]): the name of the predicate as defined by the user
+
+        Returns:
+            str: a unique name
+        """
+        suffix: str
+        if name is None:
+            suffix = f"p{self._id}"
+            self._id += 1
+        else:
+            suffix = name
+        alias: str = f"__midas_{suffix}__"
+        return alias
+
+    def make_definition(self, name: str, predicate: Predicate) -> ast.stmt:
+        """Translate the given predicate to a Python definition (or assignment)
+
+        Args:
+            name (str): the name of the predicate
+            predicate (Predicate): the predicate
+
+        Returns:
+            ast.stmt: the equivalent Python statement
+        """
+        body: ast.expr = predicate.body.accept(self)
+        if predicate.alias:
+            return ast.Assign(
+                targets=[
+                    ast.Name(id=name),
+                ],
+                value=body,
+            )
+        return self.make_func(name, [ast.Return(value=body)], predicate.type)
+
+    def make_args(self, params: ParamSpec) -> ast.arguments:
+        """Translate the given parameter spec into an `ast.arguments` node
+
+        Args:
+            params (ParamSpec): the parameter spec to translate
+
+        Returns:
+            ast.arguments: the equivalent `ast.arguments`
+        """
+        return ast.arguments(
+            posonlyargs=[
+                ast.arg(
+                    arg=param.name,
+                    annotation=ast.Constant(value=to_annotation(param.type)),
+                )
+                for param in params.pos
+            ],
+            args=[
+                ast.arg(
+                    arg=param.name,
+                    annotation=ast.Constant(value=to_annotation(param.type)),
+                )
+                for param in params.mixed
+            ],
+            kwonlyargs=[
+                ast.arg(
+                    arg=param.name,
+                    annotation=ast.Constant(value=to_annotation(param.type)),
+                )
+                for param in params.kw
+            ],
+            defaults=[],
+            kw_defaults=[],
+        )
+
+    def make_func(
+        self, name: str, inner_body: list[ast.stmt], type: Type, level: int = 0
+    ) -> ast.stmt:
+        """Generate a Python function def with the given name, body and signature
+
+        If `type` returns a function, the curried arguments are separated into
+        inner methods.
+        For example, if `type` is `(a: int) -> (b: int) -> (c: int) -> int`, the
+        following function would be generated:
+        ```python
+        def predicate(a: int):
+            def inner0(b: int):
+                def inner1(c: int):
+                    return ...
+                return inner1
+            return inner0
+        ```
+
+        Args:
+            name (str): the name of the outer function
+            inner_body (list[ast.stmt]): the body of the innermost function
+            type (Type): the function type / signature
+            level (int, optional): the current nesting level. Defaults to 0.
+
+        Raises:
+            ValueError: if `type` is not a function
+
+        Returns:
+            ast.stmt: the equivalent Python function definition
+        """
+        match type:
+            case Function(params=params, returns=Function()):
+                inner_name: str = f"inner{level}"
+                return ast.FunctionDef(
+                    name=name,
+                    args=self.make_args(params),
+                    body=[
+                        self.make_func(inner_name, inner_body, type.returns, level + 1),
+                        ast.Return(value=ast.Name(id=inner_name)),
+                    ],
+                    returns=ast.Constant(value=to_annotation(type.returns)),
+                    decorator_list=[],
+                )
+
+            case Function(params=params):
+                return ast.FunctionDef(
+                    name=name,
+                    args=self.make_args(params),
+                    body=inner_body,
+                    returns=ast.Constant(value=to_annotation(type.returns)),
+                    decorator_list=[],
+                )
+
+            case _:
+                raise ValueError(f"Expected function, got {type!r}")
+
+    def get_predicate(self, name: str) -> Optional[ast.expr]:
+        """Get a predicate's alias, and generate its definition if first reference
+
+        When calling this function for the first time for a given predicate,
+        a Python definition and an alias are generated. Subsequent calls only
+        return the alias, without re-generating the predicate's definition
+
+        Args:
+            name (str): the predicate's name
+
+        Returns:
+            Optional[ast.expr]: the predicate's alias, or `None` if it is not defined
+        """
+        if name not in self._aliases:
+            predicate: Optional[Predicate] = self.types.lookup_predicate(name)
+            if predicate is None:
+                return None
+            alias: str = self.make_alias(name)
+            self._aliases[name] = alias
+            self._definitions.append(self.make_definition(alias, predicate))
+
+        return ast.Name(id=self._aliases[name])
+
+    def visit_logical_expr(self, expr: m.LogicalExpr) -> ast.expr:
+        return ast.BoolOp(
+            op=LOGICAL_OPERATORS[expr.operator.type](),
+            values=[
+                expr.left.accept(self),
+                expr.right.accept(self),
+            ],
+        )
+
+    def visit_binary_expr(self, expr: m.BinaryExpr) -> ast.expr:
+        op: TokenType = expr.operator.type
+        if op in BINARY_OPERATORS:
+            return ast.BinOp(
+                left=expr.left.accept(self),
+                op=BINARY_OPERATORS[op](),
+                right=expr.right.accept(self),
+            )
+        if op in COMPARISON_OPERATORS:
+            return ast.Compare(
+                left=expr.left.accept(self),
+                ops=[COMPARISON_OPERATORS[op]()],
+                comparators=[expr.right.accept(self)],
+            )
+        raise ValueError(f"Unexpected binary operator {op}")
+
+    def visit_unary_expr(self, expr: m.UnaryExpr) -> ast.expr:
+        return ast.UnaryOp(
+            op=UNARY_OPERATORS[expr.operator.type](),
+            operand=expr.right.accept(self),
+        )
+
+    def visit_call_expr(self, expr: m.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: m.GetExpr) -> ast.expr:
+        return ast.Attribute(
+            value=expr.expr.accept(self),
+            attr=expr.name.lexeme,
+        )
+
+    def visit_variable_expr(self, expr: m.VariableExpr) -> ast.expr:
+        name: str = expr.name.lexeme
+        if (p := self.get_predicate(name)) is not None:
+            return p
+        return ast.Name(id=name)
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr) -> ast.expr:
+        return expr.accept(self)
+
+    def visit_literal_expr(self, expr: m.LiteralExpr) -> ast.expr:
+        return ast.Constant(value=expr.value)
+
+    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> ast.expr:
+        return ast.Name(id="_")

midas/generator/generator.py

@@ -0,0 +1,951 @@
+import ast
+import logging
+import shutil
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional, assert_never
+
+import midas.ast.midas as m
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.ast.printer import MidasPrinter
+from midas.checker.checker import TypeChecker
+from midas.checker.registry import TypesRegistry
+from midas.checker.types import (
+    AppliedType,
+    BaseType,
+    ColumnGroupBy,
+    ColumnType,
+    ConstraintType,
+    DataFrameType,
+    DerivedType,
+    FrameGroupBy,
+    Function,
+    GenericType,
+    OverloadedFunction,
+    TopType,
+    TupleType,
+    Type,
+    TypeVar,
+    UnitType,
+    UnknownType,
+)
+from midas.generator.collector import Assertion, AssertionCollector
+from midas.generator.constraints import ConstraintGenerator
+from midas.generator.stubs import StubsGenerator
+from midas.utils import TypedAST
+
+
+@dataclass
+class Scope:
+    """A simple structure to store assertions an aliases defined in a scope"""
+
+    pre_assertions: list[ast.stmt] = field(default_factory=list[ast.stmt])
+    """A list of assertions that must be generated before the scope"""
+
+    aliases: list[str] = field(default_factory=list[str])
+    """A list of aliases defined in the scope, that can be discard afterwards"""
+
+
+class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
+    """
+    A class to translate the custom Python AST back into raw `ast` nodes
+
+    This class is also responsible for generating assertions, functions for
+    predicates and other code necessary to ensure runtime safety.
+    """
+
+    IS_DATAFRAME_FUNC = "__midas_is_dataframe__"
+    IS_COLUMN_FUNC = "__midas_is_column__"
+
+    def __init__(self, workdir: Path, types: TypesRegistry) -> None:
+        self.workdir: Path = workdir.resolve()
+        self.build_dir: Path = self.workdir / "build" / "midas"
+        self.rel_src_path: Path = Path()
+        self.logger: logging.Logger = logging.getLogger("Generator")
+
+        self._typed_ast: TypedAST = TypedAST(
+            stmts=[],
+            judgements=[],
+            evaluated_casts=[],
+            assertions=AssertionCollector(),
+        )
+        self._alias_count: int = 0
+        self._predicate_count: int = 0
+        self._scopes: list[Scope] = []
+        self._aliases: list[tuple[p.Expr, ast.expr]] = []
+
+        self._constraint_generator: ConstraintGenerator = ConstraintGenerator(types)
+        self._constraints: list[tuple[m.Expr, ast.expr]] = []
+
+        self.define_is_dataframe: bool = False
+        self.define_is_column: bool = False
+
+    def set_src_path(self, path: Path):
+        """Set the current source file path
+
+        Args:
+            path (Path): the new source file path
+        """
+        self.rel_src_path = path.resolve().relative_to(self.workdir)
+
+    def generate_ast(self, typed_ast: TypedAST) -> ast.AST:
+        """Translate the given type checked AST into a Python `ast.AST`
+
+        Args:
+            typed_ast (TypedAST): the type checked Python AST
+
+        Returns:
+            ast.AST: the generated raw AST
+        """
+        self._typed_ast = typed_ast
+        body: list[ast.stmt] = self._visit_body(typed_ast.stmts, can_be_empty=True)
+        predicates: list[ast.stmt] = self._constraint_generator.get_definitions()
+        assertion_definitions: list[ast.stmt] = list(
+            typed_ast.assertions.definitions.values()
+        )
+
+        body = predicates + body
+
+        if self.define_is_dataframe:
+            body = [self._is_dataframe_definition()] + body
+
+        if self.define_is_column:
+            body = [self._is_column_definition()] + body
+
+        body = assertion_definitions + body
+
+        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,
+        type_files: Optional[list[tuple[Path, Optional[str]]]] = None,
+    ) -> Path:
+        """Generate all project files for the given source file and AST
+
+        This function calls :func:`generate_ast` to generate the output AST,
+        unparses it to runnable Python code, and also generates stubs for
+        user-defined Midas types in the same output directory
+
+        Args:
+            typed_ast (TypedAST): the type-checked AST
+            src_path (Path): the source file path
+            out_path (Optional[Path], optional): the output file path. If `None`,
+                the relative path of the source file to the working directory is
+                used to compute an equivalent path in the build directory.
+                Defaults to None.
+            type_files (Optional[list[tuple[Path, Optional[str]]]], optional):
+                the list of Midas files used to type check the AST. Defaults to None.
+
+        Raises:
+            ValueError: if `out_path` is `None` and the computed path is outside
+                the build directory
+
+        Returns:
+            Path: the actual `out_path` used
+        """
+        self.set_src_path(src_path)
+        if out_path is None:
+            if self.build_dir.exists():
+                shutil.rmtree(self.build_dir)
+            self.build_dir.mkdir(parents=True, exist_ok=True)
+            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_dir: Path = out_path.parent
+        out_dir.parent.mkdir(parents=True, exist_ok=True)
+
+        if type_files is not None:
+            for in_path, out_name in type_files:
+                if out_name is None:
+                    out_name = in_path.stem
+                self.generate_stubs(in_path, out_dir / f"{out_name}.py")
+
+        module: ast.AST = self.generate_ast(typed_ast)
+        compiled: str = ast.unparse(module)
+
+        out_path.write_text(compiled)
+        return out_path
+
+    def generate_stubs(self, in_path: Path, out_path: Path):
+        """Generate stubs from the given Midas file
+
+        Args:
+            in_path (Path): the Midas file path
+            out_path (Path): the stubs output file path
+        """
+        checker = TypeChecker()
+        checker.import_midas(in_path)
+        generator = StubsGenerator(checker.types)
+        module: ast.Module = generator.generate_stubs()
+        module = ast.fix_missing_locations(module)
+        output: str = ast.unparse(module)
+        out_path.write_text(output)
+
+    def convert(self, expr: p.Expr) -> ast.expr:
+        """Translate an expression
+
+        If the expression already has an alias, it is returned.
+        If assertions are defined for the given expression (in :attr:`TypedAST.assertions`),
+        they are materialized and added to the current scope.
+
+        Args:
+            expr (p.Expr): the expression to translate
+
+        Returns:
+            ast.expr: the translated expression
+        """
+        for expr2, alias in self._aliases:
+            if expr2 == expr:
+                return alias
+        assertions = self._typed_ast.assertions.get_assertions_for(expr)
+        if len(assertions) != 0:
+            return self._apply_assertions(expr, assertions)
+        return expr.accept(self)
+
+    def visit_binary_expr(self, expr: p.BinaryExpr) -> ast.expr:
+        return ast.BinOp(
+            left=self.convert(expr.left),
+            op=expr.operator,
+            right=self.convert(expr.right),
+        )
+
+    def visit_compare_expr(self, expr: p.CompareExpr) -> ast.expr:
+        return ast.Compare(
+            left=self.convert(expr.left),
+            ops=[expr.operator],
+            comparators=[self.convert(expr.right)],
+        )
+
+    def visit_unary_expr(self, expr: p.UnaryExpr) -> ast.expr:
+        return ast.UnaryOp(
+            op=expr.operator,
+            operand=self.convert(expr.right),
+        )
+
+    def visit_call_expr(self, expr: p.CallExpr) -> ast.expr:
+        return ast.Call(
+            func=self.convert(expr.callee),
+            args=[self.convert(arg) for arg in expr.arguments],
+            keywords=[
+                ast.keyword(arg=name, value=self.convert(arg))
+                for name, arg in expr.keywords.items()
+            ],
+        )
+
+    def visit_get_expr(self, expr: p.GetExpr) -> ast.expr:
+        return ast.Attribute(
+            value=self.convert(expr.object),
+            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=[self.convert(expr.left), self.convert(expr.right)],
+        )
+
+    def visit_cast_expr(self, expr: p.CastExpr) -> ast.expr:
+        expr2: ast.expr = self.convert(expr.expr)
+
+        if expr in self._typed_ast.evaluated_casts or expr.unsafe:
+            return expr2
+
+        alias: ast.expr = self._make_alias(expr.expr, expr2)
+
+        type: Type = self._get_expr_type(expr)
+        asserts: list[ast.stmt] = self._make_cast_asserts(expr.location, alias, type)
+        for assert_ in asserts:
+            self._add_assert(assert_)
+
+        return alias
+
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> ast.expr:
+        return ast.IfExp(
+            test=self.convert(expr.test),
+            body=self.convert(expr.if_true),
+            orelse=self.convert(expr.if_false),
+        )
+
+    def visit_list_expr(self, expr: p.ListExpr) -> ast.expr:
+        return ast.List(
+            elts=[self.convert(item) for item in expr.items],
+        )
+
+    def visit_dict_expr(self, expr: p.DictExpr) -> ast.expr:
+        return ast.Dict(
+            keys=[self.convert(key) if key is not None else None for key in expr.keys],
+            values=[self.convert(value) for value in expr.values],
+        )
+
+    def visit_subscript_expr(self, expr: p.SubscriptExpr) -> ast.expr:
+        return ast.Subscript(
+            value=self.convert(expr.object),
+            slice=self.convert(expr.index),
+        )
+
+    def visit_slice_expr(self, expr: p.SliceExpr) -> ast.expr:
+        return ast.Slice(
+            lower=self.convert(expr.lower) if expr.lower is not None else None,
+            upper=self.convert(expr.upper) if expr.upper is not None else None,
+            step=self.convert(expr.step) if expr.step is not None else None,
+        )
+
+    def visit_tuple_expr(self, expr: p.TupleExpr) -> ast.expr:
+        return ast.Tuple(
+            elts=[self.convert(item) for item in expr.items],
+        )
+
+    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=self.convert(stmt.expr),
+        )
+
+    def make_args(self, params: p.ParamSpec) -> ast.arguments:
+        """Translate a parameter spec into an `ast.arguments` node
+
+        Args:
+            params (p.ParamSpec): the parameter spec
+
+        Returns:
+            ast.arguments: the equivalent `ast.arguments`
+        """
+        return ast.arguments(
+            posonlyargs=[ast.arg(arg=param.name) for param in params.pos],
+            args=[ast.arg(arg=param.name) for param in params.mixed],
+            kwonlyargs=[ast.arg(arg=param.name) for param in params.kw],
+            defaults=[
+                self.convert(param.default)
+                for param in params.pos + params.mixed
+                if param.default is not None
+            ],
+            kw_defaults=[
+                self.convert(param.default) if param.default is not None else None
+                for param in params.kw
+            ],
+        )
+
+    def visit_function(self, stmt: p.Function) -> ast.stmt:
+        return ast.FunctionDef(
+            name=stmt.name,
+            args=self.make_args(stmt.params),
+            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=[self.convert(target) for target in stmt.targets],
+            value=self.convert(stmt.value),
+        )
+
+    def visit_return_stmt(self, stmt: p.ReturnStmt) -> ast.stmt:
+        return ast.Return(
+            value=self.convert(stmt.value) if stmt.value is not None else None,
+        )
+
+    def visit_if_stmt(self, stmt: p.IfStmt) -> ast.stmt:
+        return ast.If(
+            test=self.convert(stmt.test),
+            body=self._visit_body(stmt.body),
+            orelse=self._visit_body(stmt.orelse, can_be_empty=True),
+        )
+
+    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=self.convert(stmt.target),
+            iter=self.convert(stmt.iterator),
+            body=self._visit_body(stmt.body),
+            orelse=[],
+        )
+
+    def visit_import_stmt(self, stmt: p.ImportStmt) -> ast.stmt:
+        return ast.Import(
+            names=self._convert_imports(stmt.imports),
+        )
+
+    def visit_from_import_stmt(self, stmt: p.FromImportStmt) -> ast.stmt:
+        return ast.ImportFrom(
+            module=stmt.module,
+            names=self._convert_imports(stmt.imports),
+            level=stmt.level,
+        )
+
+    def _convert_imports(self, imports: list[p.ImportAlias]) -> list[ast.alias]:
+        """Translate a list of import aliases
+
+        Args:
+            imports (list[p.ImportAlias]): the import aliases to translate
+
+        Returns:
+            list[ast.alias]: the translated aliases
+        """
+        return [
+            ast.alias(
+                name=import_.name,
+                asname=import_.alias,
+            )
+            for import_ in imports
+        ]
+
+    def visit_raw_stmt(self, stmt: p.RawStmt) -> ast.stmt:
+        return stmt.stmt
+
+    def _visit_body(
+        self, stmts: list[p.Stmt], can_be_empty: bool = False
+    ) -> list[ast.stmt]:
+        """Translate a list of statements
+
+        Assertions generated while translating a statement are inserted before it,
+        and aliases are deleted after the statement they're used in.
+
+        Extraneous `pass` statements are automatically removed
+
+        Args:
+            stmts (list[p.Stmt]): the statements to translate
+            can_be_empty (bool, optional): if `False` and no statement is
+            generated, an `ast.Pass` statement is returned. Defaults to False.
+
+        Returns:
+            list[ast.stmt]: the generated statements
+        """
+        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)]
+        if len(generated) == 0 and not can_be_empty:
+            generated = [ast.Pass()]
+        return generated
+
+    def _make_alias(self, node: p.Expr, expr: ast.expr) -> ast.expr:
+        """Generate a unique alias for the given expression
+
+        This function creates a unique name, generates an assignment statement
+        to define the alias before the current statement, adds the alias to the
+        list of aliases defined in the current statement, and returns an
+        expression that can be used in place of `expr`
+
+        Args:
+            node (p.Expr): the AST node that generated `expr`
+            expr (ast.expr): the expression to alias
+
+        Returns:
+            ast.expr: the generated alias reference
+        """
+        name: str = f"__midas_a{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,
+            )
+        )
+        self._aliases.append((node, alias))
+        return alias
+
+    def _build_assert(self, expr: ast.expr, message: str | ast.expr) -> ast.stmt:
+        """Build an assert statement from the given test expression and message
+
+        Args:
+            expr (ast.expr): the test expression
+            message (str | ast.expr): the assert message
+
+        Returns:
+            ast.stmt: the assert statement
+        """
+        if isinstance(message, str):
+            message = ast.Constant(value=message)
+        return ast.Assert(
+            test=expr,
+            msg=message,
+        )
+
+    def _add_assert(self, assertion: ast.stmt):
+        """Append the given assertion to the current scope
+
+        Args:
+            assertion (ast.stmt): the assertion to add
+        """
+        self._scopes[-1].pre_assertions.append(assertion)
+
+    def _get_expr_type(self, query: p.Expr) -> Type:
+        """Get the type of the given expression as computed by the type checker
+
+        Args:
+            query (p.Expr): the expression
+
+        Raises:
+            RuntimeError: if no type judgment can be found for `query`
+
+        Returns:
+            Type: the type of `expr`
+        """
+        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
+    ) -> list[ast.stmt]:
+        """Generate assertions for the given cast expression
+
+        Args:
+            src_location (Location): the location of the cast expression in
+                the source file
+            expr (ast.expr): the expression being cast
+            type (Type): the target type
+
+        Returns:
+            list[ast.stmt]: the generated assertion statements
+        """
+        match type:
+            case UnknownType() | TopType():
+                return []
+
+            case BaseType(name=name):
+                return [
+                    self._build_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 DerivedType(type=base):
+                return self._make_cast_asserts(src_location, expr, base)
+
+            case UnitType():
+                return [
+                    self._build_assert(
+                        ast.Compare(
+                            left=expr,
+                            ops=[ast.Is()],
+                            comparators=[
+                                ast.Constant(value=None),
+                            ],
+                        ),
+                        self._make_cast_assert_message(src_location, expr, type),
+                    ),
+                ]
+
+            case AppliedType(body=body):
+                return self._make_cast_asserts(src_location, expr, body)
+
+            case ConstraintType(type=base, constraint=constraint):
+                asserts: list[ast.stmt] = self._make_cast_asserts(
+                    src_location, expr, base
+                )
+                asserts.append(
+                    self._make_constraint_assert(src_location, expr, constraint)
+                )
+                return asserts
+
+            case TypeVar(bound=bound):
+                # TODO: check with type from arguments / use call-site context
+                if bound is None:
+                    return []
+                return self._make_cast_asserts(src_location, expr, bound)
+
+            case TupleType(items=items):
+                asserts: list[ast.stmt] = [
+                    self._build_assert(
+                        ast.Call(
+                            func=ast.Name(id="isinstance"),
+                            args=[expr, ast.Name(id="tuple")],
+                            keywords=[],
+                        ),
+                        self._make_cast_assert_message(src_location, expr, type),
+                    ),
+                ]
+                assert isinstance(expr, ast.Tuple)
+                for item, item_type in zip(expr.elts, items):
+                    asserts.extend(
+                        self._make_cast_asserts(src_location, item, item_type)
+                    )
+                return asserts
+
+            case DataFrameType(columns=columns):
+                self.define_is_dataframe = True
+                asserts: list[ast.stmt] = [
+                    self._build_assert(
+                        ast.Call(
+                            func=ast.Name(id=self.IS_DATAFRAME_FUNC),
+                            args=[expr],
+                            keywords=[],
+                        ),
+                        self._make_cast_assert_message(
+                            src_location, expr, type, ": Not a dataframe"
+                        ),
+                    ),
+                ]
+                for column in columns:
+                    asserts.append(
+                        self._build_assert(
+                            ast.Compare(
+                                left=ast.Constant(value=column.name),
+                                ops=[ast.In()],
+                                comparators=[expr],
+                            ),
+                            self._make_cast_assert_message(
+                                src_location,
+                                expr,
+                                type,
+                                f": Missing column {column.name}",
+                            ),
+                        )
+                    )
+                    asserts.extend(
+                        self._make_cast_asserts(
+                            src_location,
+                            ast.Subscript(
+                                value=expr, slice=ast.Constant(value=column.name)
+                            ),
+                            column.type,
+                        )
+                    )
+                return asserts
+
+            case ColumnType():
+                self.define_is_column = True
+                asserts: list[ast.stmt] = [
+                    self._build_assert(
+                        ast.Call(
+                            func=ast.Name(id=self.IS_COLUMN_FUNC),
+                            args=[expr],
+                            keywords=[],
+                        ),
+                        self._make_cast_assert_message(
+                            src_location, expr, type, ": Not a column"
+                        ),
+                    ),
+                ]
+                inner_assert: Optional[ast.stmt] = self._make_column_inner_assert(
+                    src_location, expr, type
+                )
+                if inner_assert is not None:
+                    asserts.append(inner_assert)
+                return asserts
+
+            case (
+                Function()
+                | OverloadedFunction()
+                | GenericType()
+                | FrameGroupBy()
+                | ColumnGroupBy()
+            ):
+                self.logger.warning(f"Can't make assertion for type {type}")
+                return []
+
+            # Ensure exhaustiveness
+            case _:
+                assert_never(type)
+
+    def _make_cast_assert_message(
+        self,
+        location: Location,
+        expr: ast.expr,
+        type: Type,
+        extra: Optional[str] = None,
+    ) -> ast.expr:
+        """Build an AST node for a cast assertion message
+
+        The generated Python code looks like:
+        ```python
+        f"file.py:L1:1: CastError: Cannot cast {type(expr).__name__} to Type"
+        ```
+
+        Args:
+            location (Location): the location of the cast expression in the
+                source file
+            expr (ast.expr): the expression being cast
+            type (Type): the target type
+            extra (Optional[str], optional): extra text to append at the end of
+                the message. Defaults to None.
+
+        Returns:
+            ast.expr: the generated message (as an f-string)
+        """
+        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}{extra or ''}"),
+            ]
+        )
+
+    def _make_constraint_assert(
+        self, src_location: Location, expr: ast.expr, constraint: m.Expr
+    ) -> ast.stmt:
+        """Build an assertion for the given constraint on the given expression
+
+        Args:
+            src_location (Location): the location of the cast expression in the
+                source file
+            expr (ast.expr): the expression subject to `constraint`
+            constraint (m.Expr): the constraint applied on `expr`
+
+        Returns:
+            ast.stmt: the assert statement checking the constraint
+        """
+        test_func: ast.expr = self._get_constraint(constraint)
+        return self._build_assert(
+            ast.Call(
+                func=test_func,
+                args=[expr],
+                keywords=[],
+            ),
+            self._make_constraint_assert_message(src_location, constraint),
+        )
+
+    def _make_constraint_assert_message(
+        self, location: Location, constraint: m.Expr
+    ) -> ast.expr:
+        """Build an assert message for the given constraint
+
+        Args:
+            location (Location): the location of the cast expression in the
+                source file
+            constraint (m.Expr): the constraint
+
+        Returns:
+            ast.expr: the assert message
+        """
+        printer = MidasPrinter()
+        constraint_str: str = printer.print(constraint)
+        loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
+        # f"file.py:L1:1: ConstraintError: Value does not fit constraint 'v > 0'"
+        return ast.Constant(
+            f"{loc_str}: ConstraintError: Value does not fit constraint '{constraint_str}'"
+        )
+
+    def _get_constraint(self, expr: m.Expr) -> ast.expr:
+        """Get or generate a Python expression for the given constraint
+
+        Args:
+            expr (m.Expr): the constraint
+
+        Returns:
+            ast.expr: an equivalent Python expression
+        """
+        for expr2, constraint in self._constraints:
+            if expr2 == expr:
+                return constraint
+
+        constraint: ast.expr = self._constraint_generator.generate(expr)
+        self._constraints.append((expr, constraint))
+        return constraint
+
+    def _is_dataframe_definition(self) -> ast.stmt:
+        """Build a function def to check if a value is a dataframe
+
+        The function is defined as:
+        ```python
+        def IS_DATAFRAME_FUNC(obj) -> bool:
+            import pandas as pd
+            return isinstance(obj, pd.DataFrame)
+        ```
+        where `IS_DATAFRAME_FUNC` is replaced by :attr:`IS_DATAFRAME_FUNC`
+
+        Returns:
+            ast.stmt: the function def
+        """
+
+        return ast.FunctionDef(
+            name=self.IS_DATAFRAME_FUNC,
+            args=ast.arguments(
+                posonlyargs=[ast.arg(arg="obj")],
+                args=[],
+                kwonlyargs=[],
+                defaults=[],
+                kw_defaults=[],
+            ),
+            body=[
+                ast.Import(names=[ast.alias(name="pandas", asname="pd")]),
+                ast.Return(
+                    value=ast.Call(
+                        func=ast.Name(id="isinstance"),
+                        args=[
+                            ast.Name(id="obj"),
+                            ast.Attribute(
+                                value=ast.Name(id="pd"),
+                                attr="DataFrame",
+                            ),
+                        ],
+                        keywords=[],
+                    )
+                ),
+            ],
+            decorator_list=[],
+            returns=ast.Name(id="bool"),
+        )
+
+    def _is_column_definition(self) -> ast.stmt:
+        """Build a function def to check if a value is a column
+
+        The function is defined as:
+        ```python
+        def IS_COLUMN_FUNC(obj) -> bool:
+            import pandas as pd
+            return isinstance(obj, pd.Series)
+        ```
+        where `IS_COLUMN_FUNC` is replaced by :attr:`IS_COLUMN_FUNC`
+
+        Returns:
+            ast.stmt: the function def
+        """
+
+        return ast.FunctionDef(
+            name=self.IS_COLUMN_FUNC,
+            args=ast.arguments(
+                posonlyargs=[ast.arg(arg="obj")],
+                args=[],
+                kwonlyargs=[],
+                defaults=[],
+                kw_defaults=[],
+            ),
+            body=[
+                ast.Import(names=[ast.alias(name="pandas", asname="pd")]),
+                ast.Return(
+                    value=ast.Call(
+                        func=ast.Name(id="isinstance"),
+                        args=[
+                            ast.Name(id="obj"),
+                            ast.Attribute(
+                                value=ast.Name(id="pd"),
+                                attr="Series",
+                            ),
+                        ],
+                        keywords=[],
+                    )
+                ),
+            ],
+            decorator_list=[],
+            returns=ast.Name(id="bool"),
+        )
+
+    def _make_column_inner_assert(
+        self, src_location: Location, column: ast.expr, type: ColumnType
+    ) -> Optional[ast.stmt]:
+        """Build a for-loop checking the type of values inside a column
+
+        Args:
+            src_location (Location): the location of the cast expression in the
+                source file
+            column (ast.expr): the column being cast
+            type (ColumnType): the type of the column
+
+        Returns:
+            Optional[ast.stmt]: a for-loop checking the values, or `None` if no
+                assertions are necessary
+        """
+        # TODO: improve message, maybe chain contexts
+        col: ast.expr = ast.Name(id="col")
+        body: list[ast.stmt] = self._make_cast_asserts(src_location, col, type.type)
+        if len(body) == 0:
+            return None
+        return ast.For(
+            target=col,
+            iter=column,
+            body=body,
+            orelse=[],
+        )
+
+    def _convert_assertion(self, assertion: Assertion) -> ast.stmt:
+        """Generate a Python assert statement for the given assertion
+
+        Args:
+            assertion (Assertion): the assertion to translate
+
+        Returns:
+            ast.stmt: the generated assert statement
+        """
+        inputs: list[ast.expr] = []
+
+        for input in assertion.inputs:
+            converted: ast.expr = self.convert(input)
+            alias: ast.expr = self._make_alias(input, converted)
+            inputs.append(alias)
+
+        test: ast.expr = assertion.builder(*inputs)
+        location: Location = assertion.bound_expr.location
+        loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
+        return self._build_assert(
+            test, f"{loc_str}: AssertionError: {assertion.message}"
+        )
+
+    def _apply_assertions(self, expr: p.Expr, assertions: list[Assertion]) -> ast.expr:
+        """Translate the given expression, adding linked assertions to the scope
+
+        Args:
+            expr (p.Expr): the expression to translate
+            assertions (list[Assertion]): the list of assertions linked to `expr`
+
+        Returns:
+            ast.expr: the translated expression
+        """
+        for assertion in assertions:
+            assert_stmt: ast.stmt
+            assert_stmt = self._convert_assertion(assertion)
+            self._add_assert(assert_stmt)
+
+            # Mutating list in frozen dataclass
+            # Not ideal but easiest way to avoid duplicate assertions
+            self._typed_ast.assertions.remove(assertion)
+
+        return expr.accept(self)

midas/generator/stubs.py

@@ -0,0 +1,579 @@
+import ast
+from typing import Optional, assert_never
+
+import midas.ast.midas as m
+from midas.checker.registry import Member, TypesRegistry
+from midas.checker.types import (
+    AppliedType,
+    BaseType,
+    ColumnGroupBy,
+    ColumnType,
+    ConstraintType,
+    DataFrameType,
+    DerivedType,
+    FrameGroupBy,
+    Function,
+    GenericType,
+    OverloadedFunction,
+    ParamSpec,
+    TopType,
+    TupleType,
+    Type,
+    TypeVar,
+    UnitType,
+    UnknownType,
+    Variance,
+    substitute_typevars,
+)
+
+Empty = ast.Constant(value=...)
+
+
+class StubsGenerator:
+    """A class to generate Python stubs for user-defined Midas types"""
+
+    def __init__(self, types: TypesRegistry) -> None:
+        self.types: TypesRegistry = types
+        self.stubs: list[ast.stmt] = []
+        self.typing_imports: set[str] = set()
+        self.import_pandas: bool = False
+        self.protocol_idx: int = 0
+        self.stub_idx: int = 0
+        self.type_var_idx: int = 0
+        self.substitutions: dict[str, dict[str, Type]] = {}
+
+    def generate_stubs(self) -> ast.Module:
+        """Generate a Python module of stubs for all user-defined types
+
+        Returns:
+            ast.Module: the generated module
+        """
+        self.stubs = []
+        self.typing_imports = set()
+        self.import_pandas = False
+        for name, type in self.types._types.items():
+            # Skip builtin types, not just based on name so the user can override
+            # TODO: check if added members on builtin type
+            match type:
+                case BaseType(name=name_) if name == name_:
+                    continue
+                case GenericType(
+                    name=name1,
+                    body=BaseType(name=name2),
+                ) if (
+                    name == name1 == name2
+                ):
+                    continue
+            self.generate_stub(name, type)
+
+        imports: list[ast.stmt] = [
+            ast.ImportFrom(
+                module="__future__",
+                names=[ast.alias(name="annotations")],
+                level=0,
+            )
+        ]
+        if len(self.typing_imports) != 0:
+            imports.append(
+                ast.ImportFrom(
+                    module="typing",
+                    names=[
+                        ast.alias(name=name) for name in sorted(self.typing_imports)
+                    ],
+                    level=0,
+                )
+            )
+        if self.import_pandas:
+            imports.append(
+                ast.Import(
+                    names=[
+                        ast.alias(
+                            name="pandas",
+                            asname="pd",
+                        )
+                    ],
+                )
+            )
+        return ast.Module(body=imports + self.stubs, type_ignores=[])
+
+    def generate_stub(self, name: str, type: Type):
+        """Generate a stub for the given type
+
+        Args:
+            name (str): the name of the type
+            type (Type): the type
+        """
+        base_type: Type = type
+
+        # TODO: improve
+        match type:
+            case DerivedType(name=name_) | GenericType(name=name_) if name_ == name:
+                pass
+            case UnitType() if name == "None":
+                pass
+            case TopType() if name == "Any":
+                pass
+            case _:
+                alias = ast.Assign(
+                    targets=[ast.Name(id=name)], value=self.dump_type(type)
+                )
+                self.add_stub(alias)
+                return
+
+        members: dict[str, Member] = self.types._members.get(name, {})
+        if isinstance(base_type, (BaseType, TopType, UnitType)) and len(members) == 0:
+            return
+
+        bases: list[ast.expr] = []
+        substitutions: dict[str, Type] = {}
+        bases, substitutions = self.get_bases(type)
+        self.substitutions[name] = substitutions
+
+        body = self.generate_body(members, substitutions)
+        stub = ast.ClassDef(
+            name=name,
+            bases=bases,
+            body=body,
+            keywords=[],
+            decorator_list=[],
+        )
+        self.add_stub(stub)
+
+    def get_bases(self, type: Type) -> tuple[list[ast.expr], dict[str, Type]]:
+        """Get the list of class bases and type parameter substitutions for a type
+
+        Args:
+            type (Type): the type whose bases to get
+
+        Returns:
+            tuple[list[ast.expr], dict[str, Type]]: a tuple containing the list
+                of class bases (already translated to Python AST nodes), and a
+                mapping of type parameter substitutions (to replace them with
+                their generated aliases)
+        """
+        match type:
+            case DerivedType(type=base):
+                return [self.dump_type(base)], {}
+
+            case GenericType(params=params, body=body):
+                self.add_typing_import("Generic")
+                type_vars: ast.expr
+
+                params2: list[TypeVar] = self.define_type_vars(params)
+                if len(params) == 1:
+                    type_vars = ast.Name(id=params2[0].name)
+                else:
+                    type_vars = ast.Tuple(
+                        elts=[ast.Name(id=param.name) for param in params2]
+                    )
+
+                substitutions: dict[str, TypeVar] = {
+                    param.name: param2 for param, param2 in zip(params, params2)
+                }
+
+                body_bases, body_subsitutions = self.get_bases(body)
+                return (
+                    body_bases
+                    + [
+                        ast.Subscript(
+                            value=ast.Name(id="Generic"),
+                            slice=type_vars,
+                        )
+                    ],
+                    body_subsitutions | substitutions,
+                )
+
+            case ConstraintType(type=base):
+                return self.get_bases(base)
+
+            case TypeVar(bound=bound) if bound is not None:
+                return [self.dump_type(bound)], {}
+
+            case _:
+                return [], {}
+
+    def generate_body(
+        self, members: dict[str, Member], substitutions: dict[str, Type]
+    ) -> list[ast.stmt]:
+        """Generate a class body given its members
+
+        Args:
+            members (dict[str, Member]): the class members
+            substitutions (dict[str, Type]): a mapping of type parameter
+                substitutions (to replace them with their generated aliases)
+
+        Returns:
+            list[ast.stmt]: the generated class body statements
+        """
+        if len(members) == 0:
+            return [ast.Expr(value=Empty)]
+
+        body: list[ast.stmt] = []
+        for name, member in members.items():
+            type: Type = member.type
+            type = substitute_typevars(type, substitutions)
+            match member.kind:
+                case m.MemberKind.PROPERTY:
+                    body.append(
+                        ast.AnnAssign(
+                            target=ast.Name(id=name),
+                            annotation=self.dump_type(type),
+                            simple=1,
+                        )
+                    )
+                case m.MemberKind.METHOD:
+                    body.extend(self.dump_method(name, type))
+        return body
+
+    def dump_type(self, type: Type) -> ast.expr:
+        """Translate a type to a Python expression
+
+        Args:
+            type (Type): the type to translate
+
+        Returns:
+            ast.expr: the generated Python expression
+        """
+        match type:
+            case DerivedType(name=name) | GenericType(name=name) if (
+                name in self.substitutions
+            ):
+                type = substitute_typevars(type, self.substitutions[name])
+
+        match type:
+            case TopType() | UnknownType():
+                self.add_typing_import("Any")
+                return ast.Name(id="Any")
+
+            case BaseType(name=name):
+                return ast.Name(id=name)
+
+            case DerivedType(name=name):
+                return ast.Name(id=name)
+
+            case UnitType():
+                return ast.Constant(value=None)
+
+            case Function():
+                name: str = self.define_protocol(type)
+                return ast.Name(id=name)
+
+            case OverloadedFunction(overloads=overloads):
+                if len(overloads) == 1:
+                    return self.dump_type(overloads[0])
+                return ast.BinOp(
+                    left=self.dump_type(OverloadedFunction(overloads=overloads[:-1])),
+                    op=ast.BitOr(),
+                    right=self.dump_type(overloads[-1]),
+                )
+
+            case TypeVar():
+                return ast.Name(id=type.name)
+
+            case GenericType(name=name):
+                params: ast.expr
+                if len(type.params) == 1:
+                    params = self.dump_type(type.params[0])
+                else:
+                    params = ast.Tuple(
+                        elts=[self.dump_type(param) for param in type.params]
+                    )
+                return ast.Subscript(
+                    value=ast.Name(id=type.name),
+                    slice=params,
+                )
+
+            case AppliedType():
+                args: ast.expr
+                if len(type.args) == 1:
+                    args = self.dump_type(type.args[0])
+                else:
+                    args = ast.Tuple(elts=[self.dump_type(arg) for arg in type.args])
+                return ast.Subscript(
+                    value=ast.Name(id=type.name),
+                    slice=args,
+                )
+
+            case ConstraintType():
+                return self.dump_type(type.type)
+
+            case TupleType(items=items):
+                return ast.Subscript(
+                    value=ast.Name(id="tuple"),
+                    slice=ast.Tuple(
+                        elts=[self.dump_type(item) for item in items],
+                    ),
+                )
+
+            case ColumnType():
+                self.import_pandas = True
+                return ast.Attribute(
+                    value=ast.Name(id="pd"),
+                    attr="Series",
+                )
+
+            case DataFrameType():
+                self.import_pandas = True
+                return ast.Attribute(
+                    value=ast.Name(id="pd"),
+                    attr="DataFrame",
+                )
+
+            case FrameGroupBy():
+                self.import_pandas = True
+                return ast.Attribute(
+                    value=ast.Attribute(
+                        value=ast.Attribute(
+                            value=ast.Name(id="pd"),
+                            attr="api",
+                        ),
+                        attr="typing",
+                    ),
+                    attr="DataFrameGroupBy",
+                )
+
+            case ColumnGroupBy():
+                self.import_pandas = True
+                return ast.Attribute(
+                    value=ast.Attribute(
+                        value=ast.Attribute(
+                            value=ast.Name(id="pd"),
+                            attr="api",
+                        ),
+                        attr="typing",
+                    ),
+                    attr="SeriesGroupBy",
+                )
+
+            case _:
+                assert_never(type)
+
+    def dump_method(
+        self, name: str, method: Type, overloaded: bool = False
+    ) -> list[ast.stmt]:
+        """Generate definitions for a method
+
+        Args:
+            name (str): the method's name
+            method (Type): the method's type
+            overloaded (bool, optional): whether this method is part of an
+                overloaded method (used when called recursively). Defaults to False.
+
+        Returns:
+            list[ast.stmt]: the generated function definitions
+        """
+        match method:
+            case Function():
+                if overloaded:
+                    self.add_typing_import("overload")
+                return [
+                    ast.FunctionDef(
+                        name=name,
+                        args=self.dump_params(method.params, with_self=True),
+                        returns=self.dump_type(method.returns),
+                        body=[ast.Expr(value=Empty)],
+                        decorator_list=[ast.Name(id="overload")] if overloaded else [],
+                    )
+                ]
+            case OverloadedFunction(overloads=overloads):
+                stmts: list[ast.stmt] = []
+                for overload in overloads:
+                    stmts.extend(self.dump_method(name, overload, True))
+                return stmts
+            case _:
+                return [
+                    ast.AnnAssign(
+                        target=ast.Name(id=name),
+                        annotation=self.dump_type(method),
+                        simple=1,
+                    )
+                ]
+
+    def dump_params(self, params: ParamSpec, with_self: bool = False) -> ast.arguments:
+        """Generate an `ast.arguments` node for the given parameter spec
+
+        Args:
+            params (ParamSpec): the parameter spec to translate
+            with_self (bool, optional): whether to include a `self` parameter.
+                Defaults to False.
+
+        Returns:
+            ast.arguments: the generate Python AST node
+        """
+        pos: list[ast.arg] = [
+            ast.arg(
+                arg=f"_{param.pos}",
+                annotation=self.dump_type(param.type),
+            )
+            for param in params.pos
+        ]
+        mixed: list[ast.arg] = [
+            ast.arg(
+                arg=param.name,
+                annotation=self.dump_type(param.type),
+            )
+            for param in params.mixed
+        ]
+        kw: list[ast.arg] = [
+            ast.arg(
+                arg=param.name,
+                annotation=self.dump_type(param.type),
+            )
+            for param in params.kw
+        ]
+        defaults: list[ast.expr] = [
+            Empty for param in params.pos + params.mixed if not param.required
+        ]
+        kw_defaults: list[Optional[ast.expr]] = [
+            None if param.required else Empty for param in params.kw
+        ]
+        if with_self:
+            arg = ast.arg(arg="self", annotation=None)
+            if len(pos) != 0:
+                pos.insert(0, arg)
+            else:
+                mixed.insert(0, arg)
+        return ast.arguments(
+            posonlyargs=pos,
+            args=mixed,
+            kwonlyargs=kw,
+            defaults=defaults,
+            kw_defaults=kw_defaults,
+        )
+
+    def define_protocol(self, func: Function) -> str:
+        """Generate a :class:`Protocol` to use in a function stub
+
+        Args:
+            func (Function): the function signature to define
+
+        Returns:
+            str: the name of the generated protocol
+        """
+        self.add_typing_import("Protocol")
+        name: str = self.new_protocol_name()
+        protocol = ast.ClassDef(
+            name=name,
+            bases=[ast.Name(id="Protocol")],
+            keywords=[],
+            body=[
+                ast.FunctionDef(
+                    name="__call__",
+                    args=self.dump_params(func.params, with_self=True),
+                    returns=self.dump_type(func.returns),
+                    body=[ast.Expr(value=Empty)],
+                    decorator_list=[],
+                ),
+            ],
+            decorator_list=[],
+        )
+        self.add_stub(protocol)
+        return name
+
+    def new_protocol_name(self) -> str:
+        """Get a unique protocol name
+
+        Returns:
+            str: the unique protocol name
+        """
+        name: str = f"_Protocol{self.protocol_idx}"
+        self.protocol_idx += 1
+        return name
+
+    def new_stub_name(self) -> str:
+        """Get a unique stub name
+
+        Returns:
+            str: the unique stub name
+        """
+        name: str = f"_Stub_{self.stub_idx}"
+        self.stub_idx += 1
+        return name
+
+    def new_type_var_name(self) -> str:
+        """Get a unique type variable name
+
+        Returns:
+            str: the unique type variable name
+        """
+        name: str = f"_T{self.type_var_idx}"
+        self.type_var_idx += 1
+        return name
+
+    def add_stub(self, stub: ast.stmt):
+        """Append the given statement to the output
+
+        Args:
+            stub (ast.stmt): the statement to append
+        """
+        self.stubs.append(stub)
+
+    def add_typing_import(self, name: str):
+        """Add the given name to the list of names to import from `typing`
+
+        Args:
+            name (str): the name to import
+        """
+        self.typing_imports.add(name)
+
+    def define_type_vars(self, vars: list[TypeVar]) -> list[TypeVar]:
+        """Define aliases for the given type variables
+
+        Args:
+            vars (list[TypeVar]): the variables to define
+
+        Returns:
+            list[TypeVar]: new type variables named with the generated aliases
+        """
+        vars2: list[TypeVar] = []
+        for var in vars:
+            vars2.append(self.define_type_var(var))
+        return vars2
+
+    def define_type_var(self, var: TypeVar) -> TypeVar:
+        """Define a type variable alias
+
+        Args:
+            var (TypeVar): the type variable to define
+
+        Returns:
+            TypeVar: a new type variable named with a uniquely generated alias
+        """
+        name: str = self.new_type_var_name()
+        self.add_typing_import("TypeVar")
+
+        kwargs: list[ast.keyword] = []
+        if var.bound is not None:
+            kwargs.append(
+                ast.keyword(
+                    arg="bound",
+                    value=self.dump_type(var.bound),
+                )
+            )
+        if var.variance == Variance.COVARIANT:
+            kwargs.append(
+                ast.keyword(
+                    arg="covariant",
+                    value=ast.Constant(value=True),
+                )
+            )
+        elif var.variance == Variance.CONTRAVARIANT:
+            kwargs.append(
+                ast.keyword(
+                    arg="contravariant",
+                    value=ast.Constant(value=True),
+                )
+            )
+        self.add_stub(
+            ast.Assign(
+                targets=[ast.Name(id=name)],
+                value=ast.Call(
+                    func=ast.Name(id="TypeVar"),
+                    args=[
+                        ast.Constant(value=name),
+                    ],
+                    keywords=kwargs,
+                ),
+            )
+        )
+        return TypeVar(name=name, bound=None)

midas/lexer/base.py

@@ -16,9 +16,10 @@ class Lexer(ABC):
     """An abstract lexer which provides methods to easily extend it into a concrete one
 
     This implementation is based on the [_Crafting Interpreters_][1] book by Robert Nystrom,
-    more specifically on my [previous Python implementation](https://git.kb28.ch/HEL/pebble)
+    more specifically on my [previous Python implementation][2]
 
     [1]: https://craftinginterpreters.com/
+    [2]: https://git.kb28.ch/HEL/pebble
     """
 
     def __init__(self, source: str, file: Optional[str] = None) -> None:
@@ -168,6 +169,6 @@ class Lexer(ABC):
     def scan_token(self) -> None:
         """Scan a token
 
-        This function should (at least) consume the current character and produce the appropriate token(s), using `add_token`
+        This function should (at least) consume the current character and produce the appropriate token(s), using :func:`add_token`
         """
         pass

midas/lexer/midas.py

@@ -18,8 +18,6 @@ class MidasLexer(Lexer):
                 self.add_token(TokenType.LEFT_BRACE)
             case "}":
                 self.add_token(TokenType.RIGHT_BRACE)
-            case "|":
-                self.add_token(TokenType.PIPE)
             case "<":
                 self.add_token(
                     TokenType.LESS_EQUAL if self.match("=") else TokenType.LESS
@@ -34,6 +32,8 @@ class MidasLexer(Lexer):
                 )
             case "!" if self.match("="):
                 self.add_token(TokenType.BANG_EQUAL)
+            case "!":
+                self.add_token(TokenType.BANG)
             case ":":
                 self.add_token(TokenType.COLON)
             case ".":
@@ -48,16 +48,18 @@ class MidasLexer(Lexer):
                 self.add_token(TokenType.UNDERSCORE)
             case "-" if self.match(">"):
                 self.add_token(TokenType.ARROW)
-            # case "+":
-            #     self.add_token(TokenType.PLUS)
+            case "+":
+                self.add_token(TokenType.PLUS)
             case "-":
                 self.add_token(TokenType.MINUS)
-            # case "*":
-            #     self.add_token(TokenType.STAR)
+            case "*":
+                self.add_token(TokenType.STAR)
             case "/" if self.match("/"):
                 self.scan_comment()
             case "/" if self.match("*"):
                 self.scan_comment_multiline()
+            case "/":
+                self.add_token(TokenType.SLASH)
             case "\n":
                 self.add_token(TokenType.NEWLINE)
             case " " | "\r" | "\t":
@@ -69,6 +71,8 @@ class MidasLexer(Lexer):
                 ):
                     self.advance()
                 self.add_token(TokenType.WHITESPACE)
+            case '"' | "'":
+                self.scan_string(char)
             case _:
                 if char.isdigit():
                     self.scan_number()
@@ -78,6 +82,23 @@ class MidasLexer(Lexer):
                     self.error("Unexpected character")
         return None
 
+    def scan_string(self, opening: str):
+        """Scan the rest of a string and add it as a token
+
+        Args:
+            opening (str): the opening quote or double quote, to be matched
+                at the end of the string
+        """
+        while self.peek() != opening and not self.is_at_end():
+            self.advance()
+
+        if self.is_at_end():
+            self.error("Unterminated string")
+
+        self.advance()
+        value: str = self.source[self.start + 1 : self.idx - 1]
+        self.add_token(TokenType.STRING, value)
+
     def scan_number(self):
         """Scan the rest of number and add it as a token
 
@@ -87,12 +108,15 @@ class MidasLexer(Lexer):
         while self.peek().isdigit():
             self.advance()
 
+        is_float: bool = False
         if self.peek() == "." and self.peek_next().isdigit():
+            is_float = True
             self.advance()
             while self.peek().isdigit():
                 self.advance()
 
-        value: float = float(self.source[self.start : self.idx])
+        value_str: str = self.source[self.start : self.idx]
+        value: int | float = float(value_str) if is_float else int(value_str)
         self.add_token(TokenType.NUMBER, value)
 
     def scan_identifier(self):
@@ -134,6 +158,18 @@ class MidasLexer(Lexer):
         self.add_token(TokenType.COMMENT)
 
     def is_identifier_char(self, char: str, *, start: bool) -> bool:
+        """Check whether a character is a valid as part of an identifier
+
+        Identifiers can contain any alphanumerical character or underscore.
+        They cannot start with a digit.
+
+        Args:
+            char (str): the character to check
+            start (bool): whether this is the first character of the identifier
+
+        Returns:
+            bool: `True` if the character is valid, `False` otherwise
+        """
         if char == "_":
             return True
         if char.isalpha():

midas/lexer/token.py

@@ -23,13 +23,12 @@ class TokenType(Enum):
     AND = auto()
     QMARK = auto()
     DOT = auto()
-    PIPE = auto()
 
     # Operators
-    # PLUS = auto()
+    PLUS = auto()
     MINUS = auto()
-    # STAR = auto()
-    # SLASH = auto()
+    STAR = auto()
+    SLASH = auto()
     GREATER = auto()
     GREATER_EQUAL = auto()
     LESS = auto()
@@ -37,6 +36,7 @@ class TokenType(Enum):
     EQUAL = auto()
     EQUAL_EQUAL = auto()
     BANG_EQUAL = auto()
+    BANG = auto()
 
     # Literals
     IDENTIFIER = auto()
@@ -44,13 +44,17 @@ class TokenType(Enum):
     TRUE = auto()
     FALSE = auto()
     NONE = auto()
+    STRING = auto()
 
     # Keywords
     TYPE = auto()
-    OP = auto()
+    ALIAS = auto()
     PREDICATE = auto()
     EXTEND = auto()
     WHERE = auto()
+    PROP = auto()
+    DEF = auto()
+    FUNC = auto()
 
     # Misc
     COMMENT = auto()
@@ -61,13 +65,16 @@ class TokenType(Enum):
 
 KEYWORDS: dict[str, TokenType] = {
     "type": TokenType.TYPE,
-    "op": TokenType.OP,
+    "alias": TokenType.ALIAS,
     "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,
 }
 
 
@@ -98,6 +105,15 @@ class Token:
         )
 
     def location_to(self, to: Token) -> Location:
+        """Create a new :class:`Location` spanning from this token to another
+
+        Args:
+            to (Token): the end token
+
+        Returns:
+            Location: a new :class:`Location` starting at this token and ending
+                at `to`, both included
+        """
         return Location.span(self.get_location(), to.get_location())
 
     @property

midas/parser/base.py

@@ -16,6 +16,9 @@ class TokenError:
     def get_report(self) -> str:
         """Get a detailed error message
 
+        The error message is formatted as "(<position>) Error at <token>: <message>".
+        For example: "(L2:5) Error at '3': Expected ')' after arguments."
+
         Returns:
             str: the complete error message
         """
@@ -32,9 +35,10 @@ class Parser(ABC, Generic[T]):
     """An abstract parser which provides methods to easily extend it into a concrete one
 
     This implementation is based on the [_Crafting Interpreters_][1] book by Robert Nystrom,
-    more specifically on my [previous Python implementation](https://git.kb28.ch/HEL/pebble)
+    more specifically on my [previous Python implementation][2]
 
     [1]: https://craftinginterpreters.com/
+    [2]: https://git.kb28.ch/HEL/pebble
     """
 
     IGNORE: set[TokenType] = {
@@ -173,7 +177,7 @@ class Parser(ABC, Generic[T]):
             error_msg (str): the error message if the token doesn't match
 
         Raises:
-            SyntaxError: if the current token doesn't match the given type
+            ParsingError: if the current token doesn't match the given type
 
         Returns:
             Token: the current token which matched the given type

midas/parser/midas.py

@@ -2,41 +2,47 @@ from typing import Optional
 
 from midas.ast.location import Location
 from midas.ast.midas import (
+    AliasStmt,
     BinaryExpr,
-    ComplexType,
+    CallExpr,
     ConstraintType,
     Expr,
     ExtendStmt,
+    FrameType,
+    FunctionType,
     GenericType,
     GetExpr,
     GroupingExpr,
     LiteralExpr,
     LogicalExpr,
+    MemberKind,
+    MemberStmt,
     NamedType,
-    OpStmt,
+    ParamSpec,
     PredicateStmt,
-    PropertyStmt,
     Stmt,
     Type,
+    TypeParam,
     TypeStmt,
     UnaryExpr,
-    UnionType,
     VariableExpr,
     WildcardExpr,
 )
-from midas.lexer.token import Token, TokenType
+from midas.lexer.token import KEYWORDS, Token, TokenType
 from midas.parser.base import Parser
 from midas.parser.errors import ParsingError
 
 
-class MidasParser(Parser):
+class MidasParser(Parser[list[Stmt]]):
     """A simple parser for midas type definitions"""
 
     SYNC_BOUNDARY: set[TokenType] = {
+        TokenType.ALIAS,
         TokenType.TYPE,
-        TokenType.OP,
         TokenType.EXTEND,
         TokenType.PREDICATE,
+        TokenType.PROP,
+        TokenType.FUNC,
     }
 
     def parse(self) -> list[Stmt]:
@@ -66,14 +72,16 @@ class MidasParser(Parser):
     def declaration(self) -> Optional[Stmt]:
         """Try and parse a declaration
 
-        Any parsing error is caught and None is returned
+        Any parsing error is caught and `None` is returned
 
         Returns:
-            Optional[Stmt]: the parsed Midas statement, or None if a ParsingError was raised
+            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.ALIAS):
+                return self.alias_declaration()
             if self.match(TokenType.EXTEND):
                 return self.extend_declaration()
             if self.match(TokenType.PREDICATE):
@@ -86,32 +94,21 @@ class MidasParser(Parser):
     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 type declaration creates a named subtype of a type expression.
+        It can have an optional template expression after its name, wrapped in brackets, to handle type parameters.
 
-        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
-            ...
-        }
-        ```
+        A type statement consists of:
+        - the `type` keyword
+        - a name (identifier)
+        - (optional) type parameters
+        - a body, a type expression (see :func:`type_expr`)
 
         Returns:
             TypeStmt: the parsed type declaration statement
         """
         keyword: Token = self.previous()
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected type name")
-        params: list[TypeStmt.Param] = []
-        if self.check(TokenType.LEFT_BRACKET):
-            params = self.type_stmt_params()
+        name: Token = self.consume_identifier("Expected type name")
+        params: list[TypeParam] = self.type_params()
 
         self.consume(TokenType.EQUAL, "Expected '=' before type definition")
 
@@ -124,24 +121,27 @@ class MidasParser(Parser):
             type=type,
         )
 
-    def type_stmt_params(self) -> list[TypeStmt.Param]:
-        """Parse a generic template expression
+    def type_params(self) -> list[TypeParam]:
+        """Parse a list of type parameters
 
-        A template is written `[TypeExpr]`
+        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:
-            TemplateExpr: the parsed template expression
+            list[TypeParam]: the list of type parameters, if any, or an empty list
         """
-        self.consume(TokenType.LEFT_BRACKET, "Missing '[' before template expression")
-        params: list[TypeStmt.Param] = []
+        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(TokenType.IDENTIFIER, "Expected type variable")
+            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(
-                TypeStmt.Param(
+                TypeParam(
                     location=name.location_to(self.previous()),
                     name=name,
                     bound=bound,
@@ -149,32 +149,56 @@ class MidasParser(Parser):
             )
             if not self.match(TokenType.COMMA):
                 break
-        self.consume(TokenType.RIGHT_BRACKET, "Missing ']' after template expression")
+        self.consume(TokenType.RIGHT_BRACKET, "Missing ']' after type parameters")
         return params
 
+    def alias_declaration(self) -> AliasStmt:
+        """Parse an alias declaration
+
+        An alias statement consists of:
+        - the `alias` keyword
+        - a name (identifier)
+        - a body, a type expression (see :func:`type_expr`)
+
+        Returns:
+            AliasStmt: the parsed alias declaration statement
+        """
+        keyword: Token = self.previous()
+        name: Token = self.consume_identifier("Expected alias name")
+
+        self.consume(TokenType.EQUAL, "Expected '=' before alias definition")
+
+        type: Type = self.type_expr()
+
+        return AliasStmt(
+            location=keyword.location_to(self.previous()),
+            name=name,
+            type=type,
+        )
+
     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
+        A type expression can either be a function type (see :func:`function`)
+        or a constraint type (see :func:`constraint_type`)
 
         Returns:
             TypeExpr: the parsed type expression
         """
-        return self.union_type()
-
-    def union_type(self) -> Type:
-        types: list[Type] = [self.constraint_type()]
-        while self.match(TokenType.PIPE):
-            types.append(self.constraint_type())
-        if len(types) == 1:
-            return types[0]
-        return UnionType(
-            location=Location.span(types[0].location, types[-1].location),
-            types=types,
-        )
+        if self.match(TokenType.FUNC):
+            return self.function()
+        return self.constraint_type()
 
     def constraint_type(self) -> Type:
+        """Parse a constraint type expression
+
+        A constraint type consists of a base type (see :func:`base_type`),
+        optionally followed by the `where` keyword and a constraint
+        expression (see :func:`constraint`)
+
+        Returns:
+            Type: the parsed constraint type expression
+        """
         type: Type = self.base_type()
         if self.match(TokenType.WHERE):
             constraint: Expr = self.constraint()
@@ -186,82 +210,141 @@ class MidasParser(Parser):
         return type
 
     def base_type(self) -> Type:
+        """Parse a base type expression
+
+        A base type is either a parenthesized type expression (see :func:`type_expr`)
+        or a generic type (see :func:`generic_type`)
+
+        Returns:
+            Type: the parsed base type expression
+        """
         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()
+        """Parse a generic type expression
+
+        A generic type consists of a named type (see :func:`named_type`),
+        optionally followed by type arguments in brackets.
+
+        The special `Frame` type accepts a frame schema instead of type
+        arguments (see :func:`frame_type`).
+
+        Returns:
+            Type: the parsed generic type
+        """
+        type: NamedType = self.named_type()
         if self.check(TokenType.LEFT_BRACKET):
-            params: list[Type] = self.type_params()
+            if type.name.lexeme == "Frame":
+                return self.frame_type()
+            args: list[Type] = self.type_args()
             return GenericType(
                 location=Location.span(type.location, self.previous().get_location()),
                 type=type,
-                params=params,
+                args=args,
             )
         return type
 
-    def type_params(self) -> list[Type]:
-        params: list[Type] = []
-        self.consume(TokenType.LEFT_BRACKET, "Missing '[' before generic parameters")
+    def type_args(self) -> list[Type]:
+        """Parse a list of type arguments
+
+        Type arguments are a comma-separated list of type expression wrapped in brackets.
+
+        Returns:
+            list[Type]: the list of type arguments, if any, or an empty list
+        """
+        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):
-            params.append(self.type_expr())
+            args.append(self.type_expr())
             if not self.match(TokenType.COMMA):
                 break
-        self.consume(TokenType.RIGHT_BRACKET, "Missing ']' after generic parameters")
-        return params
+        self.consume(TokenType.RIGHT_BRACKET, "Missing ']' after generic arguments")
+        return args
 
-    def named_type(self) -> Type:
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected type name")
+    def named_type(self) -> NamedType:
+        """Parse a named type expression
+
+        A named type is an identifier token
+
+        Returns:
+            NamedType: the parsed named type expression
+        """
+        name: Token = self.consume_identifier("Expected type name")
         return NamedType(
             location=name.get_location(),
             name=name,
         )
 
-    def complex_type(self) -> Type:
-        """Parse a type definition body
+    def frame_type(self) -> FrameType:
+        """Parse a frame type expression
 
-        A type definition body is a set of whitespace-separated
-        property statements enclosed in curly braces
+        A frame type consists of:
+        - the `Frame` identifier
+        - an opening bracket `[`
+        - a list of comma-separated column expression consisting of:
+            - a name (token)
+            - a colon `:`
+            - a type expression (see :func:`type_expr`)
+        - a closing bracket `]`
 
         Returns:
-            list[PropertyStmt]: the parsed type properties
+            FrameType: the parsed frame type
         """
-        left: Token = self.consume(
-            TokenType.LEFT_BRACE, "Expected '{' to start type body"
+        keyword: Token = self.previous()
+        self.consume(TokenType.LEFT_BRACKET, "Expected '[' to start frame schema")
+
+        columns: list[FrameType.Column] = []
+        while not self.check(TokenType.RIGHT_BRACKET) and not self.is_at_end():
+            name: Token = self.advance()
+            self.consume(TokenType.COLON, "Expected ':' between column name and type")
+            type: Type = self.type_expr()
+            columns.append(
+                FrameType.Column(
+                    location=name.location_to(self.previous()),
+                    name=name,
+                    type=type,
                 )
-        properties: list[PropertyStmt] = []
-        names: set[str] = set()
-        while not self.check(TokenType.RIGHT_BRACE) and not self.is_at_end():
-            prop: PropertyStmt = self.property_stmt()
-            if prop.name.lexeme in names:
-                raise self.error(prop.name, "Duplicate property")
-            names.add(prop.name.lexeme)
-            properties.append(prop)
-        right: Token = self.consume(TokenType.RIGHT_BRACE, "Unclosed type body")
-        return ComplexType(
-            location=left.location_to(right),
-            properties=properties,
+            )
+            if not self.match(TokenType.COMMA):
+                break
+
+        self.consume(TokenType.RIGHT_BRACKET, "Unclosed frame schema")
+
+        return FrameType(
+            location=keyword.location_to(self.previous()),
+            columns=columns,
         )
 
     def constraint(self) -> Expr:
-        """Parse a constraint
+        """Parse a constraint expression
 
-        A constraint is basically a logical predicate
+        A constraint is an expression (see :func:`expression`)
 
         Returns:
             Expr: the parsed constraint expression
         """
+        return self.expression()
+
+    def expression(self) -> Expr:
+        """Parse an expression
+
+        An expression consists of a logical AND expression (see :func:`and_`)
+
+        Returns:
+            Expr: the parsed expression
+        """
         return self.and_()
 
     def and_(self) -> Expr:
-        """Parse a logical AND expression or a simpler expression
+        """Parse a logical AND expression
+
+        An AND consists of one or more equality expressions (see :func:`equality`)
+        separated by logical AND operators (`&`)
 
         Returns:
             Expr: the parsed expression
@@ -277,7 +360,10 @@ class MidasParser(Parser):
         return expr
 
     def equality(self) -> Expr:
-        """Parse a logical equality expression or a simpler expression
+        """Parse an equality expression
+
+        An equality consists of one or more comparison expressions (see :func:`comparison`)
+        separated by equality operators (`==`, `!=`)
 
         Returns:
             Expr: the parsed expression
@@ -293,18 +379,59 @@ class MidasParser(Parser):
         return expr
 
     def comparison(self) -> Expr:
-        """Parse a logical comparison expression or a simpler expression
+        """Parse a comparison expression
+
+        A comparison consists of one or more term expressions (see :func:`term`)
+        separated by comparison operators (`<`, `<=`, `>`, `>=`)
 
         Returns:
             Expr: the parsed expression
         """
-        expr: Expr = self.unary()
+        expr: Expr = self.term()
         while self.match(
             TokenType.LESS,
             TokenType.LESS_EQUAL,
             TokenType.GREATER,
             TokenType.GREATER_EQUAL,
         ):
+            operator: Token = self.previous()
+            right: Expr = self.term()
+            location: Location = Location.span(expr.location, right.location)
+            expr = BinaryExpr(
+                location=location, left=expr, operator=operator, right=right
+            )
+        return expr
+
+    def term(self) -> Expr:
+        """Parse a term expression
+
+        A term consists of one or more factor expressions (see :func:`factor`)
+        separated by weak arithmetic operators (`+`, `-`)
+
+        Returns:
+            Expr: the parsed expression
+        """
+        expr: Expr = self.factor()
+        while self.match(TokenType.PLUS, TokenType.MINUS):
+            operator: Token = self.previous()
+            right: Expr = self.factor()
+            location: Location = Location.span(expr.location, right.location)
+            expr = BinaryExpr(
+                location=location, left=expr, operator=operator, right=right
+            )
+        return expr
+
+    def factor(self) -> Expr:
+        """Parse a factor expression
+
+        A factor consists of one or more unary expressions (see :func:`unary`)
+        separated by strong arithmetic operators (`*`, `/`)
+
+        Returns:
+            Expr: the parsed expression
+        """
+        expr: Expr = self.unary()
+        while self.match(TokenType.STAR, TokenType.SLASH):
             operator: Token = self.previous()
             right: Expr = self.unary()
             location: Location = Location.span(expr.location, right.location)
@@ -314,29 +441,116 @@ class MidasParser(Parser):
         return expr
 
     def unary(self) -> Expr:
-        """Parse a unary expression or a simpler expression
+        """Parse a unary expression
+
+        A unary consists of a call expression (see :func:`call`) optionally
+        preceded by zero or more unary operators (`+`, `-`, `!`)
 
         Returns:
             Expr: the parsed expression
         """
-        if self.match(TokenType.MINUS):
+        if self.match(TokenType.PLUS, TokenType.MINUS, TokenType.BANG):
             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()
+        return self.call()
+
+    def call(self) -> Expr:
+        """Parse a call expression
+
+        A call consists of a reference expression (see :func:`reference`)
+        optionally followed by zero or more argument groups.
+
+        Argument groups are parenthesize, comma-separated list of arguments (see :func:`finish_call`)
+
+        Returns:
+            Expr: the parsed expression
+        """
+        expr: Expr = self.reference()
+        while self.match(TokenType.LEFT_PAREN):
+            expr = self.finish_call(expr)
+        return expr
+
+    def finish_call(self, callee: Expr) -> Expr:
+        """Parse an argument group, i.e. the arguments of a call
+
+        Arguments are either passed positionally or by name (keyword argument).
+        All positional arguments must come before any keyword argument and
+        vice-versa. Arguments are separated by commas.
+
+        A positional argument simply consists of an expression (see :func:`expression`)
+
+        A keyword argument consists of and identifier, followed by the equal `=`
+        token and an expression (see :func:`expression`).
+
+        Args:
+            callee (Expr): the callee expression
+
+        Raises:
+            ParsingError: if a positional argument is passed after a keyword
+                argument or if a keyword argument's name is invalid (i.e. not
+                an identifier)
+
+        Returns:
+            Expr: the parsed call expression
+        """
+        pos_args: list[Expr] = []
+        kw_args: dict[str, Expr] = {}
+        keywords: bool = False
+        while not self.check(TokenType.RIGHT_PAREN):
+            if self.check_identifier() and self.check_next(TokenType.EQUAL):
+                keywords = True
+                keyword: Token = self.advance()
+                self.advance()
+                value: Expr = self.expression()
+                name: str = keyword.lexeme
+                if name in kw_args:
+                    self.error(
+                        self.peek(),
+                        f"Multiple values passed for '{name}', only the last occurrence will be used",
+                    )
+                kw_args[name] = value
+            else:
+                value = self.expression()
+                if self.check(TokenType.EQUAL):
+                    error_msg: str
+                    if keywords:
+                        error_msg = "Invalid keyword argument name"
+                    else:
+                        error_msg = (
+                            "Cannot pass positional arguments after a keyword argument"
+                        )
+                    raise self.error(self.peek(), error_msg)
+                pos_args.append(value)
+
+            if not self.match(TokenType.COMMA):
+                break
+
+        r_paren: Token = self.consume(
+            TokenType.RIGHT_PAREN, "Expected ')' after arguments."
+        )
+        return CallExpr(
+            location=Location.span(callee.location, r_paren.get_location()),
+            callee=callee,
+            arguments=pos_args,
+            keywords=kw_args,
+        )
 
     def reference(self) -> Expr:
-        """Parse an attribute access expression or a simpler expression
+        """Parse a reference expression
+
+        A reference consists of a primary expression (see :func:`primary`)
+        optionally followed by zero or more attribute accesses.
+
+        An attribute access consists of a dot `.` token followed by an identifier
 
         Returns:
             Expr: the parsed expression
         """
         expr: Expr = self.primary()
         while self.match(TokenType.DOT):
-            name: Token = self.consume(
-                TokenType.IDENTIFIER, "Expected property name after '.'"
-            )
+            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
@@ -344,7 +558,12 @@ class MidasParser(Parser):
     def primary(self) -> Expr:
         """Parse a primary expression
 
-        This includes literals (booleans, numbers, etc.), wildcards, identifiers and grouped expressions
+        This includes literals (booleans, numbers, etc.), wildcards, identifiers
+        and grouped expressions
+
+        Raises:
+            ParsingError: if a primary expressions cannot be parsed from the
+                following tokens
 
         Returns:
             Expr: the parsed expression
@@ -360,100 +579,253 @@ class MidasParser(Parser):
         if self.match(TokenType.NUMBER):
             return LiteralExpr(location=token.get_location(), value=token.value)
 
-        if self.match(TokenType.IDENTIFIER):
+        if self.match(TokenType.STRING):
+            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()
+            expr: Expr = self.expression()
             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 property_stmt(self) -> PropertyStmt:
-        """Parse a property statement
+    def consume_identifier(self, message: str = "Expected identifier") -> Token:
+        """Consume the current token if it is a valid identifier or raise an error (see :func:`check_identifier`)
 
-        A type property statement is written `name: Type` or `name: Type where Condition`
+        If the current token is not a valid identifier, an error is raised
+        with the provided message
+
+        Args:
+            message (str, optional): the error message. Defaults to "Expected identifier".
+
+        Raises:
+            ParsingError: if the current token is not a valid identifier
 
         Returns:
-            PropertyStmt: the parsed property statement
+            Token: the current token which is a valid identifier
         """
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected property name")
-        self.consume(TokenType.COLON, "Expected ':' after property name")
+        if not self.match_identifier():
+            raise self.error(self.peek(), message)
+        return self.previous()
+
+    def match_identifier(self) -> bool:
+        """Consume the next token if it is a valid identifier (see :func:`check_identifier`)
+
+        Returns:
+            bool: whether a token was matched and consumed
+        """
+        return self.match(TokenType.IDENTIFIER, *KEYWORDS.values())
+
+    def check_identifier(self) -> bool:
+        """Check whether the current token is a valid identifier
+
+        A valid identifier is either an identifier token or a keyword token.
+        This function always returns False if the parser is at the EOF token
+
+        Returns:
+            bool: True if the current token is a valid identifier and not EOF
+        """
+        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 member statement is written consists of:
+        - the `prop` (for a property) or `def` (for a method) keyword
+        - an name (identifier)
+        - a colon `:`
+        - a type expression (see :func:`type_expr`)
+
+        Raises:
+            ParsingError: if the first token is neither `prop` nor `def`
+
+        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 PropertyStmt(
+        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 }`
+        An extension statement consists of:
+        - the `extend` keyword
+        - a type name (identifier)
+        - (optional) type parameters (see :func:`type_params`)
+        - an opening brace `{`
+        - zero or more member statements (see :func:`member_stmt`)
+        - a closing brace `}`
 
         Returns:
             ExtendStmt: the parsed extension statement
         """
         keyword: Token = self.previous()
-        type: Type = self.type_expr()
+        name: Token = self.consume_identifier("Expected type name")
+        params: list[TypeParam] = self.type_params()
+
         self.consume(TokenType.LEFT_BRACE, "Expected '{' to start extend body")
-        operations: list[OpStmt] = []
+        members: list[MemberStmt] = []
         while not self.is_at_end() and not self.check(TokenType.RIGHT_BRACE):
-            operations.append(self.op_declaration())
+            members.append(self.member_stmt())
         self.consume(TokenType.RIGHT_BRACE, "Unclosed extend body")
         location: Location = keyword.location_to(self.previous())
-        return ExtendStmt(location=location, type=type, operations=operations)
-
-    def op_declaration(self) -> OpStmt:
-        """Parse an operation definition
-
-        An operation is written `op name(Type) -> Type`
-
-        Returns:
-            OpStmt: the parsed operation statement
-        """
-        keyword: Token = self.consume(TokenType.OP, "Expected 'op' keyword")
-
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected operation name")
-        self.consume(TokenType.LEFT_PAREN, "Expected '(' before operand type")
-        operand: Type = self.type_expr()
-        self.consume(TokenType.RIGHT_PAREN, "Expected ')' after operand type")
-
-        self.consume(TokenType.ARROW, "Expected '->' before result type")
-        result: Type = self.type_expr()
-
-        return OpStmt(
-            location=keyword.location_to(self.previous()),
+        return ExtendStmt(
+            location=location,
             name=name,
-            operand=operand,
-            result=result,
+            params=params,
+            members=members,
         )
 
     def predicate_declaration(self) -> PredicateStmt:
         """Parse a predicate declaration
 
-        A predicate is written `predicate Name(subject: Type) = constraint_expression`
+        A predicate statement consists of:
+        - the `predicate` keyword
+        - a name (identifier)
+        - (optional) zero or more parameter specs (see :func:`function_params`)
+        - an equal sign `=`
+        - a body, a constraint expression (see :func:`constraint`)
 
         Returns:
             PredicateStmt: the parsed predicate declaration statement
         """
         keyword: Token = self.previous()
-        name: Token = self.consume(TokenType.IDENTIFIER, "Expected predicate name")
-        self.consume(TokenType.LEFT_PAREN, "Expected '(' before predicate subject")
-        subject: Token = self.consume(TokenType.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")
+
+        name: Token = self.consume_identifier("Expected predicate name")
+
+        params: list[ParamSpec] = []
+        while self.check(TokenType.LEFT_PAREN):
+            params.append(self.function_params())
+
         self.consume(TokenType.EQUAL, "Expected '=' after predicate subject")
-        condition: Expr = self.constraint()
+        body: Expr = self.constraint()
         return PredicateStmt(
             location=keyword.location_to(self.previous()),
             name=name,
-            subject=subject,
-            type=type,
-            condition=condition,
+            params=params,
+            body=body,
         )
+
+    def function(self) -> FunctionType:
+        """Parse a function type expression
+
+        A function consists of:
+        - the `fn` keyword
+        - a parameter spec (see :func:`function_params`)
+        - the arrow keyword `->`
+        - a result type expression (see :func:`type_expr`)
+
+        Returns:
+            FunctionType: the parsed function type expression
+        """
+        params: ParamSpec = self.function_params()
+
+        self.consume(TokenType.ARROW, "Expected '->' before result type")
+        result: Type = self.type_expr()
+
+        return FunctionType(
+            location=params.l_paren.location_to(self.previous()),
+            params=params,
+            returns=result,
+        )
+
+    def function_params(self) -> ParamSpec:
+        """Parse a parameter spec
+
+        A parameter spec consists of zero or more comma-separated parameters,
+        wrapped in parentheses.
+
+        Like in Python, it can contain positional-only, mixed and keyword-only
+        parameters (separated by `/` and `*`).
+
+        Each parameter has a type (see :func:`type_expr`),
+        preceded by a name (identifier) and a colon `:` (not required for
+        positional-only parameters).
+
+        Returns:
+            ParamSpec: the parsed parameter spec
+        """
+        l_paren: Token = self.consume(
+            TokenType.LEFT_PAREN, "Expected '(' before function parameters"
+        )
+        pos: list[FunctionType.Parameter] = []
+        mixed: list[FunctionType.Parameter] = []
+        kw: list[FunctionType.Parameter] = []
+
+        mixed_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 = mixed
+                    mixed = []
+                    mixed_first_tokens = []
+                    section = 1
+                case 0 | 1 if self.match(TokenType.STAR):
+                    section = 2
+                case _:
+                    # Record first token of mixed parameters for errors if unnamed
+                    if section != 2:
+                        mixed_first_tokens.append(self.peek())
+
+                    name: Optional[Token] = None
+                    if section == 2:
+                        name = self.consume_identifier(
+                            "Expected keyword parameter name"
+                        )
+                        self.consume(
+                            TokenType.COLON, "Expected ':' after parameter 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)
+                    param = FunctionType.Parameter(
+                        location=None,
+                        name=name,
+                        type=type,
+                        required=not optional,
+                    )
+                    if section == 2:
+                        kw.append(param)
+                    else:
+                        mixed.append(param)
+
+            if not self.match(TokenType.COMMA):
+                break
+
+        for param, token in zip(mixed, mixed_first_tokens):
+            if param.name is None:
+                # Not raised because we can keep parsing
+                self.error(token, "Unnamed mixed parameter")
+
+        self.consume(TokenType.RIGHT_PAREN, "Expected ')' after function parameters")
+        return ParamSpec(l_paren=l_paren, pos=pos, mixed=mixed, kw=kw)

midas/parser/python.py

@@ -10,19 +10,31 @@ from midas.ast.python import (
     CastExpr,
     CompareExpr,
     ConstraintType,
+    DictExpr,
     Expr,
     ExpressionStmt,
+    ForStmt,
     FrameColumn,
     FrameType,
+    FromImportStmt,
     Function,
     GetExpr,
     IfStmt,
+    ImportAlias,
+    ImportStmt,
+    ListExpr,
     LiteralExpr,
     LogicalExpr,
     MidasType,
+    ParamSpec,
+    RawExpr,
+    RawStmt,
     ReturnStmt,
+    SliceExpr,
     Stmt,
+    SubscriptExpr,
     TernaryExpr,
+    TupleExpr,
     TypeAssign,
     UnaryExpr,
     VariableExpr,
@@ -41,7 +53,10 @@ class UnsupportedSyntaxError(Exception):
 
 
 class PythonParser:
+    """A parser to convert raw Python `ast` nodes in custom IR nodes"""
+
     CAST_FUNCTION = "cast"
+    UNSAFE_CAST_FUNCTION = "unsafe_cast"
 
     def parse_module(self, node: ast.Module) -> list[Stmt]:
         statements: list[Stmt] = []
@@ -87,9 +102,39 @@ class PythonParser:
             case ast.If():
                 return self.parse_if(node)
 
+            case ast.Pass():
+                return None
+
+            case ast.For(orelse=[]):
+                return self.parse_for(node)
+
+            case ast.Import(names=imports):
+                return ImportStmt(
+                    location=location,
+                    imports=self._parse_imports(imports),
+                )
+
+            case ast.ImportFrom(module=module, names=imports, level=level):
+                return FromImportStmt(
+                    location=location,
+                    module=module,
+                    imports=self._parse_imports(imports),
+                    level=level,
+                )
+
             case _:
                 print(f"Unsupported statement: {ast.unparse(node)}")
-                return None
+                return RawStmt(location=location, stmt=node)
+
+    def _parse_imports(self, imports: list[ast.alias]) -> list[ImportAlias]:
+        return [
+            ImportAlias(
+                location=Location.from_ast(import_),
+                name=import_.name,
+                alias=import_.asname,
+            )
+            for import_ in imports
+        ]
 
     def parse_annotation_assign(self, node: ast.AnnAssign) -> list[Stmt]:
         statements: list[Stmt] = []
@@ -176,32 +221,31 @@ class PythonParser:
             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)
@@ -210,54 +254,58 @@ class PythonParser:
                     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
-                    ),
+                    params=self._parse_param_spec(args),
                     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(
+    def _parse_param_spec(self, args: ast.arguments) -> ParamSpec:
+        def parse_params(
+            args_list: list[ast.arg], defaults: list[Optional[Expr]]
+        ) -> list[Function.Parameter]:
+            return [
+                self._parse_function_parameter(arg, default)
+                for arg, default in zip(args_list, defaults)
+            ]
+
+        defaults: list[ast.expr] = args.defaults
+        parsed_defaults: list[Optional[Expr]] = [
+            self.parse_expr(default) for default in defaults
+        ]
+        n_pos: int = len(args.posonlyargs)
+        n_mixed: int = len(args.args)
+        n_all_pos = n_pos + n_mixed
+        parsed_defaults = [
+            None,
+        ] * (n_all_pos - len(defaults)) + parsed_defaults
+
+        pos_defaults: list[Optional[Expr]] = parsed_defaults[:n_pos]
+        mixed_defaults: list[Optional[Expr]] = parsed_defaults[n_pos:]
+        kw_defaults: list[Optional[Expr]] = [
+            self.parse_expr(default) if default is not None else None
+            for default in args.kw_defaults
+        ]
+
+        return ParamSpec(
+            pos=parse_params(args.posonlyargs, pos_defaults),
+            mixed=parse_params(args.args, mixed_defaults),
+            kw=parse_params(args.kwonlyargs, kw_defaults),
+        )
+
+    def _parse_function_parameter(
         self, arg: ast.arg, default: Optional[Expr]
-    ) -> Function.Argument:
+    ) -> Function.Parameter:
         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(
+        return Function.Parameter(
             location=loc,
             name=name,
             type=type,
@@ -270,26 +318,28 @@ class PythonParser:
             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):
+            case ast.Subscript(value=ast.Name(id=name), slice=arg):
+                args: tuple[MidasType, ...] = (
+                    tuple(self._parse_type(a) for a in arg.elts)
+                    if isinstance(arg, ast.Tuple)
+                    else (self._parse_type(arg),)
+                )
                 return BaseType(
                     location=loc,
                     base=name,
-                    param=self._parse_type(param),
+                    args=args,
                 )
 
             case ast.Name(id=name):
                 return BaseType(
                     location=loc,
                     base=name,
-                    param=None,
+                    args=(),
                 )
 
             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(
@@ -311,6 +361,13 @@ class PythonParser:
                             constraint=right_expr,
                         )
 
+            case ast.Constant(value=None):
+                return BaseType(
+                    location=loc,
+                    base="None",
+                    args=(),
+                )
+
             case _:
                 raise UnsupportedSyntaxError(type_expr)
 
@@ -323,7 +380,7 @@ class PythonParser:
                 for col in cols:
                     columns.append(self._parse_frame_column(col))
 
-            case ast.Slice() | ast.Name():
+            case ast.Slice() | ast.Name() | ast.Subscript():
                 columns.append(self._parse_frame_column(schema))
 
             case _:
@@ -334,7 +391,7 @@ class PythonParser:
     def _parse_frame_column(self, column: ast.expr) -> FrameColumn:
         loc: Location = Location.from_ast(column)
         match column:
-            case ast.Name():
+            case ast.Name() | ast.Subscript():
                 return FrameColumn(
                     location=loc,
                     name=None,
@@ -387,6 +444,9 @@ class PythonParser:
             case ast.Call(func=ast.Name(id=self.CAST_FUNCTION)):
                 return self.parse_cast(node)
 
+            case ast.Call(func=ast.Name(id=self.UNSAFE_CAST_FUNCTION)):
+                return self.parse_cast(node)
+
             case ast.Call():
                 return self.parse_call(node)
 
@@ -406,8 +466,46 @@ class PythonParser:
             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 ast.Tuple(elts=items):
+                return TupleExpr(
+                    location=location,
+                    items=tuple(self.parse_expr(item) for item in items),
+                )
+
             case _:
-                raise UnsupportedSyntaxError(node)
+                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
@@ -459,16 +557,19 @@ class PythonParser:
         return expr
 
     def parse_cast(self, node: ast.Call) -> CastExpr:
+        assert isinstance(node.func, ast.Name)
+        func: str = node.func.id
         match node:
             case ast.Call(args=[type, expr], keywords=[]):
                 return CastExpr(
                     location=Location.from_ast(node),
                     type=self._parse_type(type),
                     expr=self.parse_expr(expr),
+                    unsafe=func == self.UNSAFE_CAST_FUNCTION,
                 )
             case _:
                 raise InvalidSyntaxError(
-                    f"Invalid call to {self.CAST_FUNCTION}, expected type and expression"
+                    f"Invalid call to {func}, expected type and expression"
                 )
 
     def parse_call(self, node: ast.Call) -> CallExpr:

midas/resolver/builtin.py

@@ -1,71 +0,0 @@
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-from midas.checker.types import BaseType, Type, UnitType
-
-if TYPE_CHECKING:
-    from midas.resolver.midas import MidasResolver
-
-
-def op(ctx: MidasResolver, t1: Type, operator: str, t2: Type, t3: Type):
-    ctx.define_operation(
-        left=t1,
-        operator=operator,
-        right=t2,
-        result=t3,
-    )
-
-def basic_op(ctx: MidasResolver, type: Type, op: str):
-    ctx.define_operation(
-        left=type,
-        operator=op,
-        right=type,
-        result=type,
-    )
-
-
-def define_builtins(ctx: MidasResolver):
-    """Define builtin types and operations"""
-    unit = ctx.define_type("None", UnitType())
-    bool = ctx.define_type("bool", BaseType(name="bool"))
-    int = ctx.define_type("int", BaseType(name="int"))
-    float = ctx.define_type("float", BaseType(name="float"))
-    str = ctx.define_type("str", BaseType(name="str"))
-
-    basic_op(ctx, int, "__add__")  # int + int = int
-    basic_op(ctx, int, "__sub__")  # int - int = int
-    basic_op(ctx, int, "__mul__")  # int * int = int
-    basic_op(ctx, int, "__pow__")  # int ** int = int
-    basic_op(ctx, int, "__mod__")  # int % int = int
-    basic_op(ctx, int, "__and__")  # int & int = int
-    basic_op(ctx, int, "__or__")  # int | int = int
-    basic_op(ctx, int, "__xor__")  # int ^ int = int
-    op(ctx, int, "__lt__", int, bool)  # int < int = bool
-    op(ctx, int, "__gt__", int, bool)  # int > int = bool
-    op(ctx, int, "__le__", int, bool)  # int <= int = bool
-    op(ctx, int, "__ge__", int, bool)  # int >= int = bool
-    op(ctx, int, "__eq__", int, bool)  # int == int = bool
-    basic_op(ctx, float, "__add__")  # float + float = float
-    basic_op(ctx, float, "__sub__")  # float - float = float
-    basic_op(ctx, float, "__mul__")  # float * float = float
-    basic_op(ctx, float, "__truediv__")  # float / float = float
-    op(ctx, float, "__lt__", float, bool)  # float < float = bool
-    op(ctx, float, "__gt__", float, bool)  # float > float = bool
-    op(ctx, float, "__le__", float, bool)  # float <= float = bool
-    op(ctx, float, "__ge__", float, bool)  # float >= float = bool
-    op(ctx, float, "__eq__", float, bool)  # float == float = bool
-    basic_op(ctx, str, "__add__")  # str + str = str
-    op(ctx, str, "__eq__", str, bool)  # str == str = bool
-
-    op(ctx, int, "__lt__", float, bool)  # int < float = bool
-    op(ctx, int, "__gt__", float, bool)  # int > float = bool
-    op(ctx, int, "__le__", float, bool)  # int <= float = bool
-    op(ctx, int, "__ge__", float, bool)  # int >= float = bool
-    op(ctx, int, "__eq__", float, bool)  # int == float = bool
-
-    op(ctx, float, "__lt__", int, bool)  # float < int = bool
-    op(ctx, float, "__gt__", int, bool)  # float > int = bool
-    op(ctx, float, "__le__", int, bool)  # float <= int = bool
-    op(ctx, float, "__ge__", int, bool)  # float >= int = bool
-    op(ctx, float, "__eq__", int, bool)  # float == int = bool

midas/resolver/midas.py

@@ -1,166 +0,0 @@
-from typing import Optional
-
-import midas.ast.midas as m
-from midas.checker.types import (
-    Type,
-    UnionType,
-    UnknownType,
-)
-from midas.resolver.builtin import define_builtins
-
-
-class MidasResolver(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) -> None:
-        self._types: dict[str, Type] = {}
-        self._operations: dict[tuple[Type, str, Type], Type] = {}
-
-        define_builtins(self)
-
-    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
-        """
-        type: Optional[Type] = self._types.get(name)
-        if type is None:
-            raise NameError(f"Undefined type {name}")
-        return type
-
-    def get_operation_result(
-        self, left: Type, operator: str, right: Type
-    ) -> Optional[Type]:
-        """Get the resulting type of an operation
-
-        Args:
-            left (Type): the type of the left operand
-            operator (str): the operation name
-            right (Type): the type of the right operand
-
-        Returns:
-            Optional[Type]: the result type, or None if no matching operation was found
-        """
-        operation: tuple[Type, str, Type] = (left, operator, right)
-        result: Optional[Type] = self._operations.get(operation)
-        return result
-
-    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_operation(self, left: Type, operator: str, right: Type, result: Type):
-        """Define an operation in the registry
-
-        Args:
-            left (Type): the type of the left operand
-            operator (str): the operation name
-            right (Type): the type of the right operand
-            result (Type): the result type
-
-        Raises:
-            ValueError: if an operation is already defined with these operands and name
-        """
-        operation: tuple[Type, str, Type] = (left, operator, right)
-        if operation in self._operations:
-            raise ValueError(
-                f"Operation {operator} already defined between {left} and {right}"
-            )
-        self._operations[operation] = result
-
-    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:
-        type: Type = stmt.type.accept(self)
-        for param in stmt.params:
-            if param.bound is not None:
-                param.bound.accept(self)
-        self.define_type(stmt.name.lexeme, type)
-
-    def visit_property_stmt(self, stmt: m.PropertyStmt) -> None: ...
-
-    def visit_extend_stmt(self, stmt: m.ExtendStmt) -> None:
-        base: Type = stmt.type.accept(self)
-        for op in stmt.operations:
-            right: Type = op.operand.accept(self)
-            result: Type = op.result.accept(self)
-            self.define_operation(
-                left=base,
-                operator=op.name.lexeme,
-                right=right,
-                result=result,
-            )
-
-    def visit_op_stmt(self, stmt: m.OpStmt) -> None: ...
-
-    def visit_predicate_stmt(self, stmt: m.PredicateStmt) -> None: ...
-
-    def visit_logical_expr(self, expr: m.LogicalExpr) -> None: ...
-
-    def visit_binary_expr(self, expr: m.BinaryExpr) -> None: ...
-
-    def visit_unary_expr(self, expr: m.UnaryExpr) -> None: ...
-
-    def visit_get_expr(self, expr: m.GetExpr) -> None: ...
-
-    def visit_variable_expr(self, expr: m.VariableExpr) -> None: ...
-
-    def visit_grouping_expr(self, expr: m.GroupingExpr) -> None:
-        return 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) -> Type:
-        return self.get_type(type.name.lexeme)
-
-    def visit_generic_type(self, type: m.GenericType) -> Type:
-        type_: Type = type.type.accept(self)
-        params: list[Type] = [param.accept(self) for param in type.params]
-        # TODO
-        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_union_type(self, type: m.UnionType) -> Type:
-        types: list[Type] = [type_.accept(self) for type_ in type.types]
-        return UnionType(alternatives=types)
-
-    def visit_complex_type(self, type: m.ComplexType) -> Type:
-        for prop in type.properties:
-            prop.accept(self)
-        # TODO
-        return UnknownType()

midas/typing.py

@@ -0,0 +1,54 @@
+from typing import Generic, TypeVar
+from typing import cast as typing_cast
+
+from pandas import DataFrame, Series
+
+cast = typing_cast
+"""### Midas documentation
+Cast a value to a type.
+
+- **Compile-time**: tells the type checker that the return value has the designated type.
+- **Run-time**: generates assertions to ensure the value can be interpreted as the given type.
+
+---
+<br>
+<br>
+<br>
+
+_**Internal Python documentation**_
+"""
+
+
+unsafe_cast = typing_cast
+"""### Midas documentation
+Cast a value to a type.
+
+- **Compile-time**: tells the type checker that the return value has the designated type.
+- **Run-time**: -
+
+This operation is unsound, use at your own risk!
+
+---
+<br>
+<br>
+<br>
+
+_**Internal Python documentation**_
+"""
+
+
+T = TypeVar("T")
+
+
+class Frame(DataFrame, Generic[T]):
+    """A `Frame` is the abstract type implemented by `DataFrame`
+
+    A frame contains any number of named columns (see :class:`Column`)
+    """
+
+
+class Column(Series, Generic[T]):
+    """A `Column` is the abstract type implemented by `Series`
+
+    A column contains a any number of values of the same type
+    """

midas/utils.py

@@ -1,5 +1,10 @@
+from dataclasses import dataclass
 from typing import Any, Callable, Optional
 
+import midas.ast.python as p
+from midas.checker.types import Type
+from midas.generator.collector import AssertionCollector
+
 AllowRepeat = Callable[[object], bool]
 
 
@@ -52,3 +57,11 @@ class UniversalJSONDumper:
                 }
             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]]
+    evaluated_casts: list[p.CastExpr]
+    assertions: AssertionCollector

pyproject.toml

@@ -8,7 +8,11 @@ authors = [
     { name = "Louis Heredero", email = "louis.heredero@students.hevs.ch" },
 ]
 classifiers = ["Programming Language :: Python :: 3"]
-dependencies = ["click>=8.4.1"]
+dependencies = [
+    "black>=26.5.1",
+    "click>=8.4.1",
+    "watchdog>=6.0.0",
+]
 
 [project.urls]
 Homepage = "https://git.kbk28.ch/HEL/midas"

syntax/annotations.ebnf

@@ -1,20 +1,8 @@
-identifier ::= '[a-zA-Z][a-zA-Z_]*'
+Identifier ::= '[a-zA-Z][a-zA-Z_]*'
 
-integer ::= '\d+'
-number ::= integer ["." integer]
-boolean ::= "False" | "True"
-none ::= "None"
+TypeArgs ::= "[" (Type ("," Type)*)? "]"
 
-value ::= number | boolean | none
-lambda-value ::= "_" | value
-lambda-operator ::= ">" | "<" | ">=" | "<=" | "==" | "!="
-lambda ::= lambda-value lambda-operator lambda-value
+FrameColumn ::= ((Identifier | "_") ":")? Type
+FrameSchema ::= "[" (FrameColumn ("," FrameColumn)*)? "]"
 
-constraint ::= identifier | "(" lambda ")"
-base-type ::= identifier
-type ::= base-type { "+" constraint }
-
-column-type ::= type | "_"
-column-def ::= [ identifier ":" ] column-type
-
-frame-def ::= column-def { "," column-def }
+Type ::= "Frame" FrameSchema | Identifier TypeArgs?

syntax/annotations.typ

@@ -1,64 +1,72 @@
-#import "@preview/fervojo:0.1.1": render
+#import "@preview/fervojo:0.1.1": default-css, render
 
-#let value = ```
-{[`value` <
-  [`number` 'digit' * ! <!, ["." 'digit' * !]>],
-  [`boolean` <"False", "True">],
-  [`none` "None"]
+#let extra-css = ```css
+svg.railroad .terminal rect {
+  fill: #F7DCD4;
+}
+```
+#let css = default-css() + bytes(extra-css.text)
+
+#let type-args = ```
+{[`type-args` "[" <!, 'type'*","> "]"]}
+```
+
+#let frame-schema = ```
+{[`frame-schema` "[" <!, [[<'identifier', "_"> ":"]? 'type']*","> "]"]}
+```
+
+#let type = ```
+{[`type` <
+  ["Frame" 'frame-schema'],
+  ['identifier' <!, 'type-args'>]
 >]}
 ```
 
-#let constraint = ```
-{[`constraint` <"_", 'value'> <">", "<", ">=", "<=", "==", "!="> <"_", 'value'>]}
-```
-
-#let type-with-constraints = ```
-{[`type-with-constraints` 'identifier' <!, ["+" "(" 'constraint' ")"] * !>]}
-```
-
-#let column-def = ```
-{[`column-def` <!, ['identifier' ":"]> <"_", 'type-with-constraints'>]}
-```
-
-#let frame-def = ```
-{[`frame-def` 'column-def' * ","]}
-```
-
-#let annotation = ```
-{[`annotation` 'identifier' <!, ["[" 'frame-def' "]"]>]}
-```
-
 #let rules = (
-  value,
-  constraint,
-  type-with-constraints,
-  column-def,
-  frame-def,
-  annotation,
+  type-args: type-args,
+  frame-schema: frame-schema,
+  type: type,
+)
+
+#let inline = (
+  "type-args",
+  "frame-schema",
 )
 
 #set text(font: "Source Sans 3")
 
-= Type annotation syntax
+#title[Supported Python annotation syntax]
 
-#for rule in rules {
-  render(rule)
-}
+= Outline
 
-/*
-#let by-name = (
-  annotation: annotation,
-  frame-def: frame-def,
-  column-def: column-def,
-  type-with-constraints: type-with-constraints,
-  constraint: constraint,
-  value: value,
+#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)
@@ -66,9 +74,16 @@
   return new-rule
 }
 
-#let combined = raw(substitute(annotation.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)
+  }
+}

syntax/midas.ebnf

@@ -4,32 +4,85 @@ Identifier ::= [a-zA-Z_] [a-zA-Z_0-9]*
 Integer ::= '\d+'
 Number ::= "-"? Integer ("." Integer)?
 Boolean ::= "False" | "True"
+String ::= '(".*?")|(\'.*?\')'
 None ::= "None"
 
-Value ::= Number | Boolean | None
+Literal ::= Number | Boolean | String | None
 
+UnaryOp ::= "+" | "-" | "!"
+FactorOp ::= "*" | "/"
+TermOp ::= "+" | "-"
 ComparisonOp ::= ">" | "<" | ">=" | "<="
 EqualityOp ::= "==" | "!="
 
-Grouping ::= "(" Constraint ")"
-Primary ::= "_" | Value | Identifier | Grouping
+PosArg ::= Expression
+KwArg ::= Identifier "=" Expression
+
+PosArgs ::= PosArg ("," PosArg)*
+KwArgs ::= KwArg ("," KwArg)*
+Args ::= (
+    PosArgs
+    | KwArgs
+    | PosArgs "," KwArgs
+)
+
+Grouping ::= "(" Expression ")"
+Primary ::= "_" | Literal | Identifier | Grouping
 Reference ::= Primary ("." Identifier)*
-Unary ::= "-"? Unary | Reference
-Comparison ::= Unary (ComparisonOp Unary)*
+CallArgs ::= "(" Args ")" 
+Call ::= Reference CallArgs*
+Unary ::= UnaryOp Unary | Call
+Factor ::= Unary (FactorOp Unary)*
+Term ::= Factor (TermOp Factor)*
+Comparison ::= Term (ComparisonOp Term)*
 Equality ::= Comparison (EqualityOp Comparison)*
-Constraint ::= Equality ("&" Equality)*
+Expression ::= Equality ("&" Equality)*
+Constraint ::= Expression
 
-SimpleType ::= Identifier "?"?
-Template ::= "[" Type "]"
-Type ::= Identifier Template? "?"?
+TemplateParam ::= Identifier ("<:" Type)?
+Template ::= "[" (TemplateParam ("," TemplateParam)*)? "]"
 
-TypeProperty ::= Identifier ":" Type ("where" Constraints)?
-ComplexTypeBody ::= "{" TypeProperty* "}"
-OpDefinition ::= "op" Identifier "(" Type ")" "->" Type
-ExtendBody ::= "{" OpDefinition* "}"
+ParamType ::= Type "?"?
+PosParam ::= (Identifier ":")? ParamType
+KwParam ::= Identifier ":" ParamType
 
-TypeStatement ::= "type" Identifier Template? ("(" Type ")" ("where" Constraint)? | ComplexTypeBody)
+PosParams ::= (
+    (PosParam ("," PosParam)* ("," "/")?)
+    | "/"
+)
+MixedParams ::= KwParam ("," KwParam)
+KwParams ::= (
+    (("*", ",")? KwParam ("," KwParam)*)
+    | "*"
+)
+Params ::= (
+    PosParams
+    | MixedParams
+    | KwParams
+    | (PosParams "," MixedParams)
+    | (PosParams "," KwParams)
+    | (MixedParams "," KwParams)
+    | (PosParams "," MixedParams "," KwParams)
+)
+ParamSpec ::= "(" Params? ")"
+
+NamedType ::= Identifier
+TypeArgs ::= "[" (Type ("," Type)*)? "]"
+FrameColumn ::= TOKEN ":" Type
+FrameSchema ::= "[" (FrameColumn ("," FrameColumn)*)? "]"
+GenericType ::= "Frame" FrameSchema | NamedType TypeArgs?
+GroupedType ::= "(" Type ")"
+BaseType ::= GroupedType | GenericType
+ConstraintType ::= BaseType ("where" Constraint)?
+FuncType ::= "fn" ParamSpec "->" Type
+Type ::= ConstraintType
+
+MemberStatement ::= ("prop" | "def") Identifier ":" Type
+ExtendBody ::= "{" MemberStatement* "}"
+
+AliasStatement ::= "alias" Identifier "=" Type
+TypeStatement ::= "type" Identifier Template? "=" Type
 ExtendStatement ::= "extend" Type ExtendBody
-PredicateStatement ::= "predicate" Identifier "(" Identifier ":" Type ")" "=" Constraint
+PredicateStatement ::= "predicate" Identifier ParamSpec* "=" Constraint
 
-Statement ::= TypeStatement | ExtendStatement | PredicateStatement
+Statement ::= AliasStatement | TypeStatement | ExtendStatement | PredicateStatement

syntax/midas.typ

@@ -7,98 +7,174 @@ svg.railroad .terminal rect {
 ```
 #let css = default-css() + bytes(extra-css.text)
 
-#let value = ```
-{[`value` <
+#let literal = ```
+{[`literal` <
   [`number` 'digit' * ! <!, ["." 'digit' * !]>],
   [`boolean` <"False", "True">],
+  [`string` <["\"" 'char'*! "\""], ["'" 'char'*! "'"]>],
   [`none` "None"]
 >]}
 ```
 
 #let grouping = ```
-{[`grouping` "(" 'constraint' ")"]}
+{[`grouping` "(" 'expression' ")"]}
 ```
 
 #let primary = ```
-{[`primary` <"_", 'value', 'identifier', 'grouping'>]}
+{[`primary` <"_", 'literal', 'identifier', 'grouping'>]}
 ```
 
 #let reference = ```
 {[`reference` 'primary' <!, ["." 'identifier']*!>]}
 ```
 
+#let call-args = ```
+{[`call-args` "(" <!, <'expression', ['identifier' "=" 'expression']>*","#`Same rules as Python`> ")"]}
+```
+
+#let call = ```
+{[`call` 'reference' <!, 'call-args'*!>]}
+```
+
 #let unary = ```
-{[`unary` <[<!, "-"> 'unary'], 'reference'>]}
+{[`unary` <[<"+", "-", "!"> 'unary'], 'call'>]}
+```
+
+#let factor = ```
+{[`factor` 'unary'*<"*", "/">]}
+```
+
+#let term = ```
+{[`term` 'factor'*<"+", "-">]}
 ```
 
 #let comparison = ```
-{[`comparison` 'unary'*<">", "<", ">=", "<=">]}
+{[`comparison` 'term'*<">", "<", ">=", "<=">]}
 ```
 
 #let equality = ```
 {[`equality` 'comparison'*<"==", "!=">]}
 ```
 
-#let constraint = ```
-{[`constraint` 'equality'*"&"]}
+#let expression = ```
+{[`expression` 'equality'*"&"]}
 ```
 
-#let simple-type = ```
-{[`simple-type` 'identifier' <!, "?">]}
+#let constraint = ```
+{[`constraint` 'expression']}
+```
+
+#let template-param = ```
+{[`template-param` 'identifier' <!, ["<:" 'type']>]}
 ```
 
 #let template = ```
-{[`template` "[" 'type' "]"]}
+{[`template` "[" <!, 'template-param'*","> "]"]}
+```
+
+#let named-type = ```
+{[`named-type` 'identifier']}
+```
+
+#let type-args = ```
+{[`type-args` "[" <!, 'type'*","> "]"]}
+```
+
+#let frame-schema = ```
+{[`frame-schema` "[" <!, ['TOKEN' ":" 'type']*","> "]"]}
+```
+
+#let generic-type = ```
+{[`generic-type` <["Frame" 'frame-schema'], ['named-type' <!, 'type-args'>]>]}
+```
+
+#let grouped-type = ```
+{[`grouped-type` "(" 'type' ")"]}
+```
+
+#let base-type = ```
+{[`base-type` <'grouped-type', 'generic-type'>]}
+```
+
+#let constraint-type = ```
+{[`constraint-type` 'base-type' <!, ["where" 'constraint']>]}
+```
+
+#let pos-param = ```
+{[`pos-param` <!, ['identifier' ":"]> 'type' <!, "?">]}
+```
+
+#let kw-param = ```
+{[`kw-param` 'identifier' ":" 'type' <!, "?">]}
+```
+
+#let param-spec = ```
+{[`param-spec` "(" <!, <'pos-param', "/", "*", 'kw-param'>*",">#`Same rules as Python` ")"]}
+```
+
+#let func-type = ```
+{[`func-type` "fn" 'param-spec' "->" 'type']}
 ```
 
 #let type = ```
-{[`type` 'identifier' <!, 'template'> <!, "?">]}
+{[`type` <'func-type', 'constraint-type'>]}
 ```
 
-#let type-property = ```
-{[`type-property` 'identifier' ":" 'type' <!, ["where" 'constraint']>]}
-```
-
-#let type-body = ```
-{[`type-body` "{" <!, 'type-property'*!> "}"]}
+#let alias-statement = ```
+{[`alias-statement` "alias" 'identifier' "=" 'type']}
 ```
 
 #let type-statement = ```
-{[`type-statement` "type" 'identifier' <!, 'template'> <[["(" 'type' ")"] <!, ["where" 'constraint']>], 'type-body'>]}
+{[`type-statement` "type" 'identifier' <!, 'template'> "=" 'type']}
 ```
 
-#let op-definition = ```
-{[`op-definition` "op" 'identifier' "(" 'type' ")" "->" 'type']}
+#let member-stmt = ```
+{[`member-stmt` <"prop", "def"> 'identifier' ":" 'type']}
 ```
 
 #let extend-statement = ```
-{[`extend-statement` "extend" 'type' "{" <!, 'op-definition'*!> "}"]}
+{[`extend-statement` "extend" 'type' "{" <!, 'member-stmt'*!> "}"]}
 ```
 
 #let predicate-statement = ```
-{[`predicate-statement` "predicate" 'identifier' "(" 'identifier' ":" 'type' ")" "=" 'constraint']}
+{[`predicate-statement` "predicate" 'identifier' <!, 'param-spec'*!> "=" 'constraint']}
 ```
 
 #let statement = ```
-{[`statement` <'type-statement', 'extend-statement', 'predicate-statement'>]}
+{[`statement` <'alias-statement', 'type-statement', 'extend-statement', 'predicate-statement'>]}
 ```
 
 #let rules = (
-  value: value,
+  literal: literal,
   grouping: grouping,
   primary: primary,
   reference: reference,
+  call-args: call-args,
+  call: call,
   unary: unary,
+  factor: factor,
+  term: term,
   comparison: comparison,
   equality: equality,
+  expression: expression,
   constraint: constraint,
-  simple-type: simple-type,
+  template-param: template-param,
   template: template,
-  type: type,
   type-property: type-property,
-  type-body: type-body,
+  named-type: named-type,
+  type-args: type-args,
+  generic-type: generic-type,
+  grouped-type: grouped-type,
+  base-type: base-type,
+  constraint-type: constraint-type,
+  pos-param: pos-param,
+  kw-param: kw-param,
+  param-spec: param-spec,
+  func-type: func-type,
+  type: type,
+  alias-statement: alias-statement,
   type-statement: type-statement,
-  op-definition: op-definition,
+  member-stmt: member-stmt,
   extend-statement: extend-statement,
   predicate-statement: predicate-statement,
   statement: statement,
@@ -106,12 +182,22 @@ svg.railroad .terminal rect {
 
 #let inline = (
   "grouping",
-  "value",
+  "literal",
+  "template-param",
   "template",
-  "simple-type",
   "type-property",
-  "type-body",
-  "op-definition",
+  "call-args",
+  "type-args",
+  "named-type",
+  "grouped-type",
+  "generic-type",
+  "base-type",
+  "constraint-type",
+  "pos-param",
+  "kw-param",
+  "func-type",
+  "member-stmt",
+  "alias-statement",
   "type-statement",
   "extend-statement",
   "predicate-statement",
@@ -128,7 +214,7 @@ svg.railroad .terminal rect {
     2,
     outline(title: none),
   ),
-  height: 9cm,
+  height: 15cm,
   stroke: 1pt,
   inset: 1em,
 )

tests/__main__.py

@@ -0,0 +1,49 @@
+import sys
+from typing import Type
+
+from midas.cli.ansi import Ansi
+from tests.base import Tester, TestsSummary
+from tests.checker import CheckerTester
+from tests.generator import GeneratorTester
+from tests.midas import MidasTester
+from tests.python import PythonTester
+
+
+def print_banner(name: str):
+    horizontal: str = "+" + "-" * (len(name) + 2) + "+"
+    print(horizontal)
+    print(f"| {name} |")
+    print(horizontal)
+
+
+def run_tests(tester_cls: Type[Tester]) -> TestsSummary:
+    print_banner(tester_cls.__name__)
+    tester: Tester = tester_cls()
+    summary: TestsSummary = tester.run_all_tests()
+    print()
+    return summary
+
+
+def main():
+    testers: list[Type[Tester]] = [
+        PythonTester,
+        MidasTester,
+        CheckerTester,
+        GeneratorTester,
+    ]
+
+    summaries: list[TestsSummary] = list(
+        map(run_tests, testers)
+    )  # list to avoid early stop
+    summary: TestsSummary = TestsSummary.concat(*summaries)
+
+    if summary.success:
+        print(Ansi.FG(Ansi.BRIGHT_GREEN) + "All tests passed!" + Ansi.RESET)
+    else:
+        print(Ansi.FG(Ansi.BRIGHT_RED) + "Some tests failed!" + Ansi.RESET)
+        summary.print()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

tests/base.py

@@ -4,14 +4,54 @@ import argparse
 import difflib
 import sys
 from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Iterator, Protocol
 
+from midas.cli.ansi import Ansi
+
 
 class CaseResult(Protocol):
     def dumps(self) -> str: ...
 
 
+@dataclass
+class TestsSummary:
+    tests: list[tuple[str, bool]] = field(default_factory=list)
+
+    @property
+    def success(self) -> bool:
+        return all(map(lambda t: t[1], self.tests))
+
+    @property
+    def successes(self) -> int:
+        return sum(map(lambda t: int(t[1]), self.tests))
+
+    @property
+    def failures(self) -> int:
+        return len(self.tests) - self.successes
+
+    def add(self, name: str, success: bool):
+        self.tests.append((name, success))
+
+    @staticmethod
+    def concat(*summaries: TestsSummary) -> TestsSummary:
+        return TestsSummary(
+            tests=sum(
+                map(lambda s: s.tests, summaries),
+                start=[],
+            ),
+        )
+
+    def print(self):
+        print("Tests summary:")
+        tests: list[tuple[str, bool]] = sorted(self.tests, key=lambda t: t[0])
+        for test, success in tests:
+            print(f" - [{'.' if success else 'X'}] {test}")
+        print("-" * 20)
+        print(f"passed: {self.successes}, failed: {self.failures}")
+
+
 class Tester(ABC):
     """A test runner to check for regressions in the lexer and parser"""
 
@@ -21,6 +61,10 @@ class Tester(ABC):
     @abstractmethod
     def namespace(self) -> str: ...
 
+    @property
+    def extension(self) -> str:
+        return "json"
+
     @property
     def base_dir(self) -> Path:
         return self.CASES_DIR / self.namespace
@@ -28,37 +72,37 @@ class Tester(ABC):
     @abstractmethod
     def _list_tests(self) -> list[Path]: ...
 
-    def run_all_tests(self) -> bool:
-        paths: list[Path] = self._list_tests()
+    def run_all_tests(self) -> TestsSummary:
+        paths: list[Path] = sorted(self._list_tests())
         return self.run_tests(paths)
 
-    def run_tests(self, tests: list[Path]) -> bool:
+    def run_tests(self, tests: list[Path]) -> TestsSummary:
         rule: str = "-" * 80
         n: int = len(tests)
-        successes: int = 0
-        failures: int = 0
+
+        summary: TestsSummary = TestsSummary()
 
         print(rule)
         for i, test in enumerate(tests):
-            print(f"Case {i+1}/{n}: {test.relative_to(self.CASES_DIR)}")
+            path: Path = test.resolve().relative_to(self.CASES_DIR)
+            print(f"{Ansi.FG(Ansi.BRIGHT_CYAN)}Case {i+1}/{n}: {path}{Ansi.RESET}")
             success: bool = self._run_test(test)
-            if success:
-                successes += 1
-            else:
-                failures += 1
+            summary.add(str(path), success)
 
         print(rule)
-        print(f"Success: {successes}/{n}")
-        print(f"Failed: {failures}/{n}")
+        print(f"Success: {summary.successes}/{n}")
+        print(f"Failed: {summary.failures}/{n}")
         print(rule)
-        return failures == 0
+        return summary
 
     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
+        print(Ansi.DIM, end="")
         result: CaseResult = self._exec_case(path)
+        print(Ansi.RESET, end="")
         expected: str = result_path.read_text()
         actual: str = result.dumps()
 
@@ -78,7 +122,7 @@ class Tester(ABC):
     def _exec_case(self, path: Path) -> CaseResult: ...
 
     def update_all_tests(self):
-        paths: list[Path] = self._list_tests()
+        paths: list[Path] = sorted(self._list_tests())
         return self.update_tests(paths)
 
     def update_tests(self, tests: list[Path]):
@@ -99,7 +143,7 @@ class Tester(ABC):
         return True
 
     def _result_path(self, test_path: Path) -> Path:
-        return test_path.parent / (test_path.name + ".ref.json")
+        return test_path.parent / (test_path.name + f".ref.{self.extension}")
 
     def _print_diff(self, diff: Iterator[str]):
         for line in diff:
@@ -134,10 +178,19 @@ class Tester(ABC):
                 else:
                     tester.update_tests(args.FILE)
             case "run":
-                success: bool
+                summary: TestsSummary
                 if args.all:
-                    success = tester.run_all_tests()
+                    summary = tester.run_all_tests()
                 else:
-                    success = tester.run_tests(args.FILE)
-                if not success:
+                    summary = tester.run_tests(args.FILE)
+                if not summary.success:
+                    summary.print()
+                    sys.exit(1)
+            case None:
+                summary: TestsSummary = tester.run_all_tests()
+                if not summary.success:
+                    summary.print()
+                    sys.exit(1)
+            case _:
+                print(f"Unknown subcommand '{args.subcommand}'")
                 sys.exit(1)