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

@@ -6,4 +6,6 @@ venv
 *.pyc
 uv.lock
 .python-version
-/out
+/out
+/examples/**/build/
+/examples/**/*.pyi

README.md

@@ -51,16 +51,7 @@ This framework is being developed as part of a Bachelor's Thesis by Louis Herede
 
 ## Commands
 
-<!--
-check
-compile
-format
-highlight
-parse
-dump_registry
-types
-validate
--->
+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
 
@@ -78,6 +69,10 @@ 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
@@ -123,7 +118,7 @@ This command processes the given Midas definitions and dumps the contents of the
 midas stubs types.midas -o stubs.pyi
 ```
 
-This command generate Python stubs from a Midas definition file
+This command generates Python stubs from a Midas definition file
 
 ### Showing Type Judgements
 
@@ -143,7 +138,7 @@ This command lets you validate a Midas definition file by running the parser and
 
 ## Tests
 
-Several snapshot tests are available to assert the good behaviour of the parsers and type checker. They can be run as follows:
+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
@@ -152,7 +147,16 @@ uv run -m tests.checker run -a
 uv run -m tests.generator run -a
 ```
 
-**Available subcommands:**
+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`

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/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
 

gen/midas.py

@@ -29,9 +29,9 @@ class MemberKind(Enum):
 @dataclass(frozen=True, kw_only=True)
 class ParamSpec:
     l_paren: Token
-    pos: list[FunctionType.Argument]
-    mixed: list[FunctionType.Argument]
-    kw: list[FunctionType.Argument]
+    pos: list[FunctionType.Parameter]
+    mixed: list[FunctionType.Parameter]
+    kw: list[FunctionType.Parameter]
 
 
 ###<
@@ -44,6 +44,11 @@ class TypeStmt:
     type: Type
 
 
+class AliasStmt:
+    name: Token
+    type: Type
+
+
 class MemberStmt:
     name: Token
     type: Type
@@ -131,21 +136,12 @@ class ConstraintType:
     constraint: Expr
 
 
-class ComplexType:
-    members: list[MemberStmt]
-
-
-class ExtensionType:
-    base: Type
-    extension: ComplexType
-
-
 class FunctionType:
     params: ParamSpec
     returns: Type
 
     @dataclass(frozen=True, kw_only=True)
-    class Argument:
+    class Parameter:
         location: Optional[Location] = None
         name: Optional[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
@@ -92,6 +110,16 @@ class ForStmt:
     body: list[Stmt]
 
 
+class ImportStmt:
+    imports: list[ImportAlias]
+
+
+class FromImportStmt:
+    module: Optional[str]
+    imports: list[ImportAlias]
+    level: int
+
+
 class RawStmt:
     stmt: ast.stmt
 
@@ -174,6 +202,10 @@ class SliceExpr:
     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

@@ -30,9 +30,9 @@ class MemberKind(Enum):
 @dataclass(frozen=True, kw_only=True)
 class ParamSpec:
     l_paren: Token
-    pos: list[FunctionType.Argument]
-    mixed: list[FunctionType.Argument]
-    kw: list[FunctionType.Argument]
+    pos: list[FunctionType.Parameter]
+    mixed: list[FunctionType.Parameter]
+    kw: list[FunctionType.Parameter]
 
 
 ##############
@@ -51,6 +51,9 @@ class Stmt(ABC):
         @abstractmethod
         def visit_type_stmt(self, stmt: TypeStmt) -> T: ...
 
+        @abstractmethod
+        def visit_alias_stmt(self, stmt: AliasStmt) -> T: ...
+
         @abstractmethod
         def visit_member_stmt(self, stmt: MemberStmt) -> T: ...
 
@@ -71,6 +74,15 @@ class TypeStmt(Stmt):
         return visitor.visit_type_stmt(self)
 
 
+@dataclass(frozen=True)
+class AliasStmt(Stmt):
+    name: Token
+    type: Type
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_alias_stmt(self)
+
+
 @dataclass(frozen=True)
 class MemberStmt(Stmt):
     name: Token
@@ -244,12 +256,6 @@ class Type(ABC):
         @abstractmethod
         def visit_constraint_type(self, type: ConstraintType) -> T: ...
 
-        @abstractmethod
-        def visit_complex_type(self, type: ComplexType) -> T: ...
-
-        @abstractmethod
-        def visit_extension_type(self, type: ExtensionType) -> T: ...
-
         @abstractmethod
         def visit_function_type(self, type: FunctionType) -> T: ...
 
@@ -283,30 +289,13 @@ class ConstraintType(Type):
         return visitor.visit_constraint_type(self)
 
 
-@dataclass(frozen=True)
-class ComplexType(Type):
-    members: list[MemberStmt]
-
-    def accept(self, visitor: Type.Visitor[T]) -> T:
-        return visitor.visit_complex_type(self)
-
-
-@dataclass(frozen=True)
-class ExtensionType(Type):
-    base: Type
-    extension: ComplexType
-
-    def accept(self, visitor: Type.Visitor[T]) -> T:
-        return visitor.visit_extension_type(self)
-
-
 @dataclass(frozen=True)
 class FunctionType(Type):
     params: ParamSpec
     returns: Type
 
     @dataclass(frozen=True, kw_only=True)
-    class Argument:
+    class Parameter:
         location: Optional[Location] = None
         name: Optional[Token]
         type: Type

midas/ast/printer.py

@@ -1,868 +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_param(param)
-            self._write_line("type", last=True)
-            with self._child_level(single=True):
-                stmt.type.accept(self)
-
-    def _print_type_param(self, param: m.TypeParam) -> None:
-        self._write_line("Param")
-        with self._child_level():
-            self._write_line(f'name: "{param.name.lexeme}"')
-            self._write_optional_child("bound", param.bound, last=True)
-
-    def visit_member_stmt(self, stmt: m.MemberStmt):
-        self._write_line("MemberStmt")
-        with self._child_level():
-            self._write_line(f"kind: {stmt.kind.name}")
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line("type", last=True)
-            with self._child_level(single=True):
-                stmt.type.accept(self)
-
-    def visit_extend_stmt(self, stmt: m.ExtendStmt) -> None:
-        self._write_line("ExtendStmt")
-        with self._child_level():
-            self._write_line("params")
-            with self._child_level():
-                for i, param in enumerate(stmt.params):
-                    self._idx = i
-                    if i == len(stmt.params) - 1:
-                        self._mark_last()
-                    self._print_type_param(param)
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line("params")
-            with self._child_level():
-                for i, param in enumerate(stmt.params):
-                    self._idx = i
-                    if i == len(stmt.params) - 1:
-                        self._mark_last()
-                    self._print_type_param(param)
-            self._write_line("members", last=True)
-            with self._child_level():
-                for i, member in enumerate(stmt.members):
-                    self._idx = i
-                    if i == len(stmt.members) - 1:
-                        self._mark_last()
-                    member.accept(self)
-
-    def visit_predicate_stmt(self, stmt: m.PredicateStmt):
-        self._write_line("PredicateStmt")
-        with self._child_level():
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line("params")
-            with self._child_level():
-                for i, spec in enumerate(stmt.params):
-                    self._idx = i
-                    if i == len(stmt.params) - 1:
-                        self._mark_last()
-                    self._visit_param_spec(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_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: m.GetExpr):
-        self._write_line("GetExpr")
-        with self._child_level():
-            self._write_line("expr")
-            with self._child_level(single=True):
-                expr.expr.accept(self)
-            self._write_line(f'name: "{expr.name.lexeme}"', last=True)
-
-    def visit_variable_expr(self, expr: m.VariableExpr):
-        self._write_line("VariableExpr")
-        with self._child_level():
-            self._write_line(f'name: "{expr.name.lexeme}"', last=True)
-
-    def visit_grouping_expr(self, expr: m.GroupingExpr):
-        self._write_line("GroupingExpr")
-        with self._child_level():
-            self._write_line("expr", last=True)
-            with self._child_level(single=True):
-                expr.expr.accept(self)
-
-    def visit_literal_expr(self, expr: m.LiteralExpr) -> None:
-        self._write_line("LiteralExpr")
-        with self._child_level():
-            self._write_line(f"value: {expr.value}", last=True)
-
-    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> None:
-        self._write_line("WildcardExpr")
-
-    def visit_named_type(self, type: m.NamedType) -> None:
-        self._write_line("NamedType")
-        with self._child_level():
-            self._write_line(f'name: "{type.name.lexeme}"', last=True)
-
-    def visit_generic_type(self, type: m.GenericType) -> None:
-        self._write_line("GenericType")
-        with self._child_level():
-            self._write_line("type")
-            with self._child_level():
-                type.type.accept(self)
-            self._write_line("args", last=True)
-            with self._child_level():
-                for i, param in enumerate(type.args):
-                    self._idx = i
-                    if i == len(type.args) - 1:
-                        self._mark_last()
-                    param.accept(self)
-
-    def visit_constraint_type(self, type: m.ConstraintType) -> None:
-        self._write_line("ConstraintType")
-        with self._child_level():
-            self._write_line("type")
-            with self._child_level(single=True):
-                type.type.accept(self)
-            self._write_line("constraint", last=True)
-            with self._child_level(single=True):
-                type.constraint.accept(self)
-
-    def visit_complex_type(self, type: m.ComplexType) -> None:
-        self._write_line("ComplexType")
-        with self._child_level():
-            self._write_line("members", last=True)
-            with self._child_level():
-                for i, member in enumerate(type.members):
-                    self._idx = i
-                    if i == len(type.members) - 1:
-                        self._mark_last()
-                    member.accept(self)
-
-    def visit_extension_type(self, type: m.ExtensionType) -> None:
-        self._write_line("ExtensionType")
-        with self._child_level():
-            self._write_line("base")
-            with self._child_level(single=True):
-                type.base.accept(self)
-            self._write_line("extension", last=True)
-            with self._child_level(single=True):
-                type.extension.accept(self)
-
-    def visit_function_type(self, type: m.FunctionType) -> None:
-        self._write_line("FunctionType")
-        with self._child_level():
-            self._write_line("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_line("pos")
-            with self._child_level():
-                for i, arg in enumerate(spec.pos):
-                    self._idx = i
-                    if i == len(spec.pos) - 1:
-                        self._mark_last()
-                    self._print_function_arg(arg)
-
-            self._write_line("mixed")
-            with self._child_level():
-                for i, arg in enumerate(spec.mixed):
-                    self._idx = i
-                    if i == len(spec.mixed) - 1:
-                        self._mark_last()
-                    self._print_function_arg(arg)
-
-            self._write_line("kw", last=True)
-            with self._child_level():
-                for i, arg in enumerate(spec.kw):
-                    self._idx = i
-                    if i == len(spec.kw) - 1:
-                        self._mark_last()
-                    self._print_function_arg(arg)
-
-    def _print_function_arg(self, arg: m.FunctionType.Argument) -> None:
-        self._write_line("Argument")
-        with self._child_level():
-            name: str = "None"
-            if arg.name is not None:
-                name = f'"{arg.name.lexeme}"'
-            self._write_line(f"name: {name}")
-            self._write_line("type")
-            with self._child_level(single=True):
-                arg.type.accept(self)
-            self._write_line(f"required: {arg.required}", last=True)
-
-    def visit_frame_type(self, type: m.FrameType) -> None:
-        self._write_line("FrameType")
-        with self._child_level(single=True):
-            self._write_line("columns")
-            with self._child_level():
-                for i, column in enumerate(type.columns):
-                    self._idx = i
-                    if i == len(type.columns) - 1:
-                        self._mark_last()
-                    self._print_frame_column(column)
-
-    def _print_frame_column(self, column: m.FrameType.Column) -> None:
-        self._write_line("Column")
-        with self._child_level():
-            self._write_line(f'name: "{column.name.lexeme}"')
-            self._write_line("type")
-            with self._child_level(single=True):
-                column.type.accept(self)
-
-
-class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str]):
-    def __init__(self, indent: int = 4):
-        self.indent: int = indent
-        self.level: int = 0
-
-    def indented(self, text: str) -> str:
-        return " " * (self.level * self.indent) + text
-
-    def print(self, expr: m.Expr | m.Stmt | m.Type) -> str:
-        self.level = 0
-        return expr.accept(self)
-
-    def visit_type_stmt(self, stmt: m.TypeStmt) -> str:
-        template: str = ""
-        if len(stmt.params) != 0:
-            params: list[str] = [self._print_type_param(param) for param in stmt.params]
-            template = f"[{', '.join(params)}]"
-        res: str = f"type {stmt.name.lexeme}{template} = {stmt.type.accept(self)}"
-        return self.indented(res)
-
-    def _print_type_param(self, param: m.TypeParam) -> str:
-        res: str = param.name.lexeme
-        if param.bound is not None:
-            res += "<:" + param.bound.accept(self)
-        return res
-
-    def visit_member_stmt(self, stmt: m.MemberStmt):
-        keyword: str = {
-            m.MemberKind.PROPERTY: "prop",
-            m.MemberKind.METHOD: "def",
-        }.get(stmt.kind, "")
-        res: str = f"{keyword} {stmt.name.lexeme}: {stmt.type.accept(self)}"
-        return self.indented(res)
-
-    def visit_extend_stmt(self, stmt: m.ExtendStmt):
-        template: str = ""
-        if len(stmt.params) != 0:
-            params: list[str] = [self._print_type_param(param) for param in stmt.params]
-            template = f"[{', '.join(params)}]"
-        res: str = self.indented(f"extend {stmt.name.lexeme}{template}")
-        res += " {\n"
-        self.level += 1
-        for member in stmt.members:
-            res += member.accept(self) + "\n"
-        self.level -= 1
-        res += self.indented("}")
-        return res
-
-    def visit_predicate_stmt(self, stmt: m.PredicateStmt):
-        name: str = stmt.name.lexeme
-        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}")
-
-    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 "_"
-
-    def visit_named_type(self, type: m.NamedType) -> str:
-        return type.name.lexeme
-
-    def visit_generic_type(self, type: m.GenericType) -> str:
-        res: str = type.type.accept(self)
-        if len(type.args) != 0:
-            args: list[str] = [param.accept(self) for param in type.args]
-            res += f"[{', '.join(args)}]"
-        return res
-
-    def visit_constraint_type(self, type: m.ConstraintType) -> str:
-        res: str = type.type.accept(self)
-        res += " where " + type.constraint.accept(self)
-        return res
-
-    def visit_complex_type(self, type: m.ComplexType) -> str:
-        res: str = "{\n"
-        self.level += 1
-        for member in type.members:
-            res += member.accept(self)
-            res += "\n"
-        self.level -= 1
-        res += self.indented("}")
-        return res
-
-    def visit_extension_type(self, type: m.ExtensionType) -> str:
-        return f"{type.base.accept(self)} & {type.extension.accept(self)}"
-
-    def visit_function_type(self, type: m.FunctionType) -> str:
-        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_args: list[str] = [self._print_arg(arg) for arg in spec.pos]
-        mixed_args: list[str] = [self._print_arg(arg) for arg in spec.mixed]
-        kw_args: list[str] = [self._print_arg(arg) for arg in spec.kw]
-        args: list[str] = pos_args
-
-        if len(pos_args) != 0:
-            args.append("/")
-        args += mixed_args
-        if len(kw_args) != 0:
-            args.append("*")
-        args += kw_args
-        return f"({', '.join(args)})"
-
-    def _print_arg(self, arg: m.FunctionType.Argument) -> str:
-        res: str = ""
-        if arg.name is not None:
-            res += arg.name.lexeme
-            res += ": "
-        res += arg.type.accept(self)
-        if not arg.required:
-            res += "?"
-        return res
-
-    def visit_frame_type(self, type: m.FrameType) -> str:
-        res: str = self.indented("Frame[")
-        if len(type.columns) != 0:
-            res += "\n"
-            self.level += 1
-            columns: list[str] = []
-            for column in type.columns:
-                columns.append(self.indented(self._print_frame_column(column)))
-            res += ",\n".join(columns)
-            self.level -= 1
-            res += "\n"
-        res += "]"
-        return res
-
-    def _print_frame_column(self, column: m.FrameType.Column) -> str:
-        return f"{column.name.lexeme}: {column.type.accept(self)}"
-
-
-class PythonAstPrinter(
-    AstPrinter,
-    p.MidasType.Visitor[None],
-    p.Stmt.Visitor[None],
-    p.Expr.Visitor[None],
-):
-    def visit_base_type(self, node: p.BaseType) -> None:
-        self._write_line("BaseType")
-        with self._child_level():
-            self._write_line(f"base: {node.base}")
-            self._write_optional_child("param", node.param, last=True)
-
-    def visit_constraint_type(self, node: p.ConstraintType) -> None:
-        self._write_line("ConstraintType")
-        with self._child_level():
-            self._write_line("type")
-            with self._child_level(single=True):
-                node.type.accept(self)
-            self._write_line(f"constraint: {ast.unparse(node.constraint)}", last=True)
-
-    def visit_frame_column(self, node: p.FrameColumn) -> None:
-        self._write_line("FrameColumn")
-        with self._child_level():
-            self._write_line(f"name: {node.name}")
-            self._write_optional_child("type", node.type, last=True)
-
-    def visit_frame_type(self, node: p.FrameType) -> None:
-        self._write_line("FrameType")
-        with self._child_level():
-            self._write_line("columns", last=True)
-            with self._child_level():
-                for i, col in enumerate(node.columns):
-                    self._idx = i
-                    if i == len(node.columns) - 1:
-                        self._mark_last()
-                    col.accept(self)
-
-    def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
-        stmt.expr.accept(self)
-
-    def visit_function(self, stmt: p.Function) -> None:
-        self._write_line("Function")
-        with self._child_level():
-            self._write_line(f"name: {stmt.name}")
-
-            self._write_line("posonlyargs")
-            with self._child_level():
-                for i, arg in enumerate(stmt.posonlyargs):
-                    self._idx = i
-                    if i == len(stmt.posonlyargs) - 1:
-                        self._mark_last()
-                    self._print_argument(arg)
-
-            self._write_line("args")
-            with self._child_level():
-                for i, arg in enumerate(stmt.args):
-                    self._idx = i
-                    if i == len(stmt.args) - 1:
-                        self._mark_last()
-                    self._print_argument(arg)
-
-            self._write_line("kwonlyargs")
-            with self._child_level():
-                for i, arg in enumerate(stmt.kwonlyargs):
-                    self._idx = i
-                    if i == len(stmt.kwonlyargs) - 1:
-                        self._mark_last()
-                    self._print_argument(arg)
-
-            self._write_optional_child("returns", stmt.returns)
-            self._write_line("body", last=True)
-            with self._child_level():
-                for i, body_stmt in enumerate(stmt.body):
-                    self._idx = i
-                    if i == len(stmt.body) - 1:
-                        self._mark_last()
-                    body_stmt.accept(self)
-
-    def _print_argument(self, arg: p.Function.Argument) -> None:
-        self._write_line("FunctionArgument")
-        with self._child_level():
-            self._write_line(f"name: {arg.name}")
-            self._write_optional_child("type", arg.type, last=True)
-
-    def visit_type_assign(self, stmt: p.TypeAssign) -> None:
-        self._write_line("TypeAssign")
-        with self._child_level():
-            self._write_line(f"name: {stmt.name}")
-            self._write_line("type", last=True)
-            with self._child_level(single=True):
-                stmt.type.accept(self)
-
-    def visit_assign_stmt(self, stmt: p.AssignStmt) -> None:
-        self._write_line("AssignStmt")
-        with self._child_level():
-            self._write_line("targets")
-            with self._child_level():
-                for i, target in enumerate(stmt.targets):
-                    self._idx = i
-                    if i == len(stmt.targets) - 1:
-                        self._mark_last()
-                    target.accept(self)
-            self._write_line("value", last=True)
-            with self._child_level(single=True):
-                stmt.value.accept(self)
-
-    def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
-        self._write_line("ReturnStmt")
-        with self._child_level():
-            self._write_optional_child("value", stmt.value, last=True)
-
-    def visit_if_stmt(self, stmt: p.IfStmt) -> None:
-        self._write_line("IfStmt")
-        with self._child_level():
-            self._write_line("test")
-            with self._child_level(single=True):
-                stmt.test.accept(self)
-            self._write_line("body")
-            with self._child_level():
-                for i, body_stmt in enumerate(stmt.body):
-                    self._idx = i
-                    if i == len(stmt.body) - 1:
-                        self._mark_last()
-                    body_stmt.accept(self)
-            self._write_line("orelse", last=True)
-            with self._child_level():
-                for i, else_stmt in enumerate(stmt.orelse):
-                    self._idx = i
-                    if i == len(stmt.orelse) - 1:
-                        self._mark_last()
-                    else_stmt.accept(self)
-
-    def visit_pass(self, stmt: p.Pass) -> None:
-        self._write_line("Pass")
-
-    def visit_for_stmt(self, stmt: p.ForStmt) -> None:
-        self._write_line("ForStmt")
-        with self._child_level():
-            self._write_line("target")
-            with self._child_level(single=True):
-                stmt.target.accept(self)
-            self._write_line("iterator")
-            with self._child_level(single=True):
-                stmt.iterator.accept(self)
-            self._write_line("body", last=True)
-            with self._child_level():
-                for i, body_stmt in enumerate(stmt.body):
-                    self._idx = i
-                    if i == len(stmt.body) - 1:
-                        self._mark_last()
-                    body_stmt.accept(self)
-
-    def visit_raw_stmt(self, stmt: p.RawStmt) -> None:
-        self._write_line("RawStmt")
-        with self._child_level(single=True):
-            self._write_line(f"stmt: {ast.unparse(stmt.stmt)}")
-
-    def visit_binary_expr(self, expr: p.BinaryExpr) -> None:
-        self._write_line("BinaryExpr")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level(single=True):
-                expr.left.accept(self)
-
-            self._write_line(f"operator: {expr.operator.__class__.__name__}")
-
-            self._write_line("right", last=True)
-            with self._child_level(single=True):
-                expr.right.accept(self)
-
-    def visit_compare_expr(self, expr: p.CompareExpr) -> None:
-        self._write_line("CompareExpr")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level(single=True):
-                expr.left.accept(self)
-
-            self._write_line(f"operator: {expr.operator.__class__.__name__}")
-
-            self._write_line("right", last=True)
-            with self._child_level(single=True):
-                expr.right.accept(self)
-
-    def visit_unary_expr(self, expr: p.UnaryExpr) -> None:
-        self._write_line("UnaryExpr")
-        with self._child_level():
-            self._write_line(f"operator: {expr.operator.__class__.__name__}")
-
-            self._write_line("right", last=True)
-            with self._child_level(single=True):
-                expr.right.accept(self)
-
-    def visit_call_expr(self, expr: p.CallExpr) -> None:
-        self._write_line("CallExpr")
-        with self._child_level():
-            self._write_line("callee")
-            with self._child_level(single=True):
-                expr.callee.accept(self)
-
-            self._write_line("arguments")
-            with self._child_level():
-                for i, arg in enumerate(expr.arguments):
-                    self._idx = i
-                    if i == len(expr.arguments) - 1:
-                        self._mark_last()
-                    arg.accept(self)
-
-            self._write_line("keywords", last=True)
-            with self._child_level():
-                for i, (name, arg) in enumerate(expr.keywords.items()):
-                    self._idx = i
-                    if i == len(expr.keywords) - 1:
-                        self._mark_last()
-                    self._write_line(name)
-                    with self._child_level(single=True):
-                        arg.accept(self)
-
-    def visit_get_expr(self, expr: p.GetExpr) -> None:
-        self._write_line("GetExpr")
-        with self._child_level():
-            self._write_line("object")
-            with self._child_level(single=True):
-                expr.object.accept(self)
-            self._write_line(f"name: {expr.name}", last=True)
-
-    def visit_literal_expr(self, expr: p.LiteralExpr) -> None:
-        self._write_line("LiteralExpr")
-        with self._child_level(single=True):
-            self._write_line(f"value: {expr.value!r}")
-
-    def visit_variable_expr(self, expr: p.VariableExpr) -> None:
-        self._write_line("VariableExpr")
-        with self._child_level(single=True):
-            self._write_line(f"name: {expr.name}")
-
-    def visit_logical_expr(self, expr: p.LogicalExpr) -> None:
-        self._write_line("LogicalExpr")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level(single=True):
-                expr.left.accept(self)
-
-            self._write_line(f"operator: {expr.operator.__class__.__name__}")
-
-            self._write_line("right", last=True)
-            with self._child_level(single=True):
-                expr.right.accept(self)
-
-    def visit_cast_expr(self, expr: p.CastExpr) -> None:
-        self._write_line("CastExpr")
-        with self._child_level():
-            self._write_line("type")
-            with self._child_level(single=True):
-                expr.type.accept(self)
-            self._write_line("expr")
-            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_line("items", last=True)
-            with self._child_level():
-                for i, item in enumerate(expr.items):
-                    self._idx = i
-                    if i == len(expr.items) - 1:
-                        self._mark_last()
-                    item.accept(self)
-
-    def visit_dict_expr(self, expr: p.DictExpr) -> None:
-        self._write_line("DictExpr")
-        with self._child_level():
-            self._write_line("keys")
-            with self._child_level():
-                for i, key in enumerate(expr.keys):
-                    self._idx = i
-                    if i == len(expr.keys) - 1:
-                        self._mark_last()
-                    if key is None:
-                        self._write_line("None")
-                    else:
-                        key.accept(self)
-            self._write_line("values", last=True)
-            with self._child_level():
-                for i, value in enumerate(expr.values):
-                    self._idx = i
-                    if i == len(expr.values) - 1:
-                        self._mark_last()
-                    value.accept(self)
-
-    def visit_subscript_expr(self, expr: p.SubscriptExpr) -> None:
-        self._write_line("SubscriptExpr")
-        with self._child_level():
-            self._write_line("object")
-            with self._child_level(single=True):
-                expr.object.accept(self)
-            self._write_line("index", last=True)
-            with self._child_level(single=True):
-                expr.index.accept(self)
-
-    def visit_slice_expr(self, expr: p.SliceExpr) -> None:
-        self._write_line("SliceExpr")
-        with self._child_level():
-            self._write_optional_child("lower", expr.lower)
-            self._write_optional_child("upper", expr.upper)
-            self._write_optional_child("step", expr.step, last=True)
-
-    def visit_raw_expr(self, expr: p.RawExpr) -> None:
-        self._write_line("RawExpr")
-        with self._child_level(single=True):
-            self._write_line(f"expr: {ast.unparse(expr.expr)}")

midas/ast/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,27 @@ 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 #
@@ -44,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)
@@ -113,6 +134,12 @@ class Stmt(ABC):
         @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: ...
 
@@ -128,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)
 
@@ -205,6 +224,24 @@ class ForStmt(Stmt):
         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
@@ -268,6 +305,9 @@ class Expr(ABC):
         @abstractmethod
         def visit_slice_expr(self, expr: SliceExpr) -> T: ...
 
+        @abstractmethod
+        def visit_tuple_expr(self, expr: TupleExpr) -> T: ...
+
         @abstractmethod
         def visit_raw_expr(self, expr: RawExpr) -> T: ...
 
@@ -402,6 +442,14 @@ class SliceExpr(Expr):
         return visitor.visit_slice_expr(self)
 
 
+@dataclass(frozen=True)
+class TupleExpr(Expr):
+    items: tuple[Expr, ...]
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_tuple_expr(self)
+
+
 @dataclass(frozen=True)
 class RawExpr(Expr):
     expr: ast.expr

midas/checker/builtins.midas

@@ -178,4 +178,100 @@ extend dict[K, V] {
     // 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

@@ -15,10 +15,14 @@ if TYPE_CHECKING:
 
 
 BUILTIN_SUBTYPES: dict[str, set[str]] = {
-    "object": {"float", "list", "dict", "str"},
+    "object": {"float", "list", "dict", "str", "bytes", "tuple"},
     "float": {"int"},
-    "int": {"bool"},
 }
+"""
+Hard-coded subtype relationships between builtin types
+
+Circular dependencies and diamond inheritance MUST be avoided
+"""
 
 
 def define_builtins(reg: TypesRegistry):
@@ -26,12 +30,15 @@ def define_builtins(reg: TypesRegistry):
     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(

midas/checker/checker.py

@@ -10,6 +10,11 @@ from midas.utils import TypedAST
 
 
 class TypeChecker:
+    """Type checking dispatcher
+
+    Contains a typer for Midas and one for Python, as well as the types registry
+    """
+
     def __init__(self):
         self.types: TypesRegistry = TypesRegistry()
         self.reporter: Reporter = Reporter()

midas/checker/diagnostic.py

@@ -14,6 +14,12 @@ class DiagnosticType(StrEnum):
 
 @dataclass(frozen=True)
 class Diagnostic:
+    """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
@@ -21,6 +27,18 @@ class Diagnostic:
 
     @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 (

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

@@ -2,6 +2,7 @@ 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
@@ -11,10 +12,18 @@ 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
@@ -22,16 +31,51 @@ class Evaluator(m.Expr.Visitor[Any]):
         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:
@@ -56,6 +100,8 @@ class Evaluator(m.Expr.Visitor[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:
@@ -80,8 +126,12 @@ class Evaluator(m.Expr.Visitor[Any]):
     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
 
@@ -94,7 +144,7 @@ class Evaluator(m.Expr.Visitor[Any]):
 
         match callee:
             case Predicate():
-                return self._evaluate_predicate(callee, args, kwargs)
+                return self._evaluate_predicate(expr.location, callee, args, kwargs)
             case _ if callable(callee):
                 return callee(*args, **kwargs)
             case _:
@@ -131,8 +181,22 @@ class Evaluator(m.Expr.Visitor[Any]):
         return self.get_value("_")
 
     def _evaluate_predicate(
-        self, predicate: Predicate, args: list[Any], kwargs: dict[str, Any]
+        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)
@@ -140,7 +204,7 @@ class Evaluator(m.Expr.Visitor[Any]):
             self.scopes.append({})
         match predicate.type:
             case Function(returns=Function() as inner):
-                self._map_args(predicate.type, args, kwargs)
+                self._map_args(location, predicate.type, args, kwargs)
                 res = PartialPredicate(
                     type=inner,
                     body=predicate.body,
@@ -149,7 +213,7 @@ class Evaluator(m.Expr.Visitor[Any]):
                 )
 
             case Function():
-                self._map_args(predicate.type, args, kwargs)
+                self._map_args(location, predicate.type, args, kwargs)
                 res = self.evaluate(predicate.body)
 
             case _:
@@ -157,16 +221,45 @@ class Evaluator(m.Expr.Visitor[Any]):
         self.scopes.pop()
         return res
 
-    def _map_args(self, function: Function, args: list[Any], kwargs: dict[str, Any]):
-        positional: list[Function.Argument] = function.pos_args + function.args
-        keywords: dict[str, Function.Argument] = {
-            arg.name: arg for arg in function.args + function.kw_args
+    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):
-            param: Function.Argument = positional[i]
+            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():
-            param: Function.Argument = keywords[name]
+            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/frame_methods.py

@@ -1,198 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Optional
-
-from midas.ast.location import Location
-from midas.checker.registry import TypesRegistry
-from midas.checker.reporter import FileReporter
-from midas.checker.types import (
-    ColumnType,
-    DataFrameType,
-    Function,
-    OverloadedFunction,
-    TopType,
-    Type,
-    UnknownType,
-    unfold_type,
-)
-
-if TYPE_CHECKING:
-    from midas.checker.python import PythonTyper, TypedExpr
-
-
-@staticmethod
-def frame_method(*names: str):
-    def wrapper(func):
-        names_: tuple[str, ...] = names
-        if len(names_) == 0:
-            names_ = (func.__name__,)
-        setattr(func, "__method_names__", names_)
-        return func
-
-    return wrapper
-
-
-@dataclass(frozen=True, kw_only=True)
-class Call:
-    location: Location
-    frame: DataFrameType
-    positional: list[TypedExpr]
-    keywords: dict[str, TypedExpr]
-
-
-class _MethodRegistryMeta(type):
-    _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 MethodRegistry(metaclass=_MethodRegistryMeta):
-    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
-
-    def call(
-        self,
-        method: str,
-        call: Call,
-    ) -> Type:
-        func: Optional[Callable[..., Type]] = self._methods.get(method)
-        if func is None:
-            self.reporter.warning(call.location, f"Unknown method {method}")
-            return UnknownType()
-        return func(self, call)
-
-    @frame_method("add", "__add__")
-    def add(
-        self,
-        call: Call,
-    ) -> Type:
-        # TODO: support add with scalar, sequence, Series, dict
-        # TODO: check operation exists on inner column types
-
-        new_columns: list[DataFrameType.Column] = []
-
-        by_name: dict[str, DataFrameType.Column] = {}
-        frame2: Optional[DataFrameType] = None
-        if len(call.positional) != 0:
-            other: Type = call.positional[0][1]
-            unfolded_other: Type = unfold_type(other)
-            if isinstance(unfolded_other, DataFrameType):
-                frame2 = unfolded_other
-                by_name = {
-                    col.name: col for col in frame2.columns if col.name is not None
-                }
-
-        in_frame1: set[str] = set()
-        for column in call.frame.columns:
-            if column.name is not None:
-                in_frame1.add(column.name)
-
-            col_type1: Type = column.type
-            col_type: Type = ColumnType(type=UnknownType())
-            if column.name in by_name:
-                column2 = by_name[column.name]
-                col_type2: Type = column2.type
-                if self.types.are_equivalent(col_type2, col_type1):
-                    col_type = col_type1
-
-            new_column = DataFrameType.Column(
-                index=column.index,
-                name=column.name,
-                type=col_type,
-            )
-            new_columns.append(new_column)
-
-        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()),
-                    )
-                )
-
-        signature = Function(
-            args=[
-                Function.Argument(
-                    pos=0,
-                    name="other",
-                    type=DataFrameType(columns=[]),
-                    required=True,
-                ),
-            ],
-            returns=DataFrameType(columns=new_columns),
-        )
-
-        return (
-            self.typer._get_call_result(
-                location=call.location,
-                callee=signature,
-                positional=call.positional,
-                keywords=call.keywords,
-            )
-            or UnknownType()
-        )
-
-    @frame_method()
-    def mean(self, call: Call) -> Type:
-        with_axis = Function(
-            kw_args=[
-                Function.Argument(
-                    pos=0,
-                    name="axis",
-                    type=self.types.get_type("int"),
-                    required=False,
-                )
-            ],
-            returns=ColumnType(type=TopType()),
-        )
-        without_axis = Function(
-            kw_args=[
-                Function.Argument(
-                    pos=0,
-                    name="axis",
-                    type=self.types.get_type("None"),
-                    required=True,
-                )
-            ],
-            returns=TopType(),
-        )
-        overload = OverloadedFunction(
-            overloads=[
-                with_axis,
-                without_axis,
-            ]
-        )
-        return (
-            self.typer._get_call_result(
-                location=call.location,
-                callee=overload,
-                positional=call.positional,
-                keywords=call.keywords,
-            )
-            or UnknownType()
-        )

midas/checker/frames.py

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

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

@@ -1,74 +1,46 @@
 import logging
-from dataclasses import dataclass
 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 (
-    AliasType,
-    AppliedType,
     ColumnType,
-    ComplexType,
     ConstraintType,
     DataFrameType,
-    ExtensionType,
+    DerivedType,
     Function,
     GenericType,
-    OverloadedFunction,
+    ParamSpec,
     Predicate,
     Type,
     TypeVar,
     UnknownType,
-    unfold_type,
 )
 from midas.checker.variance import VarianceInferrer
 from midas.lexer.midas import MidasLexer
-from midas.lexer.token import Token
+from midas.lexer.token import Token, TokenType
 from midas.parser.midas import MidasParser
 
 
-@dataclass(frozen=True, kw_only=True)
-class TypedParamSpec:
-    pos: list[Function.Argument]
-    mixed: list[Function.Argument]
-    kw: list[Function.Argument]
-
-
-TypedExpr = tuple[m.Expr, Type]
-
-
-class ReturnException(Exception):
-    pass
-
-
-@dataclass(frozen=True, kw_only=True)
-class MappedArgument:
-    expr: m.Expr
-    type: Type
-    argument: Function.Argument
-
-
-@dataclass(frozen=True, kw_only=True)
-class OverloadCandidate:
-    function: Function
-    mapped: list[MappedArgument]
-
-
 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] = {}
@@ -83,8 +55,25 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
 
         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]):
-        self.reporter = self.reporter.for_file(path)
+        """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)
@@ -94,6 +83,14 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         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
 
@@ -114,6 +111,21 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         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)
@@ -141,6 +153,11 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
                 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}")
@@ -154,11 +171,18 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         if len(params) != 0:
             type = GenericType(name=name, params=params, body=type)
         else:
-            type = AliasType(name=name, type=type)
+            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:
@@ -185,9 +209,7 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
                 self._predicate_params[param.name.lexeme] = param.type.accept(self)
 
         type: Type = self.type_of(stmt.body)
-        params: list[TypedParamSpec] = [
-            self._visit_param_spec(spec) for spec in stmt.params
-        ]
+        params: list[ParamSpec] = [self._visit_param_spec(spec) for spec in stmt.params]
 
         if not self._is_valid_predicate(type):
             self.reporter.error(
@@ -198,9 +220,7 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
             type = self._bool
             for spec in reversed(params):
                 type = Function(
-                    pos_args=spec.pos,
-                    args=spec.mixed,
-                    kw_args=spec.kw,
+                    params=spec,
                     returns=type,
                 )
         self._predicate_params = {}
@@ -214,6 +234,16 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         )
 
     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)
@@ -239,7 +269,11 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         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
+        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)
@@ -252,15 +286,20 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
             )
             return UnknownType()
 
-        result: Optional[Type] = self._get_call_result(
-            location,
-            operation,
-            [(right_expr, right)],
-            {},
+        result: CallResult = self.dispatcher.get_result(
+            location=location,
+            callee=operation,
+            positional=[(right_expr, right)],
+            keywords={},
         )
-        return result or UnknownType()
+        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}")
@@ -278,31 +317,29 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
             )
             return UnknownType()
 
-        result: Optional[Type] = self._get_call_result(
-            expr.location,
-            operation,
-            [],
-            {},
+        result: CallResult = self.dispatcher.get_result(
+            location=expr.location,
+            callee=operation,
+            positional=[],
+            keywords={},
         )
-        return result or UnknownType()
+        return result.result
 
     def visit_call_expr(self, expr: m.CallExpr) -> Type:
         callee: Type = expr.callee.accept(self)
-        positional: list[TypedExpr] = [
+        positional: list[tuple[m.Expr, Type]] = [
             (arg, self.type_of(arg)) for arg in expr.arguments
         ]
-        keywords: dict[str, TypedExpr] = {
+        keywords: dict[str, tuple[m.Expr, Type]] = {
             name: (arg, self.type_of(arg)) for name, arg in expr.keywords.items()
         }
-        return (
-            self._get_call_result(
-                expr.location,
-                callee,
-                positional,
-                keywords,
-            )
-            or UnknownType()
+        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)
@@ -349,6 +386,16 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
             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:
@@ -358,49 +405,50 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
             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=type.type.accept(self),
+            type=base_type,
             constraint=type.constraint,
         )
 
-    def visit_complex_type(self, type: m.ComplexType) -> ComplexType:
-        return ComplexType(
-            members={
-                member.name.lexeme: member.type.accept(self) for member in type.members
-            }
-        )
-
-    def visit_extension_type(self, type: m.ExtensionType) -> Type:
-        return ExtensionType(
-            base=type.base.accept(self),
-            extension=self.visit_complex_type(type.extension),
-        )
-
     def visit_function_type(self, type: m.FunctionType) -> Type:
-        params: TypedParamSpec = self._visit_param_spec(type.params)
         return Function(
-            pos_args=params.pos,
-            args=params.mixed,
-            kw_args=params.kw,
+            params=self._visit_param_spec(type.params),
             returns=type.returns.accept(self),
         )
 
-    def _visit_param_spec(self, spec: m.ParamSpec) -> TypedParamSpec:
+    def _visit_param_spec(self, spec: m.ParamSpec) -> ParamSpec:
         n_pos: int = len(spec.pos)
         n_mixed: int = len(spec.mixed)
 
-        def process_arg(arg: m.FunctionType.Argument, i: int) -> Function.Argument:
-            return Function.Argument(
+        def process_param(
+            param: m.FunctionType.Parameter, i: int
+        ) -> Function.Parameter:
+            return Function.Parameter(
                 pos=i,
-                name=arg.name.lexeme if arg.name is not None else str(i),
-                type=arg.type.accept(self),
-                required=arg.required,
+                name=param.name.lexeme if param.name is not None else str(i),
+                type=param.type.accept(self),
+                required=param.required,
             )
 
-        return TypedParamSpec(
-            pos=[process_arg(arg, i) for i, arg in enumerate(spec.pos)],
-            mixed=[process_arg(arg, i + n_pos) for i, arg in enumerate(spec.mixed)],
-            kw=[process_arg(arg, i + n_pos + n_mixed) for i, arg in enumerate(spec.kw)],
+        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:
@@ -426,343 +474,3 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
             self._local_variables[name] = var
             vars.append(var)
         return vars
-
-    def _get_call_result(
-        self,
-        location: Location,
-        callee: Type,
-        positional: list[TypedExpr],
-        keywords: dict[str, TypedExpr],
-        report_errors: bool = True,
-    ) -> Optional[Type]:
-        """Get the result type of a function call
-
-        If the function has overloads, the function will try to resolve the
-        appropriate signature.
-        Argument types are matched to the defined parameters.
-        The function doesn't take the raw expression as a parameter to accommodate
-        for desugared calls such as for operators.
-
-        Args:
-            location (Location): the call location
-            callee (Type): the called function
-            positional (list[TypedExpr]): the list positional arguments
-            keywords (dict[str, TypedExpr]): the map of keyword arguments
-            report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
-
-        Returns:
-            Type: the return type of the call, or `None` if either
-            the call is invalid or no overload matched the arguments uniquely
-        """
-        match callee:
-            case Function() as function:
-                valid: bool
-                mapped: list[MappedArgument]
-                valid, mapped = self.map_call_arguments(
-                    function, location, positional, keywords
-                )
-                valid = valid and self._are_arguments_valid(mapped, report_errors)
-                if not valid:
-                    return None
-                return function.returns
-
-            case OverloadedFunction(overloads=overloads):
-                function = self._match_overload(
-                    overloads, location, positional, keywords, report_errors
-                )
-                if function is None:
-                    return None
-                return function.returns
-
-            case AppliedType(body=body):
-                return self._get_call_result(
-                    location, body, positional, keywords, report_errors
-                )
-
-            case UnknownType():
-                return UnknownType()
-
-            case _:
-                if report_errors:
-                    self.reporter.error(location, f"{callee} is not callable")
-                return None
-
-    def _are_arguments_valid(
-        self,
-        arguments: list[MappedArgument],
-        report_errors: bool = True,
-    ) -> bool:
-        """Check whether the passed argument types correspond to their matched parameter definitions
-
-        Args:
-            arguments (list[MappedArgument]): the list of argument/parameter pairs
-            report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
-
-        Returns:
-            bool: True if all arguments fit the matching parameter definitions, False otherwise
-        """
-        valid: bool = True
-        for arg in arguments:
-            if not self.types.is_subtype(arg.type, arg.argument.type):
-                if report_errors:
-                    self.reporter.error(
-                        arg.expr.location,
-                        f"Wrong type for argument '{arg.argument.name}', expected {arg.argument.type}, got {arg.type}",
-                    )
-                valid = False
-        return valid
-
-    def _match_overload(
-        self,
-        overloads: list[Type],
-        location: Location,
-        positional: list[TypedExpr],
-        keywords: dict[str, TypedExpr],
-        report_errors: bool = True,
-    ) -> Optional[Function]:
-        """Try and resolve the appropriate overload for the given arguments
-
-        Args:
-            overloads (list[Type]): the list of possible overloads
-            location (Location): the call location
-            positional (list[TypedExpr]): the list of positional arguments
-            keywords (dict[str, TypedExpr]): the map of keywords arguments
-            report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
-
-        Returns:
-            Optional[Function]: the resolved function signature if it can be
-            determined unambiguously, or `None`.
-        """
-        candidates: list[OverloadCandidate] = []
-        for overload in overloads:
-            function: Type = unfold_type(overload)
-            if not isinstance(function, Function):
-                if report_errors:
-                    self.logger.error(
-                        f"Overload is not a function: {overload} is {function}"
-                    )
-                continue
-            valid, mapped = self.map_call_arguments(
-                function=function,
-                location=location,
-                positional=positional,
-                keywords=keywords,
-                report_errors=False,
-            )
-            if valid and self._are_arguments_valid(mapped, report_errors=False):
-                candidates.append(
-                    OverloadCandidate(
-                        function=function,
-                        mapped=mapped,
-                    )
-                )
-
-        pos_types: str = ", ".join(str(type) for _, type in positional)
-        kw_types: str = ", ".join(
-            f"{name}: {type}" for name, (_, type) in keywords.items()
-        )
-        for_args: str = f"for arguments pos=[{pos_types}] and kw={{{kw_types}}}"
-
-        n_candidates: int = len(candidates)
-
-        # Exactly 1 match -> return it
-        if n_candidates == 1:
-            return candidates[0].function
-
-        # No match -> invalid call
-        if n_candidates == 0:
-            overloads_str: str = ", ".join(map(str, overloads))
-            if report_errors:
-                self.reporter.error(
-                    location,
-                    f"No matching overload in [{overloads_str}] {for_args}",
-                )
-            return None
-
-        # Multiple matches -> see if one <: all others (more specific)
-        for i1, c1 in enumerate(candidates):
-            mapped1: list[MappedArgument] = c1.mapped
-            best_match: bool = True
-            for i2, c2 in enumerate(candidates):
-                if i1 == i2:
-                    continue
-                mapped2: list[MappedArgument] = c2.mapped
-                if not self._are_mapped_subtypes(mapped1, mapped2):
-                    best_match = False
-                    break
-                self.logger.debug(f"{c1.function} is a full overload of {c2.function}")
-            if best_match:
-                return c1.function
-
-        candidates_str: str = ", ".join(
-            str(candidate.function) for candidate in candidates
-        )
-        if report_errors:
-            self.reporter.error(
-                location,
-                f"Multiple matching overloads {for_args}: {candidates_str}",
-            )
-        return None
-
-    def map_call_arguments(
-        self,
-        function: Function,
-        location: Location,
-        positional: list[TypedExpr],
-        keywords: dict[str, TypedExpr],
-        report_errors: bool = True,
-    ) -> tuple[bool, list[MappedArgument]]:
-        """Map call arguments to a function's parameters as defined in its signature
-
-        This method maps positional-only, keyword-only and mixed parameter definitions
-        with the arguments passed at the call site
-
-        Any mismatched, missing or unexpected argument is reported as a diagnostic,
-        unless `report_errors` is set to `False`
-
-        Args:
-            function (Function): the function definition
-            location (Location): the call location
-            positional (list[TypedExpr]): the list of positional arguments
-            keywords (dict[str, TypedExpr]): the map of keyword arguments
-            report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
-
-        Returns:
-            tuple[bool, list[MappedArgument]]: a boolean reporting whether
-            the call is valid and the list of mapped arguments
-        """
-        set_args: set[str] = set()
-
-        required_positional: list[str] = [
-            arg.name for arg in function.pos_args + function.args if arg.required
-        ]
-        required_keyword: list[str] = [
-            arg.name for arg in function.kw_args if arg.required
-        ]
-
-        mapped: list[MappedArgument] = []
-
-        pos_params: list[Function.Argument] = list(function.pos_args)
-        mixed_params: list[Function.Argument] = list(function.args)
-        kw_params: dict[str, Function.Argument] = {
-            arg.name: arg for arg in function.kw_args
-        }
-
-        valid_call: bool = True
-
-        # TODO: handle *args and **kwargs sinks
-        for arg in positional:
-            param: Function.Argument
-            if len(pos_params) != 0:
-                param = pos_params.pop(0)
-            elif len(mixed_params) != 0:
-                param = mixed_params.pop(0)
-            else:
-                if report_errors:
-                    self.reporter.error(
-                        arg[0].location, "Too many positional arguments"
-                    )
-                valid_call = False
-                break
-            name: str = param.name
-            if name in required_positional:
-                required_positional.remove(name)
-            if name in required_keyword:
-                required_keyword.remove(name)
-            set_args.add(name)
-            mapped.append(
-                MappedArgument(
-                    expr=arg[0],
-                    type=arg[1],
-                    argument=param,
-                )
-            )
-
-        kw_params.update({arg.name: arg for arg in mixed_params})
-        for name, arg in keywords.items():
-            param: Function.Argument
-            if name not in kw_params:
-                if report_errors:
-                    if name in set_args:
-                        self.reporter.error(
-                            arg[0].location, f"Multiple values for argument '{name}'"
-                        )
-                    else:
-                        self.reporter.error(
-                            arg[0].location, f"Unknown keyword argument '{name}'"
-                        )
-                valid_call = False
-                continue
-            param = kw_params.pop(name)
-            if name in required_positional:
-                required_positional.remove(name)
-            if name in required_keyword:
-                required_keyword.remove(name)
-            set_args.add(name)
-            mapped.append(
-                MappedArgument(
-                    expr=arg[0],
-                    type=arg[1],
-                    argument=param,
-                )
-            )
-
-        def join_args(args: list[str]) -> str:
-            args = list(map(lambda a: f"'{a}'", args))
-            if len(args) == 0:
-                return ""
-            if len(args) == 1:
-                return args[0]
-            return ", ".join(args[:-1]) + " and " + args[-1]
-
-        if len(required_positional) != 0:
-            plural: str = "" if len(required_positional) == 1 else "s"
-            args: str = join_args(required_positional)
-            if report_errors:
-                self.reporter.error(
-                    location,
-                    f"Missing required positional argument{plural}: {args}",
-                )
-            valid_call = False
-
-        if len(required_keyword) != 0:
-            plural: str = "" if len(required_keyword) == 1 else "s"
-            args: str = join_args(required_keyword)
-            if report_errors:
-                self.reporter.error(
-                    location,
-                    f"Missing required keyword argument{plural}: {args}",
-                )
-            valid_call = False
-
-        return valid_call, mapped
-
-    def _are_mapped_subtypes(
-        self, mapped1: list[MappedArgument], mapped2: list[MappedArgument]
-    ) -> bool:
-        """Check whether the given argument mappings are subtype/supertype of one another
-
-        This function checks whether the argument mappings `mapped1` are subtypes
-        of `mapped2`. If any of the parameter type in `mapped1` is not a subtype
-        of the corresponding parameter in `mapped2`, `False` is returned.
-
-        This is used to check whether a given overload is
-        a more specific function/ a subtype of another.
-
-        Args:
-            mapped1 (list[MappedArgument]): the first argument mappings (subtype)
-            mapped2 (list[MappedArgument]): the second argument mappings (supertype)
-
-        Returns:
-            bool: `True` if `mapped1` is a subtype of `mapped2`, `False` otherwise
-        """
-        by_expr: dict[m.Expr, Type] = {}
-        for arg in mapped1:
-            by_expr[arg.expr] = arg.argument.type
-
-        for arg in mapped2:
-            type2: Type = arg.argument.type
-            type1: Type = by_expr[arg.expr]
-            if not self.types.is_subtype(type1, type2):
-                return False
-        return True

midas/checker/operators.py

@@ -41,7 +41,7 @@ PY_UNARY_METHODS: dict[Type[ast.unaryop], str] = {
 
 
 MIDAS_BINARY_METHODS: dict[TokenType, str] = {
-    # TokenType.PLUS: "__add__",
+    TokenType.PLUS: "__add__",
     TokenType.MINUS: "__sub__",
     TokenType.STAR: "__mul__",
     TokenType.SLASH: "__truediv__",

midas/checker/preamble.py

@@ -1,9 +1,18 @@
 from dataclasses import dataclass
-from typing import Callable, Optional
+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, TopType, Type, TypeVar, UnitType
+from midas.checker.types import (
+    Function,
+    GenericType,
+    OverloadedFunction,
+    ParamSpec,
+    TopType,
+    Type,
+    TypeVar,
+    UnitType,
+)
 
 
 @dataclass(frozen=True)
@@ -14,11 +23,14 @@ class Param:
 
 
 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] = {}
+        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)
@@ -34,7 +46,7 @@ class Preamble(Environment):
         # TODO: use sink
         self._def_function(
             name="print",
-            pos=[Param("object", TopType())],
+            pos=[Param("object", TopType(), required=False)],
             returns=UnitType(),
             py_function=print,
         )
@@ -64,11 +76,48 @@ class Preamble(Environment):
             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"),
+        )
 
-    def _list_of(self, item_type: Type) -> Type:
-        return self._types.apply_generic(self._types.get_type("list"), [item_type])
+        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 _def_type_constructor(self, name: str, py_function: Optional[Callable] = None):
+    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,
@@ -87,9 +136,9 @@ class Preamble(Environment):
         returns: Type = UnitType(),
         type_vars: list[TypeVar] = [],
     ) -> Type:
-        def map_args(params: list[Param], offset: int) -> list[Function.Argument]:
+        def map_params(params: list[Param], offset: int) -> list[Function.Parameter]:
             return [
-                Function.Argument(
+                Function.Parameter(
                     pos=i + offset,
                     name=param.name,
                     type=param.type,
@@ -99,9 +148,11 @@ class Preamble(Environment):
             ]
 
         function = Function(
-            pos_args=map_args(pos, 0),
-            args=map_args(mixed, len(pos)),
-            kw_args=map_args(kw, len(pos) + len(mixed)),
+            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:
@@ -121,7 +172,7 @@ class Preamble(Environment):
         kw: list[Param] = [],
         returns: Type = UnitType(),
         type_vars: list[TypeVar] = [],
-        py_function: Optional[Callable] = None,
+        py_function: Optional[Callable[..., Any]] = None,
     ):
         function: Type = self._make_function(
             name=name,
@@ -135,5 +186,31 @@ class Preamble(Environment):
         if py_function is not None:
             self._python_funcs[name] = py_function
 
-    def get_py_func(self, name: str) -> Optional[Callable]:
+    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

@@ -5,19 +5,18 @@ from typing import Optional
 from midas.ast.midas import MemberKind
 from midas.checker.builtins import BUILTIN_SUBTYPES
 from midas.checker.types import (
-    AliasType,
     AppliedType,
     BaseType,
     ColumnType,
-    ComplexType,
     ConstraintType,
     DataFrameType,
-    ExtensionType,
+    DerivedType,
     Function,
     GenericType,
     OverloadedFunction,
     Predicate,
     TopType,
+    TupleType,
     Type,
     TypeVar,
     UnknownType,
@@ -28,11 +27,15 @@ from midas.checker.types import (
 
 @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] = {}
@@ -80,6 +83,25 @@ class TypesRegistry:
         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]
@@ -108,10 +130,37 @@ class TypesRegistry:
             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`
 
@@ -145,19 +194,11 @@ class TypesRegistry:
                     return True
                 return self.is_subtype(type1, bound)
 
-            case (AliasType(type=base1), _):
+            case (DerivedType(type=base1), _):
                 return self.is_subtype(base1, type2)
 
             case (BaseType(name=name1), BaseType(name=name2)):
-                return name1 in BUILTIN_SUBTYPES.get(name2, set())
-
-            case (ComplexType(properties=props1), ComplexType(properties=props2)):
-                for k, t in props2.items():
-                    if k not in props1:
-                        return False
-                    if not self.is_subtype(props1[k], t):
-                        return False
-                return True
+                return self.is_builtin_subtype(name1, name2)
 
             case (DataFrameType(columns=columns1), DataFrameType(columns=columns2)):
                 # TODO: check order?
@@ -208,6 +249,15 @@ class TypesRegistry:
         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
@@ -224,101 +274,121 @@ class TypesRegistry:
         if not self.is_subtype(func1.returns, func2.returns):
             return False
 
-        pos1: list[Function.Argument] = func1.pos_args
-        mixed1: list[Function.Argument] = func1.args
-        kw1: dict[str, Function.Argument] = {a.name: a for a in func1.kw_args}
-        pos2: list[Function.Argument] = func2.pos_args
-        mixed2: list[Function.Argument] = func2.args
-        kw2: dict[str, Function.Argument] = {a.name: a for a in func2.kw_args}
+        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.Argument] = {arg.pos: arg for arg in mixed2}
-        mixed_by_name: dict[str, Function.Argument] = {arg.name: arg for arg in mixed2}
+        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.Argument, sup: Function.Argument) -> bool:
+        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 arg1 in pos1:
-            arg2: Function.Argument
-            if arg1.pos < len(pos2):
-                arg2 = pos2[arg1.pos]
-            elif arg1.pos in mixed_by_pos:
-                arg2 = mixed_by_pos[arg1.pos]
-            elif not arg1.required:
+        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(arg2, arg1):
+            if not is_arg_subtype(param2, param1):
                 return False
 
-        for name, arg1 in kw1.items():
-            arg2: Function.Argument
+        for name, param1 in kw1.items():
+            param2: Function.Parameter
             if name in kw2:
-                arg2 = kw2[name]
+                param2 = kw2[name]
             elif name in mixed_by_name:
-                arg2 = mixed_by_name[name]
-            elif not arg1.required:
+                param2 = mixed_by_name[name]
+            elif not param1.required:
                 continue
             else:
                 return False
-            if not is_arg_subtype(arg2, arg1):
+            if not is_arg_subtype(param2, param1):
                 return False
 
-        for arg1 in mixed1:
-            pos_arg2: Optional[Function.Argument] = None
-            kw_arg2: Optional[Function.Argument] = None
-            if arg1.name in kw2:
-                kw_arg2 = kw2[arg1.name]
-            elif arg1.name in mixed_by_name:
-                kw_arg2 = mixed_by_name[arg1.name]
-            if arg1.pos < len(pos2):
-                pos_arg2 = pos2[arg1.pos]
-            elif arg1.pos in mixed_by_pos:
-                pos_arg2 = mixed_by_pos[arg1.pos]
+        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_arg2 is None and kw_arg2 is None and arg1.required:
+            if pos_param2 is None and kw_param2 is None and param1.required:
                 return False
 
             # Matching keyword argument
-            if kw_arg2 is not None and not is_arg_subtype(kw_arg2, arg1):
+            if kw_param2 is not None and not is_arg_subtype(kw_param2, param1):
                 return False
 
             # Matching positional argument
-            if pos_arg2 is not None and not is_arg_subtype(pos_arg2, arg1):
+            if pos_param2 is not None and not is_arg_subtype(pos_param2, param1):
                 return False
 
-        mixed_positions: set[int] = {a.pos for a in mixed1}
-        mixed_names: set[str] = {a.name for a in mixed1}
-        for arg2 in pos2:
-            if not arg2.required:
+        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 arg2.pos >= len(pos1) and arg2.pos not in mixed_positions:
+            if param2.pos >= len(pos1) and param2.pos not in mixed_positions:
                 return False
 
-        for name, arg2 in kw2.items():
-            if not arg2.required:
+        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 arg2 in mixed2:
-            if arg2.required:
+        for param2 in mixed2:
+            if param2.required:
                 continue
-            pos_match: bool = arg2.pos < len(pos1) or arg2.pos in mixed_positions
-            kw_match: bool = arg2.name in kw1 or arg2.name in mixed_names
+            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 AliasType(name=name, type=base):
-                return AliasType(name=name, type=self.apply_generic(base, args))
+            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)
@@ -346,6 +416,9 @@ class TypesRegistry:
                     body=substitute_typevars(body, substitutions),
                 )
 
+            case BaseType(name="tuple"):
+                return TupleType(items=tuple(args))
+
             case _:
                 raise ValueError(f"{type} is not a generic type")
 
@@ -378,6 +451,19 @@ class TypesRegistry:
         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:
@@ -385,7 +471,7 @@ class TypesRegistry:
                         return self._members[name][member_name].type
                 return None
 
-            case AliasType(name=name, type=base):
+            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
@@ -410,20 +496,6 @@ class TypesRegistry:
                     member_type2 = substitute_typevars(member_type2, substitutions)
                 return member_type2
 
-            case ComplexType(members=members):
-                if member_name in members:
-                    return members[member_name]
-                self.logger.debug(f"No member '{member_name}' in {type}")
-                return None
-
-            case ExtensionType(base=base, extension=ComplexType(members=members)):
-                if member_name in members:
-                    return members[member_name]
-                self.logger.debug(
-                    f"No member '{member_name}' on {type}, looking up in base"
-                )
-                return self.lookup_member(base, member_name)
-
             case ConstraintType(type=base):
                 return self.lookup_member(base, member_name)
 
@@ -438,4 +510,75 @@ class TypesRegistry:
                 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

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from contextlib import contextmanager
 from typing import Optional
 
 from midas.ast.location import Location
@@ -7,6 +8,8 @@ from midas.checker.diagnostic import Diagnostic, DiagnosticType
 
 
 class Reporter:
+    """Helper class to store diagnostics"""
+
     def __init__(self):
         self.diagnostics: list[Diagnostic] = []
 
@@ -17,6 +20,14 @@ class Reporter:
         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,
@@ -27,21 +38,63 @@ class Reporter:
         )
 
     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,
@@ -49,6 +102,12 @@ class FileReporter:
         )
 
     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,
@@ -56,6 +115,12 @@ class FileReporter:
         )
 
     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,
@@ -63,6 +128,12 @@ class FileReporter:
         )
 
     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,

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.reporter: FileReporter = reporter
 
     def resolve(self, *objects: p.Stmt | p.Expr) -> None:
         """Resolve the given statements or expressions"""
@@ -25,29 +28,29 @@ 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",
             )
-        scope[name] = False
+        else:
+            scope[name] = False
 
     def define(self, name: str) -> None:
         """Define a variable in the current scope
@@ -77,7 +80,15 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
                 self.locals[expr] = i
                 return
 
-    def is_defined(self, name: str) -> bool:
+    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
@@ -93,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()
@@ -104,14 +119,12 @@ 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)
@@ -121,9 +134,9 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
     def _visit_assign(self, target: p.Expr):
         match target:
             case p.VariableExpr(name=name):
-                if not self.is_defined(name):
-                    self.declare(name)
-                    self.define(name)
+                if not self.is_declared(name):
+                    self.declare(target.location, name)
+                self.define(name)
                 target.accept(self)
 
             case p.GetExpr():
@@ -133,7 +146,9 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
                 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:
@@ -150,23 +165,40 @@ 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)
-        self.end_scope()
+        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.begin_scope()
         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
 
@@ -196,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)
 
@@ -236,5 +269,9 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
         if expr.step is not None:
             self.resolve(expr.step)
 
+    def visit_tuple_expr(self, expr: p.TupleExpr) -> None:
+        for item in expr.items:
+            self.resolve(item)
+
     def visit_raw_expr(self, expr: p.RawExpr) -> None:
         pass

midas/checker/types.py

@@ -10,12 +10,16 @@ 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:
@@ -23,7 +27,9 @@ class BaseType:
 
 
 @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
 
@@ -33,77 +39,83 @@ class AliasType:
 
 @dataclass(frozen=True, kw_only=True)
 class UnknownType:
+    """An unknown type"""
+
     def __str__(self) -> str:
         return "<Unknown>"
 
 
 @dataclass(frozen=True, kw_only=True)
 class UnitType:
+    """The unit type (`None`)"""
+
     def __str__(self) -> str:
         return "None"
 
 
 @dataclass(frozen=True, kw_only=True)
 class Function:
-    pos_args: list[Argument] = field(default_factory=list)
-    args: list[Argument] = field(default_factory=list)
-    kw_args: list[Argument] = field(default_factory=list)
+    """A function type"""
+
+    params: ParamSpec
     returns: Type
 
     def __str__(self) -> str:
-        args: list[str] = []
-        if len(self.pos_args) != 0:
-            args += list(map(str, self.pos_args))
-            args.append("/")
-
-        if len(self.args) != 0:
-            args += list(map(str, self.args))
-
-        if len(self.kw_args) != 0:
-            args.append("*")
-            args += list(map(str, self.kw_args))
-
-        return f"({', '.join(args)}) -> {self.returns}"
+        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 "?"
-            return f"{self.name}: {self.type}{opt}"
+            param: str = f"{self.name}: {self.type}{opt}"
+            if self.unsupported:
+                param = f"({param})"
+            return param
+
+
+@dataclass(frozen=True, kw_only=True)
+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 OverloadedFunction:
+    """A list of method overloads"""
+
     overloads: list[Type]
 
     def __str__(self) -> str:
         return "<overloaded function>"
 
 
-@dataclass(frozen=True, kw_only=True)
-class ComplexType:
-    members: dict[str, Type]
-
-    def __str__(self) -> str:
-        props: list[str] = [f"{name}: {type}" for name, type in self.members.items()]
-        return f"{{{', '.join(props)}}}"
-
-
-@dataclass(frozen=True, kw_only=True)
-class ExtensionType:
-    base: Type
-    extension: ComplexType
-
-    def __str__(self) -> str:
-        return f"{self.base} & {self.extension}"
-
-
 class Variance(StrEnum):
+    """The variance of a :class:`TypeVar`"""
+
     INVARIANT = "INVARIANT"
     COVARIANT = "COVARIANT"
     CONTRAVARIANT = "CONTRAVARIANT"
@@ -111,6 +123,8 @@ class Variance(StrEnum):
 
 @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
@@ -128,6 +142,8 @@ class TypeVar:
 
 @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
@@ -138,6 +154,8 @@ class GenericType:
 
 @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
@@ -148,6 +166,8 @@ class AppliedType:
 
 @dataclass(frozen=True, kw_only=True)
 class ConstraintType:
+    """A type with a constraint expression"""
+
     type: Type
     constraint: m.Expr
 
@@ -158,6 +178,8 @@ class ConstraintType:
 
 @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:
@@ -166,6 +188,8 @@ class TupleType:
 
 @dataclass(frozen=True, kw_only=True)
 class ColumnType:
+    """A column type containing items of a given unique type"""
+
     type: Type
 
     def __str__(self) -> str:
@@ -174,6 +198,8 @@ class ColumnType:
 
 @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:
@@ -187,13 +213,53 @@ class DataFrameType:
         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:
-    def sub_argument(arg: Function.Argument):
-        return Function.Argument(
-            pos=arg.pos,
-            name=arg.name,
-            type=substitute_typevars(arg.type, substitutions),
-            required=arg.required,
+    """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):
@@ -213,19 +279,17 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
         case BaseType():
             return type
 
-        case AliasType(name=name, type=type2):
-            return AliasType(name=name, type=substitute_typevars(type2, substitutions))
+        case DerivedType(name=name, type=type2):
+            return DerivedType(
+                name=name, type=substitute_typevars(type2, substitutions)
+            )
 
         case Function(
-            pos_args=pos_args,
-            args=args,
-            kw_args=kw_args,
+            params=params,
             returns=returns,
         ):
             return Function(
-                pos_args=list(map(sub_argument, pos_args)),
-                args=list(map(sub_argument, args)),
-                kw_args=list(map(sub_argument, kw_args)),
+                params=sub_param_spec(params),
                 returns=substitute_typevars(returns, substitutions),
             )
 
@@ -237,24 +301,6 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
                 ]
             )
 
-        case ComplexType(members=members):
-            members2: dict[str, Type] = {
-                name: substitute_typevars(prop, substitutions)
-                for name, prop in members.items()
-            }
-            return ComplexType(members=members2)
-
-        case ExtensionType(base=base, extension=ComplexType(members=members)):
-            return ExtensionType(
-                base=substitute_typevars(base, substitutions),
-                extension=ComplexType(
-                    members={
-                        name: substitute_typevars(prop, substitutions)
-                        for name, prop in members.items()
-                    }
-                ),
-            )
-
         case AppliedType(name=name, args=args, body=body):
             return AppliedType(
                 name=name,
@@ -303,11 +349,20 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
                 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
@@ -316,22 +371,39 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> 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 AliasType(type=ref_type):
+        case DerivedType(type=ref_type):
             return unfold_type(ref_type)
         case _:
             return type
 
 
 def to_annotation(type: Type) -> str:
-    def _args_annotation(func: Function) -> str:
-        if len(func.kw_args) != 0:
+    """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 "..."
 
-        args: str = ", ".join(
-            to_annotation(arg.type) for arg in func.pos_args + func.args
+        params: str = ", ".join(
+            to_annotation(param.type) for param in spec.pos + spec.mixed
         )
-        return f"[{args}]"
+        return f"[{params}]"
 
     match type:
         case TopType():
@@ -340,7 +412,7 @@ def to_annotation(type: Type) -> str:
         case BaseType(name=name):
             return name
 
-        case AliasType(name=name):
+        case DerivedType(name=name):
             return name
 
         case UnknownType():
@@ -349,16 +421,13 @@ def to_annotation(type: Type) -> str:
         case UnitType():
             return "None"
 
-        case Function(returns=returns):
-            params_annot: str = _args_annotation(type)
+        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 ComplexType() | ExtensionType():
-            raise NotImplementedError
-
         case TypeVar(name=name):
             return name
 
@@ -380,12 +449,20 @@ def to_annotation(type: Type) -> str:
         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
@@ -394,13 +471,11 @@ class Predicate:
 Type = (
     TopType
     | BaseType
-    | AliasType
+    | DerivedType
     | UnknownType
     | UnitType
     | Function
     | OverloadedFunction
-    | ComplexType
-    | ExtensionType
     | TypeVar
     | GenericType
     | AppliedType
@@ -408,4 +483,6 @@ Type = (
     | TupleType
     | ColumnType
     | DataFrameType
+    | FrameGroupBy
+    | ColumnGroupBy
 )

midas/checker/unifier.py

@@ -4,8 +4,11 @@ from typing import Optional
 from midas.checker.registry import TypesRegistry
 from midas.checker.types import (
     AppliedType,
+    ColumnType,
+    DataFrameType,
     Function,
     GenericType,
+    ParamSpec,
     TopType,
     Type,
     TypeVar,
@@ -16,6 +19,14 @@ 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")
@@ -26,26 +37,37 @@ class Unifier:
         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(
-            pos_args=[
-                Function.Argument(
-                    pos=i,
-                    name=str(i),
-                    type=arg,
-                    required=True,
-                )
-                for i, arg in enumerate(positional)
-            ],
-            args=[],
-            kw_args=[
-                Function.Argument(
-                    pos=len(positional) + i,
-                    name=name,
-                    type=arg,
-                    required=True,
-                )
-                for i, (name, arg) in enumerate(keywords.items())
-            ],
+            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)
@@ -56,6 +78,18 @@ class Unifier:
         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)
@@ -77,6 +111,22 @@ class Unifier:
         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):
@@ -98,8 +148,32 @@ class Unifier:
 
                 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.Argument, Function.Argument]] = (
+                mapped: list[tuple[Function.Parameter, Function.Parameter]] = (
                     self.map_params(template, concrete)
                 )
                 substitutions: dict[str, Type] = {}
@@ -122,6 +196,18 @@ class Unifier:
                 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():
@@ -135,19 +221,32 @@ class Unifier:
 
     def map_params(
         self, func1: Function, func2: Function
-    ) -> list[tuple[Function.Argument, Function.Argument]]:
-        pos1: list[Function.Argument] = func1.pos_args
-        mixed1: list[Function.Argument] = func1.args
-        kw1: list[Function.Argument] = func1.kw_args
+    ) -> list[tuple[Function.Parameter, Function.Parameter]]:
+        """Map parameters of two functions
 
-        pos2: list[Function.Argument] = func2.pos_args
-        mixed2: list[Function.Argument] = func2.args
-        kw2: list[Function.Argument] = func2.kw_args
+        Args:
+            func1 (Function): the first function
+            func2 (Function): the second function
 
-        mapped: list[tuple[Function.Argument, Function.Argument]] = []
+        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
 
-        by_pos2: dict[int, Function.Argument] = {arg.pos: arg for arg in pos2 + mixed2}
-        by_name2: dict[str, Function.Argument] = {arg.name: arg for arg in mixed2 + kw2}
+        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:

midas/checker/variance.py

@@ -16,14 +16,27 @@ 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)
@@ -32,6 +45,18 @@ class Tracker:
         ]
 
     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
@@ -46,11 +71,22 @@ class Tracker:
 
 
 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)
@@ -71,20 +107,36 @@ class VarianceInferrer:
         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(pos_args=pos_args, args=mixed_args, kw_args=kw_args):
-                all_args: list[Function.Argument] = pos_args + mixed_args + kw_args
-                for arg in all_args:
+            case Function(params=spec):
+                all_params: list[Function.Parameter] = spec.pos + spec.mixed + spec.kw
+                for param in all_params:
                     self.walk(
-                        arg.type,
+                        param.type,
                         -polarity,
                         base_name,
-                        path + [f"arg:'{arg.name}'"],
+                        path + [f"param:'{param.name}'"],
                     )
 
                 self.walk(type.returns, polarity, base_name, path + ["return"])

midas/cli/commands/compile.py

@@ -5,7 +5,7 @@
 
 import sys
 from pathlib import Path
-from typing import TextIO
+from typing import Optional, TextIO
 
 import click
 
@@ -19,18 +19,23 @@ 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()
-    for types_file in types:
-        checker.import_midas(Path(types_file.name).resolve())
+    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()
@@ -43,4 +48,4 @@ def compile(
         sys.exit(1)
 
     generator = Generator(workdir=source_path.parent, types=checker.types)
-    generator.generate(typed_ast, source_path)
+    generator.generate(typed_ast, source_path, type_files=type_files)

midas/cli/commands/registry.py

@@ -11,14 +11,14 @@ 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 AliasType, AppliedType, BaseType, GenericType, Type
+from midas.checker.types import AppliedType, BaseType, DerivedType, GenericType, Type
 
 
 def base_type(type: Type) -> Type:
     match type:
         case BaseType():
             return type
-        case AliasType(type=base):
+        case DerivedType(type=base):
             return base
         case AppliedType(body=body):
             return body

midas/cli/commands/stubs.py

@@ -1,7 +1,7 @@
 import ast
 import time
 from pathlib import Path
-from typing import TextIO
+from typing import Optional, TextIO
 
 import black
 import click
@@ -38,15 +38,17 @@ class Handler(FileSystemEventHandler):
 
 @click.command(help="Generate stubs from Midas definitions")
 @click.argument("file", type=click.File("r"))
-@click.option("-o", "--output", type=click.File("w"), default="-")
+@click.option("-o", "--output", type=click.File("w"))
 @click.option("-w", "--watch", is_flag=True)
 def stubs(
     file: TextIO,
-    output: TextIO,
+    output: Optional[TextIO],
     watch: bool,
 ):
     source_path: Path = Path(file.name).resolve()
-    out_path: Path = Path(output.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:

midas/cli/highlighter.py

@@ -134,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")
@@ -157,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)
@@ -247,6 +250,10 @@ class PythonHighlighter(
         if expr.step is not None:
             expr.step.accept(self)
 
+    def visit_tuple_expr(self, expr: p.TupleExpr) -> None:
+        for item in expr.items:
+            item.accept(self)
+
     def visit_raw_expr(self, expr: p.RawExpr) -> None: ...
 
     def visit_raw_stmt(self, stmt: p.RawStmt) -> None: ...
@@ -331,21 +338,11 @@ class MidasHighlighter(
         type.type.accept(self)
         type.constraint.accept(self)
 
-    def visit_complex_type(self, type: m.ComplexType) -> None:
-        self.wrap(type, "complex-type")
-        for member in type.members:
-            member.accept(self)
-
     def visit_function_type(self, type: m.FunctionType) -> None:
         self.wrap(type, "function")
         self._visit_param_spec(type.params)
         type.returns.accept(self)
 
-    def visit_extension_type(self, type: m.ExtensionType) -> None:
-        self.wrap(type, "extension")
-        type.base.accept(self)
-        type.extension.accept(self)
-
     def _visit_param_spec(self, spec: m.ParamSpec) -> None:
         for param in spec.pos + spec.mixed + spec.kw:
             param.type.accept(self)

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

@@ -7,8 +7,7 @@ span {
 
     &.named-type,
     &.generic-type,
-    &.constraint-type,
-    &.complex-type {
+    &.constraint-type {
         --col: 150, 150, 150;
     }
 

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/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

@@ -5,6 +5,7 @@ import midas.ast.midas as m
 from midas.checker.registry import TypesRegistry
 from midas.checker.types import (
     Function,
+    ParamSpec,
     Predicate,
     Type,
     to_annotation,
@@ -17,14 +18,14 @@ LOGICAL_OPERATORS: dict[TokenType, type[ast.boolop]] = {
 }
 
 BINARY_OPERATORS: dict[TokenType, type[ast.operator]] = {
-    # TokenType.PLUS: ast.Add,
+    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.PLUS: ast.UAdd,
     TokenType.MINUS: ast.USub,
 }
 
@@ -39,6 +40,8 @@ COMPARISON_OPERATORS: dict[TokenType, type[ast.cmpop]] = {
 
 
 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
@@ -46,24 +49,37 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
         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(
-                    pos_args=[],
-                    args=[
-                        Function.Argument(
-                            pos=0,
-                            name="_",
-                            type=self.types.get_type("Any"),
-                            required=True,
-                        )
-                    ],
-                    kw_args=[],
+                    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)
@@ -74,6 +90,14 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
                 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}"
@@ -84,6 +108,15 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
         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(
@@ -94,28 +127,36 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
             )
         return self.make_func(name, [ast.Return(value=body)], predicate.type)
 
-    def make_args(self, func: Function) -> ast.arguments:
+    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=arg.name,
-                    annotation=ast.Constant(value=to_annotation(arg.type)),
+                    arg=param.name,
+                    annotation=ast.Constant(value=to_annotation(param.type)),
                 )
-                for arg in func.pos_args
+                for param in params.pos
             ],
             args=[
                 ast.arg(
-                    arg=arg.name,
-                    annotation=ast.Constant(value=to_annotation(arg.type)),
+                    arg=param.name,
+                    annotation=ast.Constant(value=to_annotation(param.type)),
                 )
-                for arg in func.args
+                for param in params.mixed
             ],
             kwonlyargs=[
                 ast.arg(
-                    arg=arg.name,
-                    annotation=ast.Constant(value=to_annotation(arg.type)),
+                    arg=param.name,
+                    annotation=ast.Constant(value=to_annotation(param.type)),
                 )
-                for arg in func.kw_args
+                for param in params.kw
             ],
             defaults=[],
             kw_defaults=[],
@@ -124,12 +165,39 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
     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(returns=Function()):
+            case Function(params=params, returns=Function()):
                 inner_name: str = f"inner{level}"
                 return ast.FunctionDef(
                     name=name,
-                    args=self.make_args(type),
+                    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)),
@@ -138,10 +206,10 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
                     decorator_list=[],
                 )
 
-            case Function():
+            case Function(params=params):
                 return ast.FunctionDef(
                     name=name,
-                    args=self.make_args(type),
+                    args=self.make_args(params),
                     body=inner_body,
                     returns=ast.Constant(value=to_annotation(type.returns)),
                     decorator_list=[],
@@ -151,6 +219,18 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
                 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:

midas/generator/generator.py

@@ -9,16 +9,17 @@ 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 (
-    AliasType,
     AppliedType,
     BaseType,
+    ColumnGroupBy,
     ColumnType,
-    ComplexType,
     ConstraintType,
     DataFrameType,
-    ExtensionType,
+    DerivedType,
+    FrameGroupBy,
     Function,
     GenericType,
     OverloadedFunction,
@@ -29,17 +30,31 @@ from midas.checker.types import (
     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__"
 
@@ -53,10 +68,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
             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]] = []
@@ -64,11 +81,29 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         self.define_is_dataframe: bool = False
         self.define_is_column: bool = False
 
-    def generate_ast(self, typed_ast: TypedAST, src_path: Path) -> ast.AST:
-        self.rel_src_path = src_path.resolve().relative_to(self.workdir)
+    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)
+        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
 
@@ -78,15 +113,43 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         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
+        self,
+        typed_ast: TypedAST,
+        src_path: Path,
+        out_path: Optional[Path] = None,
+        type_files: Optional[list[tuple[Path, Optional[str]]]] = None,
     ) -> Path:
-        module: ast.AST = self.generate_ast(typed_ast, src_path)
-        compiled: str = ast.unparse(module)
+        """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)
@@ -98,43 +161,90 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 raise ValueError(
                     f"Directory traversal, {self.rel_src_path} points outside of parent directory"
                 )
-            out_path.parent.mkdir(parents=True, exist_ok=True)
+        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=expr.left.accept(self),
+            left=self.convert(expr.left),
             op=expr.operator,
-            right=expr.right.accept(self),
+            right=self.convert(expr.right),
         )
 
     def visit_compare_expr(self, expr: p.CompareExpr) -> ast.expr:
         return ast.Compare(
-            left=expr.left.accept(self),
+            left=self.convert(expr.left),
             ops=[expr.operator],
-            comparators=[expr.right.accept(self)],
+            comparators=[self.convert(expr.right)],
         )
 
     def visit_unary_expr(self, expr: p.UnaryExpr) -> ast.expr:
         return ast.UnaryOp(
             op=expr.operator,
-            operand=expr.right.accept(self),
+            operand=self.convert(expr.right),
         )
 
     def visit_call_expr(self, expr: p.CallExpr) -> ast.expr:
         return ast.Call(
-            func=expr.callee.accept(self),
-            args=[arg.accept(self) for arg in expr.arguments],
+            func=self.convert(expr.callee),
+            args=[self.convert(arg) for arg in expr.arguments],
             keywords=[
-                ast.keyword(arg=name, value=arg.accept(self))
+                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=expr.object.accept(self),
+            value=self.convert(expr.object),
             attr=expr.name,
         )
 
@@ -147,16 +257,16 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
     def visit_logical_expr(self, expr: p.LogicalExpr) -> ast.expr:
         return ast.BoolOp(
             op=expr.operator,
-            values=[expr.left.accept(self), expr.right.accept(self)],
+            values=[self.convert(expr.left), self.convert(expr.right)],
         )
 
     def visit_cast_expr(self, expr: p.CastExpr) -> ast.expr:
-        expr2: ast.expr = expr.expr.accept(self)
+        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(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)
@@ -167,33 +277,38 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
 
     def visit_ternary_expr(self, expr: p.TernaryExpr) -> ast.expr:
         return ast.IfExp(
-            test=expr.test.accept(self),
-            body=expr.if_true.accept(self),
-            orelse=expr.if_false.accept(self),
+            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=[item.accept(self) for item in expr.items],
+            elts=[self.convert(item) for item in expr.items],
         )
 
     def visit_dict_expr(self, expr: p.DictExpr) -> ast.expr:
         return ast.Dict(
-            keys=[key.accept(self) if key is not None else None for key in expr.keys],
-            values=[value.accept(self) for value in expr.values],
+            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=expr.object.accept(self),
-            slice=expr.index.accept(self),
+            value=self.convert(expr.object),
+            slice=self.convert(expr.index),
         )
 
     def visit_slice_expr(self, expr: p.SliceExpr) -> ast.expr:
         return ast.Slice(
-            lower=expr.lower.accept(self) if expr.lower is not None else None,
-            upper=expr.upper.accept(self) if expr.upper is not None else None,
-            step=expr.step.accept(self) if expr.step is not None else None,
+            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:
@@ -201,28 +316,37 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
 
     def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> ast.stmt:
         return ast.Expr(
-            value=stmt.expr.accept(self),
+            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=ast.arguments(
-                posonlyargs=[ast.arg(arg=arg.name) for arg in stmt.posonlyargs],
-                vararg=None,
-                args=[ast.arg(arg=arg.name) for arg in stmt.args],
-                kwonlyargs=[ast.arg(arg=arg.name) for arg in stmt.kwonlyargs],
-                kwarg=None,
-                defaults=[
-                    arg.default.accept(self)
-                    for arg in stmt.posonlyargs + stmt.args
-                    if arg.default is not None
-                ],
-                kw_defaults=[
-                    arg.default.accept(self) if arg.default is not None else None
-                    for arg in stmt.kwonlyargs
-                ],
-            ),
+            args=self.make_args(stmt.params),
             body=self._visit_body(stmt.body),
             decorator_list=[],
         )
@@ -233,20 +357,20 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
 
     def visit_assign_stmt(self, stmt: p.AssignStmt) -> ast.stmt:
         return ast.Assign(
-            targets=[target.accept(self) for target in stmt.targets],
-            value=stmt.value.accept(self),
+            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=stmt.value.accept(self) if stmt.value is not None else None,
+            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=stmt.test.accept(self),
+            test=self.convert(stmt.test),
             body=self._visit_body(stmt.body),
-            orelse=self._visit_body(stmt.orelse),
+            orelse=self._visit_body(stmt.orelse, can_be_empty=True),
         )
 
     def visit_pass(self, stmt: p.Pass) -> ast.stmt:
@@ -254,16 +378,62 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
 
     def visit_for_stmt(self, stmt: p.ForStmt) -> ast.stmt:
         return ast.For(
-            target=stmt.target.accept(self),
-            iter=stmt.iterator.accept(self),
+            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]) -> list[ast.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()
@@ -281,9 +451,25 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         # 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, expr: ast.expr) -> ast.expr:
+    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
@@ -294,9 +480,19 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 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(
@@ -305,9 +501,25 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         )
 
     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
@@ -316,8 +528,19 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
     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():
+            case UnknownType() | TopType():
                 return []
 
             case BaseType(name=name):
@@ -332,7 +555,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                     )
                 ]
 
-            case AliasType(type=base):
+            case DerivedType(type=base):
                 return self._make_cast_asserts(src_location, expr, base)
 
             case UnitType():
@@ -440,16 +663,19 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                         ),
                     ),
                 ]
-                asserts.append(self._make_column_inner_assert(src_location, expr, type))
+                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 (
-                TopType()
-                | Function()
+                Function()
                 | OverloadedFunction()
-                | ComplexType()
-                | ExtensionType()
                 | GenericType()
+                | FrameGroupBy()
+                | ColumnGroupBy()
             ):
                 self.logger.warning(f"Can't make assertion for type {type}")
                 return []
@@ -465,6 +691,24 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[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(
@@ -488,6 +732,17 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
     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(
@@ -495,12 +750,22 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 args=[expr],
                 keywords=[],
             ),
-            self._make_constraint_assert_message(src_location, expr, constraint),
+            self._make_constraint_assert_message(src_location, constraint),
         )
 
     def _make_constraint_assert_message(
-        self, location: Location, expr: ast.expr, constraint: m.Expr
+        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}"
@@ -510,6 +775,14 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         )
 
     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
@@ -519,10 +792,18 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         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(
@@ -555,10 +836,18 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         )
 
     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(
@@ -592,12 +881,71 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
 
     def _make_column_inner_assert(
         self, src_location: Location, column: ast.expr, type: ColumnType
-    ) -> ast.stmt:
+    ) -> 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=self._make_cast_asserts(src_location, col, type.type),
+            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

@@ -4,17 +4,18 @@ from typing import Optional, assert_never
 import midas.ast.midas as m
 from midas.checker.registry import Member, TypesRegistry
 from midas.checker.types import (
-    AliasType,
     AppliedType,
     BaseType,
+    ColumnGroupBy,
     ColumnType,
-    ComplexType,
     ConstraintType,
     DataFrameType,
-    ExtensionType,
+    DerivedType,
+    FrameGroupBy,
     Function,
     GenericType,
     OverloadedFunction,
+    ParamSpec,
     TopType,
     TupleType,
     Type,
@@ -29,6 +30,8 @@ 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] = []
@@ -40,6 +43,11 @@ class StubsGenerator:
         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
@@ -89,8 +97,29 @@ class StubsGenerator:
         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
@@ -111,8 +140,19 @@ class StubsGenerator:
         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 AliasType(type=base):
+            case DerivedType(type=base):
                 return [self.dump_type(base)], {}
 
             case GenericType(params=params, body=body):
@@ -155,6 +195,16 @@ class StubsGenerator:
     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)]
 
@@ -176,8 +226,16 @@ class StubsGenerator:
         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 AliasType(name=name) | GenericType(name=name) if (
+            case DerivedType(name=name) | GenericType(name=name) if (
                 name in self.substitutions
             ):
                 type = substitute_typevars(type, self.substitutions[name])
@@ -190,7 +248,7 @@ class StubsGenerator:
             case BaseType(name=name):
                 return ast.Name(id=name)
 
-            case AliasType(name=name):
+            case DerivedType(name=name):
                 return ast.Name(id=name)
 
             case UnitType():
@@ -209,14 +267,6 @@ class StubsGenerator:
                     right=self.dump_type(overloads[-1]),
                 )
 
-            case ComplexType():
-                name: str = self.new_stub_name()
-                self.generate_stub(name, type)
-                return ast.Name(id=name)
-
-            case ExtensionType():
-                raise NotImplementedError
-
             case TypeVar():
                 return ast.Name(id=type.name)
 
@@ -255,14 +305,11 @@ class StubsGenerator:
                     ),
                 )
 
-            case ColumnType(type=inner):
+            case ColumnType():
                 self.import_pandas = True
-                return ast.Subscript(
-                    value=ast.Attribute(
-                        value=ast.Name(id="pd"),
-                        attr="Series",
-                    ),
-                    slice=self.dump_type(inner),
+                return ast.Attribute(
+                    value=ast.Name(id="pd"),
+                    attr="Series",
                 )
 
             case DataFrameType():
@@ -272,12 +319,49 @@ class StubsGenerator:
                     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:
@@ -285,7 +369,7 @@ class StubsGenerator:
                 return [
                     ast.FunctionDef(
                         name=name,
-                        args=self.dump_args(method, with_self=True),
+                        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 [],
@@ -305,24 +389,43 @@ class StubsGenerator:
                     )
                 ]
 
-    def dump_args(self, func: Function, with_self: bool = False) -> ast.arguments:
+    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"_{arg.pos}", annotation=self.dump_type(arg.type))
-            for arg in func.pos_args
+            ast.arg(
+                arg=f"_{param.pos}",
+                annotation=self.dump_type(param.type),
+            )
+            for param in params.pos
         ]
         mixed: list[ast.arg] = [
-            ast.arg(arg=arg.name, annotation=self.dump_type(arg.type))
-            for arg in func.args
+            ast.arg(
+                arg=param.name,
+                annotation=self.dump_type(param.type),
+            )
+            for param in params.mixed
         ]
         kw: list[ast.arg] = [
-            ast.arg(arg=arg.name, annotation=self.dump_type(arg.type))
-            for arg in func.kw_args
+            ast.arg(
+                arg=param.name,
+                annotation=self.dump_type(param.type),
+            )
+            for param in params.kw
         ]
         defaults: list[ast.expr] = [
-            Empty for arg in func.pos_args + func.args if not arg.required
+            Empty for param in params.pos + params.mixed if not param.required
         ]
         kw_defaults: list[Optional[ast.expr]] = [
-            None if arg.required else Empty for arg in func.kw_args
+            None if param.required else Empty for param in params.kw
         ]
         if with_self:
             arg = ast.arg(arg="self", annotation=None)
@@ -339,6 +442,14 @@ class StubsGenerator:
         )
 
     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(
@@ -348,7 +459,7 @@ class StubsGenerator:
             body=[
                 ast.FunctionDef(
                     name="__call__",
-                    args=self.dump_args(func, with_self=True),
+                    args=self.dump_params(func.params, with_self=True),
                     returns=self.dump_type(func.returns),
                     body=[ast.Expr(value=Empty)],
                     decorator_list=[],
@@ -360,33 +471,74 @@ class StubsGenerator:
         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")
 

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

@@ -32,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 ".":
@@ -46,8 +48,8 @@ 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 "*":
@@ -81,6 +83,12 @@ class MidasLexer(Lexer):
         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()
 
@@ -100,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):
@@ -147,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

@@ -25,7 +25,7 @@ class TokenType(Enum):
     DOT = auto()
 
     # Operators
-    # PLUS = auto()
+    PLUS = auto()
     MINUS = auto()
     STAR = auto()
     SLASH = auto()
@@ -36,6 +36,7 @@ class TokenType(Enum):
     EQUAL = auto()
     EQUAL_EQUAL = auto()
     BANG_EQUAL = auto()
+    BANG = auto()
 
     # Literals
     IDENTIFIER = auto()
@@ -47,6 +48,7 @@ class TokenType(Enum):
 
     # Keywords
     TYPE = auto()
+    ALIAS = auto()
     PREDICATE = auto()
     EXTEND = auto()
     WHERE = auto()
@@ -63,6 +65,7 @@ class TokenType(Enum):
 
 KEYWORDS: dict[str, TokenType] = {
     "type": TokenType.TYPE,
+    "alias": TokenType.ALIAS,
     "predicate": TokenType.PREDICATE,
     "extend": TokenType.EXTEND,
     "where": TokenType.WHERE,
@@ -102,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,13 +2,12 @@ from typing import Optional
 
 from midas.ast.location import Location
 from midas.ast.midas import (
+    AliasStmt,
     BinaryExpr,
     CallExpr,
-    ComplexType,
     ConstraintType,
     Expr,
     ExtendStmt,
-    ExtensionType,
     FrameType,
     FunctionType,
     GenericType,
@@ -34,10 +33,11 @@ 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.EXTEND,
         TokenType.PREDICATE,
@@ -72,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):
@@ -92,23 +94,14 @@ 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
@@ -159,30 +152,53 @@ class MidasParser(Parser):
         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
         """
-        base: Type
         if self.match(TokenType.FUNC):
-            base = self.function()
-        else:
-            base = self.constraint_type()
-        if self.match(TokenType.AND):
-            extension: ComplexType = self.complex_type()
-            return ExtensionType(
-                location=Location.span(base.location, extension.location),
-                base=base,
-                extension=extension,
-            )
-        return base
+            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()
@@ -194,17 +210,33 @@ 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:
+        """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):
             if type.name.lexeme == "Frame":
@@ -218,6 +250,13 @@ class MidasParser(Parser):
         return type
 
     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):
@@ -228,41 +267,34 @@ class MidasParser(Parser):
         return args
 
     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) -> ComplexType:
-        """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:
-            ComplexType: the parsed complex type
+            FrameType: the parsed frame type
         """
-        left: Token = self.consume(
-            TokenType.LEFT_BRACE, "Expected '{' to start type body"
-        )
-        members: list[MemberStmt] = []
-        # TODO: add keyword to differentiate properties and methods,
-        # and allow multiple methods with the same name but not properties
-        names: set[str] = set()
-        while not self.check(TokenType.RIGHT_BRACE) and not self.is_at_end():
-            member: MemberStmt = self.member_stmt()
-            # if member.name.lexeme in names:
-            #    raise self.error(member.name, "Duplicate property")
-            # names.add(member.name.lexeme)
-            members.append(member)
-        right: Token = self.consume(TokenType.RIGHT_BRACE, "Unclosed type body")
-        return ComplexType(
-            location=left.location_to(right),
-            members=members,
-        )
-
-    def frame_type(self) -> FrameType:
         keyword: Token = self.previous()
         self.consume(TokenType.LEFT_BRACKET, "Expected '[' to start frame schema")
 
@@ -289,9 +321,9 @@ class MidasParser(Parser):
         )
 
     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
@@ -299,10 +331,20 @@ class MidasParser(Parser):
         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
@@ -318,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
@@ -334,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)
@@ -355,12 +441,15 @@ 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)
@@ -368,16 +457,48 @@ class MidasParser(Parser):
         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.match(TokenType.RIGHT_PAREN):
+        while not self.check(TokenType.RIGHT_PAREN):
             if self.check_identifier() and self.check_next(TokenType.EQUAL):
                 keywords = True
                 keyword: Token = self.advance()
@@ -393,13 +514,14 @@ class MidasParser(Parser):
             else:
                 value = self.expression()
                 if self.check(TokenType.EQUAL):
+                    error_msg: str
                     if keywords:
-                        raise self.error(self.peek(), "Invalid keyword argument name")
+                        error_msg = "Invalid keyword argument name"
                     else:
-                        raise self.error(
-                            self.peek(),
-                            "Cannot pass positional arguments after a keyword argument",
+                        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):
@@ -416,7 +538,12 @@ class MidasParser(Parser):
         )
 
     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
@@ -431,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
@@ -457,21 +589,48 @@ class MidasParser(Parser):
             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 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`)
+
+        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:
+            Token: the current token which is a valid identifier
+        """
         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
@@ -480,7 +639,14 @@ class MidasParser(Parser):
     def member_stmt(self) -> MemberStmt:
         """Parse a member statement
 
-        A type member statement is written `prop name: Type` or `def name: Type`
+        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
@@ -507,7 +673,13 @@ class MidasParser(Parser):
     def extend_declaration(self) -> ExtendStmt:
         """Parse an extension definition
 
-        An extension is written `extend Type { operations }` or `extend[S <: T, U] 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
@@ -532,7 +704,12 @@ class MidasParser(Parser):
     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
@@ -543,7 +720,7 @@ class MidasParser(Parser):
 
         params: list[ParamSpec] = []
         while self.check(TokenType.LEFT_PAREN):
-            params.append(self.function_args())
+            params.append(self.function_params())
 
         self.consume(TokenType.EQUAL, "Expected '=' after predicate subject")
         body: Expr = self.constraint()
@@ -555,7 +732,18 @@ class MidasParser(Parser):
         )
 
     def function(self) -> FunctionType:
-        params: ParamSpec = self.function_args()
+        """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()
@@ -566,36 +754,53 @@ class MidasParser(Parser):
             returns=result,
         )
 
-    def function_args(self) -> ParamSpec:
+    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_args: list[FunctionType.Argument] = []
-        args: list[FunctionType.Argument] = []
-        kw_args: list[FunctionType.Argument] = []
+        pos: list[FunctionType.Parameter] = []
+        mixed: list[FunctionType.Parameter] = []
+        kw: list[FunctionType.Parameter] = []
 
-        args_first_tokens: list[Token] = []
+        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_args = args
-                    args = []
-                    args_first_tokens = []
+                    pos = mixed
+                    mixed = []
+                    mixed_first_tokens = []
                     section = 1
                 case 0 | 1 if self.match(TokenType.STAR):
                     section = 2
                 case _:
-                    # Record first token of mixed argument for errors if unnamed
+                    # Record first token of mixed parameters for errors if unnamed
                     if section != 2:
-                        args_first_tokens.append(self.peek())
+                        mixed_first_tokens.append(self.peek())
 
                     name: Optional[Token] = None
                     if section == 2:
-                        name = self.consume_identifier("Expected keyword argument name")
+                        name = self.consume_identifier(
+                            "Expected keyword parameter name"
+                        )
                         self.consume(
-                            TokenType.COLON, "Expected ':' after argument name"
+                            TokenType.COLON, "Expected ':' after parameter name"
                         )
                     elif self.check_identifier() and self.check_next(TokenType.COLON):
                         name = self.advance()
@@ -603,24 +808,24 @@ class MidasParser(Parser):
 
                     type: Type = self.type_expr()
                     optional: bool = self.match(TokenType.QMARK)
-                    arg = FunctionType.Argument(
+                    param = FunctionType.Parameter(
                         location=None,
                         name=name,
                         type=type,
                         required=not optional,
                     )
                     if section == 2:
-                        kw_args.append(arg)
+                        kw.append(param)
                     else:
-                        args.append(arg)
+                        mixed.append(param)
 
             if not self.match(TokenType.COMMA):
                 break
 
-        for arg, token in zip(args, args_first_tokens):
-            if arg.name is None:
+        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 argument")
+                self.error(token, "Unnamed mixed parameter")
 
         self.consume(TokenType.RIGHT_PAREN, "Expected ')' after function parameters")
-        return ParamSpec(l_paren=l_paren, pos=pos_args, mixed=args, kw=kw_args)
+        return ParamSpec(l_paren=l_paren, pos=pos, mixed=mixed, kw=kw)

midas/parser/python.py

@@ -16,13 +16,17 @@ from midas.ast.python import (
     ForStmt,
     FrameColumn,
     FrameType,
+    FromImportStmt,
     Function,
     GetExpr,
     IfStmt,
+    ImportAlias,
+    ImportStmt,
     ListExpr,
     LiteralExpr,
     LogicalExpr,
     MidasType,
+    ParamSpec,
     RawExpr,
     RawStmt,
     ReturnStmt,
@@ -30,6 +34,7 @@ from midas.ast.python import (
     Stmt,
     SubscriptExpr,
     TernaryExpr,
+    TupleExpr,
     TypeAssign,
     UnaryExpr,
     VariableExpr,
@@ -48,6 +53,8 @@ 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"
 
@@ -101,10 +108,34 @@ class PythonParser:
             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 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] = []
         loc: Location = Location.from_ast(node)
@@ -211,27 +242,10 @@ class PythonParser:
         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,
-                ),
+                args=args,
                 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)
@@ -240,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,
@@ -300,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(
@@ -345,7 +365,7 @@ class PythonParser:
                 return BaseType(
                     location=loc,
                     base="None",
-                    param=None,
+                    args=(),
                 )
 
             case _:
@@ -360,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 _:
@@ -371,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,
@@ -477,6 +497,12 @@ class PythonParser:
                     step=self.parse_expr(step) if step is not None else None,
                 )
 
+            case ast.Tuple(elts=items):
+                return TupleExpr(
+                    location=location,
+                    items=tuple(self.parse_expr(item) for item in items),
+                )
+
             case _:
                 print(f"Unsupported expression: {ast.unparse(node)}")
                 return RawExpr(location=location, expr=node)

midas/typing.py

@@ -1,6 +1,8 @@
 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.
@@ -38,14 +40,14 @@ _**Internal Python documentation**_
 T = TypeVar("T")
 
 
-class Frame(Generic[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(Generic[T]):
+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

@@ -3,6 +3,7 @@ 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]
 
@@ -63,3 +64,4 @@ class TypedAST:
     stmts: list[p.Stmt]
     judgements: list[tuple[p.Expr, Type]]
     evaluated_casts: list[p.CastExpr]
+    assertions: AssertionCollector

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,40 +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
 
 TemplateParam ::= Identifier ("<:" Type)?
 Template ::= "[" (TemplateParam ("," TemplateParam)*)? "]"
 
+ParamType ::= Type "?"?
+PosParam ::= (Identifier ":")? ParamType
+KwParam ::= Identifier ":" ParamType
+
+PosParams ::= (
+    (PosParam ("," PosParam)* ("," "/")?)
+    | "/"
+)
+MixedParams ::= KwParam ("," KwParam)
+KwParams ::= (
+    (("*", ",")? KwParam ("," KwParam)*)
+    | "*"
+)
+Params ::= (
+    PosParams
+    | MixedParams
+    | KwParams
+    | (PosParams "," MixedParams)
+    | (PosParams "," KwParams)
+    | (MixedParams "," KwParams)
+    | (PosParams "," MixedParams "," KwParams)
+)
+ParamSpec ::= "(" Params? ")"
 
-TypeProperty ::= Identifier ":" Type
-ComplexType ::= "{" TypeProperty* "}"
 NamedType ::= Identifier
-TypeParams ::= "[" (Type ("," Type)*)? "]"
-GenericType ::= NamedType TypeParams?
+TypeArgs ::= "[" (Type ("," Type)*)? "]"
+FrameColumn ::= TOKEN ":" Type
+FrameSchema ::= "[" (FrameColumn ("," FrameColumn)*)? "]"
+GenericType ::= "Frame" FrameSchema | NamedType TypeArgs?
 GroupedType ::= "(" Type ")"
-BaseType ::= GroupedType | ComplexType | GenericType
+BaseType ::= GroupedType | GenericType
 ConstraintType ::= BaseType ("where" Constraint)?
+FuncType ::= "fn" ParamSpec "->" Type
 Type ::= ConstraintType
 
-OpDefinition ::= "op" Identifier "(" Type ")" "->" Type
-ExtendBody ::= "{" OpDefinition* "}"
+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,40 +7,61 @@ 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 expression = ```
+{[`expression` 'equality'*"&"]}
+```
+
 #let constraint = ```
-{[`constraint` 'equality'*"&"]}
+{[`constraint` 'expression']}
 ```
 
 #let template-param = ```
@@ -51,24 +72,20 @@ svg.railroad .terminal rect {
 {[`template` "[" <!, 'template-param'*","> "]"]}
 ```
 
-#let type-property = ```
-{[`type-property` 'identifier' ":" 'type']}
-```
-
-#let complex-type = ```
-{[`complex-type` "{" <!, 'type-property'*!> "}"]}
-```
-
 #let named-type = ```
 {[`named-type` 'identifier']}
 ```
 
-#let type-params = ```
-{[`type-params` "[" <!, 'type'*","> "]"]}
+#let type-args = ```
+{[`type-args` "[" <!, 'type'*","> "]"]}
+```
+
+#let frame-schema = ```
+{[`frame-schema` "[" <!, ['TOKEN' ":" 'type']*","> "]"]}
 ```
 
 #let generic-type = ```
-{[`generic-type` 'named-type' <!, 'type-params'>]}
+{[`generic-type` <["Frame" 'frame-schema'], ['named-type' <!, 'type-args'>]>]}
 ```
 
 #let grouped-type = ```
@@ -76,59 +93,88 @@ svg.railroad .terminal rect {
 ```
 
 #let base-type = ```
-{[`base-type` <'grouped-type', 'complex-type', 'generic-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` 'constraint-type']}
+{[`type` <'func-type', 'constraint-type'>]}
+```
+
+#let alias-statement = ```
+{[`alias-statement` "alias" 'identifier' "=" 'type']}
 ```
 
 #let type-statement = ```
 {[`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,
   template-param: template-param,
   template: template,
   type-property: type-property,
-  complex-type: complex-type,
   named-type: named-type,
-  type-params: type-params,
+  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,
@@ -136,18 +182,22 @@ svg.railroad .terminal rect {
 
 #let inline = (
   "grouping",
-  "value",
+  "literal",
   "template-param",
   "template",
   "type-property",
-  "complex-type",
-  "type-params",
+  "call-args",
+  "type-args",
   "named-type",
   "grouped-type",
   "generic-type",
   "base-type",
   "constraint-type",
-  "op-definition",
+  "pos-param",
+  "kw-param",
+  "func-type",
+  "member-stmt",
+  "alias-statement",
   "type-statement",
   "extend-statement",
   "predicate-statement",
@@ -164,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"""
 
@@ -32,37 +72,37 @@ class Tester(ABC):
     @abstractmethod
     def _list_tests(self) -> list[Path]: ...
 
-    def run_all_tests(self) -> bool:
+    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.resolve().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()
 
@@ -138,16 +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:
-                print("No subcommand provided. Available subcommands: run, update")
-                sys.exit(1)
+                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)

tests/cases/checker/03_functions.py.ref.json

@@ -124,7 +124,7 @@
           22
         ]
       },
-      "message": "Multiple values for argument 'b'"
+      "message": "Multiple values for parameter 'b'"
     },
     {
       "type": "Error",
@@ -152,7 +152,7 @@
           12
         ]
       },
-      "message": "Unknown keyword argument 'a'"
+      "message": "Unknown keyword parameter 'a'"
     },
     {
       "type": "Error",
@@ -194,7 +194,7 @@
           17
         ]
       },
-      "message": "Unknown keyword argument 'g'"
+      "message": "Unknown keyword parameter 'g'"
     },
     {
       "type": "Error",
@@ -277,36 +277,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -351,36 +356,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -443,36 +453,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -539,36 +554,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -649,36 +669,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -762,36 +787,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -850,36 +880,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -929,36 +964,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -1034,36 +1074,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -1150,36 +1195,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -1266,36 +1316,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }
@@ -1382,36 +1437,41 @@
         "name": "foo"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "args": [
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [
-          {
-            "pos": 2,
-            "name": "c",
-            "type": {
-              "name": "str"
-            },
-            "required": true
-          }
-        ],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "mixed": [
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": [
+            {
+              "pos": 2,
+              "name": "c",
+              "type": {
+                "name": "str"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ]
+        },
         "returns": {
           "name": "bool"
         }

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

@@ -24,7 +24,7 @@
         "type": {
           "_type": "BaseType",
           "base": "Meter",
-          "param": null
+          "args": []
         },
         "expr": {
           "_type": "LiteralExpr",
@@ -62,7 +62,7 @@
         "type": {
           "_type": "BaseType",
           "base": "Second",
-          "param": null
+          "args": []
         },
         "expr": {
           "_type": "LiteralExpr",

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

@@ -136,26 +136,30 @@
         "name": "maximum"
       },
       "type": {
-        "pos_args": [],
-        "args": [
-          {
-            "pos": 0,
-            "name": "a",
-            "type": {
-              "name": "float"
+        "params": {
+          "pos": [],
+          "mixed": [
+            {
+              "pos": 0,
+              "name": "a",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
             },
-            "required": true
-          },
-          {
-            "pos": 1,
-            "name": "b",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [],
+            {
+              "pos": 1,
+              "name": "b",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": []
+        },
         "returns": {
           "name": "float"
         }

tests/cases/checker/07_variance.py

@@ -11,16 +11,16 @@ from _ import (
     Unused,
 )
 
-unused: Unused
-covariant: Covariant
-contravariant: Contravariant
-invariant: Invariant
-coco: Coco
-cocontra: Cocontra
-contraco: Contraco
-contracontra: Contracontra
-t1: T1
-t2: T2
+unused: Unused = object()
+covariant: Covariant = object()
+contravariant: Contravariant = object()
+invariant: Invariant = object()
+coco: Coco = object()
+cocontra: Cocontra = object()
+contraco: Contraco = object()
+contracontra: Contracontra = object()
+t1: T1 = object()
+t2: T2 = object()
 
 # Dummy print to prudce judgements for the expressions
 print(
@@ -36,17 +36,17 @@ print(
     t2,
 )
 
-cov1: Covariant[float]
-cov2: Covariant[int]
+cov1: Covariant[float] = object()
+cov2: Covariant[int] = object()
 cov1 = cov2  # Ok because int <: float => Covariant[int] <: Covariant[float]
 cov2 = cov1  # Invalid
 
-contra1: Contravariant[float]
-contra2: Contravariant[int]
+contra1: Contravariant[float] = object()
+contra2: Contravariant[int] = object()
 contra1 = contra2  # Invalid
 contra2 = contra1  # Ok because int <: float => Covariant[float] <: Covariant[int]
 
-inv1: Invariant[float]
-inv2: Invariant[int]
+inv1: Invariant[float] = object()
+inv2: Invariant[int] = object()
 inv1 = inv2  # Invalid
 inv2 = inv1  # Invalid

tests/cases/checker/08_unification.py.ref.json

@@ -120,16 +120,19 @@
         "name": "bool"
       },
       "type": {
-        "pos_args": [
-          {
-            "pos": 0,
-            "name": "object",
-            "type": {},
-            "required": false
-          }
-        ],
-        "args": [],
-        "kw_args": [],
+        "params": {
+          "pos": [
+            {
+              "pos": 0,
+              "name": "object",
+              "type": {},
+              "required": false,
+              "unsupported": false
+            }
+          ],
+          "mixed": [],
+          "kw": []
+        },
         "returns": {
           "name": "bool"
         }
@@ -377,18 +380,21 @@
         "name": "double"
       },
       "type": {
-        "pos_args": [],
-        "args": [
-          {
-            "pos": 0,
-            "name": "value",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [],
+        "params": {
+          "pos": [],
+          "mixed": [
+            {
+              "pos": 0,
+              "name": "value",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": []
+        },
         "returns": {
           "name": "float"
         }
@@ -439,54 +445,61 @@
           }
         ],
         "body": {
-          "pos_args": [
-            {
-              "pos": 0,
-              "name": "transform",
-              "type": {
-                "pos_args": [
-                  {
-                    "pos": 0,
-                    "name": "v",
-                    "type": {
-                      "name": "T",
-                      "bound": null,
-                      "variance": "INVARIANT"
-                    },
-                    "required": true
-                  }
-                ],
-                "args": [],
-                "kw_args": [],
-                "returns": {
-                  "name": "U",
-                  "bound": null,
-                  "variance": "INVARIANT"
-                }
-              },
-              "required": true
-            },
-            {
-              "pos": 1,
-              "name": "iterable",
-              "type": {
-                "name": "list",
-                "args": [
-                  {
-                    "name": "T",
+          "params": {
+            "pos": [
+              {
+                "pos": 0,
+                "name": "transform",
+                "type": {
+                  "params": {
+                    "pos": [
+                      {
+                        "pos": 0,
+                        "name": "v",
+                        "type": {
+                          "name": "T",
+                          "bound": null,
+                          "variance": "INVARIANT"
+                        },
+                        "required": true,
+                        "unsupported": false
+                      }
+                    ],
+                    "mixed": [],
+                    "kw": []
+                  },
+                  "returns": {
+                    "name": "U",
                     "bound": null,
                     "variance": "INVARIANT"
                   }
-                ],
-                "body": {
-                  "name": "list"
-                }
+                },
+                "required": true,
+                "unsupported": false
               },
-              "required": true
-            }
-          ],
-          "args": [],
-          "kw_args": [],
+              {
+                "pos": 1,
+                "name": "iterable",
+                "type": {
+                  "name": "list",
+                  "args": [
+                    {
+                      "name": "T",
+                      "bound": null,
+                      "variance": "INVARIANT"
+                    }
+                  ],
+                  "body": {
+                    "name": "list"
+                  }
+                },
+                "required": true,
+                "unsupported": false
+              }
+            ],
+            "mixed": [],
+            "kw": []
+          },
           "returns": {
             "name": "list",
             "args": [
@@ -548,18 +561,21 @@
         "name": "double"
       },
       "type": {
-        "pos_args": [],
-        "args": [
-          {
-            "pos": 0,
-            "name": "value",
-            "type": {
-              "name": "float"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [],
+        "params": {
+          "pos": [],
+          "mixed": [
+            {
+              "pos": 0,
+              "name": "value",
+              "type": {
+                "name": "float"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": []
+        },
         "returns": {
           "name": "float"
         }
@@ -610,54 +626,61 @@
           }
         ],
         "body": {
-          "pos_args": [
-            {
-              "pos": 0,
-              "name": "transform",
-              "type": {
-                "pos_args": [
-                  {
-                    "pos": 0,
-                    "name": "v",
-                    "type": {
-                      "name": "T",
-                      "bound": null,
-                      "variance": "INVARIANT"
-                    },
-                    "required": true
-                  }
-                ],
-                "args": [],
-                "kw_args": [],
-                "returns": {
-                  "name": "U",
-                  "bound": null,
-                  "variance": "INVARIANT"
-                }
-              },
-              "required": true
-            },
-            {
-              "pos": 1,
-              "name": "iterable",
-              "type": {
-                "name": "list",
-                "args": [
-                  {
-                    "name": "T",
+          "params": {
+            "pos": [
+              {
+                "pos": 0,
+                "name": "transform",
+                "type": {
+                  "params": {
+                    "pos": [
+                      {
+                        "pos": 0,
+                        "name": "v",
+                        "type": {
+                          "name": "T",
+                          "bound": null,
+                          "variance": "INVARIANT"
+                        },
+                        "required": true,
+                        "unsupported": false
+                      }
+                    ],
+                    "mixed": [],
+                    "kw": []
+                  },
+                  "returns": {
+                    "name": "U",
                     "bound": null,
                     "variance": "INVARIANT"
                   }
-                ],
-                "body": {
-                  "name": "list"
-                }
+                },
+                "required": true,
+                "unsupported": false
               },
-              "required": true
-            }
-          ],
-          "args": [],
-          "kw_args": [],
+              {
+                "pos": 1,
+                "name": "iterable",
+                "type": {
+                  "name": "list",
+                  "args": [
+                    {
+                      "name": "T",
+                      "bound": null,
+                      "variance": "INVARIANT"
+                    }
+                  ],
+                  "body": {
+                    "name": "list"
+                  }
+                },
+                "required": true,
+                "unsupported": false
+              }
+            ],
+            "mixed": [],
+            "kw": []
+          },
           "returns": {
             "name": "list",
             "args": [
@@ -709,18 +732,21 @@
         "name": "is_odd"
       },
       "type": {
-        "pos_args": [],
-        "args": [
-          {
-            "pos": 0,
-            "name": "value",
-            "type": {
-              "name": "int"
-            },
-            "required": true
-          }
-        ],
-        "kw_args": [],
+        "params": {
+          "pos": [],
+          "mixed": [
+            {
+              "pos": 0,
+              "name": "value",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": []
+        },
         "returns": {
           "name": "bool"
         }
@@ -771,54 +797,61 @@
           }
         ],
         "body": {
-          "pos_args": [
-            {
-              "pos": 0,
-              "name": "transform",
-              "type": {
-                "pos_args": [
-                  {
-                    "pos": 0,
-                    "name": "v",
-                    "type": {
-                      "name": "T",
-                      "bound": null,
-                      "variance": "INVARIANT"
-                    },
-                    "required": true
-                  }
-                ],
-                "args": [],
-                "kw_args": [],
-                "returns": {
-                  "name": "U",
-                  "bound": null,
-                  "variance": "INVARIANT"
-                }
-              },
-              "required": true
-            },
-            {
-              "pos": 1,
-              "name": "iterable",
-              "type": {
-                "name": "list",
-                "args": [
-                  {
-                    "name": "T",
+          "params": {
+            "pos": [
+              {
+                "pos": 0,
+                "name": "transform",
+                "type": {
+                  "params": {
+                    "pos": [
+                      {
+                        "pos": 0,
+                        "name": "v",
+                        "type": {
+                          "name": "T",
+                          "bound": null,
+                          "variance": "INVARIANT"
+                        },
+                        "required": true,
+                        "unsupported": false
+                      }
+                    ],
+                    "mixed": [],
+                    "kw": []
+                  },
+                  "returns": {
+                    "name": "U",
                     "bound": null,
                     "variance": "INVARIANT"
                   }
-                ],
-                "body": {
-                  "name": "list"
-                }
+                },
+                "required": true,
+                "unsupported": false
               },
-              "required": true
-            }
-          ],
-          "args": [],
-          "kw_args": [],
+              {
+                "pos": 1,
+                "name": "iterable",
+                "type": {
+                  "name": "list",
+                  "args": [
+                    {
+                      "name": "T",
+                      "bound": null,
+                      "variance": "INVARIANT"
+                    }
+                  ],
+                  "body": {
+                    "name": "list"
+                  }
+                },
+                "required": true,
+                "unsupported": false
+              }
+            ],
+            "mixed": [],
+            "kw": []
+          },
           "returns": {
             "name": "list",
             "args": [

tests/cases/checker/09_frame_ops.py

@@ -0,0 +1,118 @@
+# type: ignore
+# ruff: disable [F821]
+import pandas as pd
+
+df1 = cast(Frame[i:int, a:int, b:float], pd.DataFrame())
+df2 = cast(Frame[i:int, a:int, b:float], pd.DataFrame())
+
+_: Any
+
+# Arithmetic
+_ = df1 + df2
+_ = df1 - df2
+_ = df1 * df2
+_ = df1 / df2
+_ = df1 // df2
+_ = df1 % df2
+_ = df1**df2
+
+# Comparisons
+_ = df1 < df2
+_ = df1 > df2
+_ = df1 <= df2
+_ = df1 >= df2
+_ = df1 != df2
+_ = df1 == df2
+
+# Aggregate
+_ = df1.kurt()
+_ = df1.kurtosis()
+_ = df1.max()
+_ = df1.mean()
+_ = df1.median()
+_ = df1.min()
+_ = df1.mode()
+_ = df1.prod()
+_ = df1.product()
+_ = df1.std()
+_ = df1.sum()
+_ = df1.var()
+
+# Groupby
+df_gb = df1.groupby(by="i")
+
+_ = df_gb.kurt()
+_ = df_gb.max()
+_ = df_gb.mean()
+_ = df_gb.median()
+_ = df_gb.min()
+_ = df_gb.prod()
+_ = df_gb.std()
+_ = df_gb.sum()
+_ = df_gb.var()
+
+
+# Columns
+
+col1 = df1["a"]
+col2 = df1["a"]
+
+# Arithmetic
+_ = col1 + col2
+_ = col1 - col2
+_ = col1 * col2
+_ = col1 / col2
+_ = col1 // col2
+_ = col1 % col2
+_ = col1**col2
+
+# Comparisons
+_ = col1 < col2
+_ = col1 > col2
+_ = col1 <= col2
+_ = col1 >= col2
+_ = col1 != col2
+_ = col1 == col2
+
+# Aggregate
+_ = col1.kurt()
+_ = col1.kurtosis()
+_ = col1.max()
+_ = col1.mean()
+_ = col1.median()
+_ = col1.min()
+_ = col1.mode()
+_ = col1.prod()
+_ = col1.product()
+_ = col1.std()
+_ = col1.sum()
+_ = col1.var()
+
+# Groupby
+col_gb = col1.groupby(level=0)
+
+_ = col_gb.kurt()
+_ = col_gb.max()
+_ = col_gb.mean()
+_ = col_gb.median()
+_ = col_gb.min()
+_ = col_gb.prod()
+_ = col_gb.std()
+_ = col_gb.sum()
+_ = col_gb.var()
+
+# Attributes
+_ = df1.ndim  # int
+_ = df1.size  # int
+_ = df1.shape  # (int, int)
+_ = col1.ndim  # int
+_ = col1.size  # int
+_ = col1.shape  # (int)
+_ = col1.T  # Column[int]
+
+
+# Misc
+_ = df1.head()
+_ = df1.tail()
+_ = col1.head()
+_ = col1.tail()

tests/cases/checker/10_variable_defined.py

@@ -0,0 +1,45 @@
+# type: ignore
+# ruff: disable[F821]
+
+_: Any
+
+_ = undeclared
+
+declared: int
+_ = declared
+
+half_defined1: int
+half_defined2: int
+if False:
+    half_defined1 = 0
+else:
+    half_defined2 = 1
+_ = half_defined1
+_ = half_defined2
+
+fully_defined: int
+if False:
+    fully_defined = 0
+else:
+    fully_defined = 1
+_ = fully_defined
+
+defined: int = 0
+_ = defined
+
+no_annotation = 0
+_ = no_annotation
+
+self_ref1 = self_ref1
+self_ref2: int = self_ref2
+
+
+def fact(n: int) -> int:
+    if n <= 1:
+        return 1
+    return n * fact(n - 1)
+
+
+for i in [1, 2, 3]:
+    _ = i
+_ = i

tests/cases/checker/10_variable_defined.py.ref.json

@@ -0,0 +1,651 @@
+{
+  "diagnostics": [
+    {
+      "type": "Error",
+      "location": {
+        "start": [
+          9,
+          4
+        ],
+        "end": [
+          9,
+          12
+        ]
+      },
+      "message": "Variable 'declared' is declared but may not be defined"
+    },
+    {
+      "type": "Error",
+      "location": {
+        "start": [
+          17,
+          4
+        ],
+        "end": [
+          17,
+          17
+        ]
+      },
+      "message": "Variable 'half_defined1' is declared but may not be defined"
+    },
+    {
+      "type": "Error",
+      "location": {
+        "start": [
+          18,
+          4
+        ],
+        "end": [
+          18,
+          17
+        ]
+      },
+      "message": "Variable 'half_defined2' is declared but may not be defined"
+    },
+    {
+      "type": "Error",
+      "location": {
+        "start": [
+          34,
+          17
+        ],
+        "end": [
+          34,
+          26
+        ]
+      },
+      "message": "Variable 'self_ref2' is declared but may not be defined"
+    },
+    {
+      "type": "Warning",
+      "location": {
+        "start": [
+          6,
+          4
+        ],
+        "end": [
+          6,
+          14
+        ]
+      },
+      "message": "Unknown variable"
+    },
+    {
+      "type": "Warning",
+      "location": {
+        "start": [
+          33,
+          12
+        ],
+        "end": [
+          33,
+          21
+        ]
+      },
+      "message": "Unknown variable"
+    },
+    {
+      "type": "Warning",
+      "location": {
+        "start": [
+          45,
+          4
+        ],
+        "end": [
+          45,
+          5
+        ]
+      },
+      "message": "Unknown variable"
+    }
+  ],
+  "judgments": [
+    {
+      "location": {
+        "from": "L6:4",
+        "to": "L6:14"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "undeclared"
+      },
+      "type": {}
+    },
+    {
+      "location": {
+        "from": "L9:4",
+        "to": "L9:12"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "declared"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L13:3",
+        "to": "L13:8"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": false
+      },
+      "type": {
+        "name": "bool"
+      }
+    },
+    {
+      "location": {
+        "from": "L14:20",
+        "to": "L14:21"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 0
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L16:20",
+        "to": "L16:21"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 1
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L17:4",
+        "to": "L17:17"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "half_defined1"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L18:4",
+        "to": "L18:17"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "half_defined2"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L21:3",
+        "to": "L21:8"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": false
+      },
+      "type": {
+        "name": "bool"
+      }
+    },
+    {
+      "location": {
+        "from": "L22:20",
+        "to": "L22:21"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 0
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L24:20",
+        "to": "L24:21"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 1
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L25:4",
+        "to": "L25:17"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "fully_defined"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L27:15",
+        "to": "L27:16"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 0
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L28:4",
+        "to": "L28:11"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "defined"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L30:16",
+        "to": "L30:17"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 0
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L31:4",
+        "to": "L31:17"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "no_annotation"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L33:12",
+        "to": "L33:21"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "self_ref1"
+      },
+      "type": {}
+    },
+    {
+      "location": {
+        "from": "L34:17",
+        "to": "L34:26"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "self_ref2"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L38:7",
+        "to": "L38:8"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "n"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L38:12",
+        "to": "L38:13"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 1
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L38:7",
+        "to": "L38:13"
+      },
+      "expr": {
+        "_type": "CompareExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "n"
+        },
+        "operator": "<=",
+        "right": {
+          "_type": "LiteralExpr",
+          "value": 1
+        }
+      },
+      "type": {
+        "name": "bool"
+      }
+    },
+    {
+      "location": {
+        "from": "L39:15",
+        "to": "L39:16"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 1
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L40:11",
+        "to": "L40:12"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "n"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L40:20",
+        "to": "L40:21"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "n"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L40:24",
+        "to": "L40:25"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 1
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L40:20",
+        "to": "L40:25"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "n"
+        },
+        "operator": "-",
+        "right": {
+          "_type": "LiteralExpr",
+          "value": 1
+        }
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L40:15",
+        "to": "L40:19"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "fact"
+      },
+      "type": {
+        "params": {
+          "pos": [],
+          "mixed": [
+            {
+              "pos": 0,
+              "name": "n",
+              "type": {
+                "name": "int"
+              },
+              "required": true,
+              "unsupported": false
+            }
+          ],
+          "kw": []
+        },
+        "returns": {
+          "name": "int"
+        }
+      }
+    },
+    {
+      "location": {
+        "from": "L40:15",
+        "to": "L40:26"
+      },
+      "expr": {
+        "_type": "CallExpr",
+        "callee": {
+          "_type": "VariableExpr",
+          "name": "fact"
+        },
+        "arguments": [
+          {
+            "_type": "BinaryExpr",
+            "left": {
+              "_type": "VariableExpr",
+              "name": "n"
+            },
+            "operator": "-",
+            "right": {
+              "_type": "LiteralExpr",
+              "value": 1
+            }
+          }
+        ],
+        "keywords": {}
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L40:11",
+        "to": "L40:26"
+      },
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "n"
+        },
+        "operator": "*",
+        "right": {
+          "_type": "CallExpr",
+          "callee": {
+            "_type": "VariableExpr",
+            "name": "fact"
+          },
+          "arguments": [
+            {
+              "_type": "BinaryExpr",
+              "left": {
+                "_type": "VariableExpr",
+                "name": "n"
+              },
+              "operator": "-",
+              "right": {
+                "_type": "LiteralExpr",
+                "value": 1
+              }
+            }
+          ],
+          "keywords": {}
+        }
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L43:10",
+        "to": "L43:11"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 1
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L43:13",
+        "to": "L43:14"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 2
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L43:16",
+        "to": "L43:17"
+      },
+      "expr": {
+        "_type": "LiteralExpr",
+        "value": 3
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L43:9",
+        "to": "L43:18"
+      },
+      "expr": {
+        "_type": "ListExpr",
+        "items": [
+          {
+            "_type": "LiteralExpr",
+            "value": 1
+          },
+          {
+            "_type": "LiteralExpr",
+            "value": 2
+          },
+          {
+            "_type": "LiteralExpr",
+            "value": 3
+          }
+        ]
+      },
+      "type": {
+        "name": "list",
+        "args": [
+          {
+            "name": "int"
+          }
+        ],
+        "body": {
+          "name": "list"
+        }
+      }
+    },
+    {
+      "location": {
+        "from": "L43:4",
+        "to": "L43:5"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "i"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L44:8",
+        "to": "L44:9"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "i"
+      },
+      "type": {
+        "name": "int"
+      }
+    },
+    {
+      "location": {
+        "from": "L45:4",
+        "to": "L45:5"
+      },
+      "expr": {
+        "_type": "VariableExpr",
+        "name": "i"
+      },
+      "type": {}
+    }
+  ]
+}

tests/cases/generator/02_constraints.py.ref.txt

@@ -18,7 +18,7 @@ Module(
             ops=[
               Gt()],
             comparators=[
-              Constant(value=0.0)]))],
+              Constant(value=0)]))],
       decorator_list=[],
       returns=Constant(value='bool')),
     FunctionDef(
@@ -39,7 +39,7 @@ Module(
             ops=[
               Gt()],
             comparators=[
-              Constant(value=0.0)]))],
+              Constant(value=0)]))],
       decorator_list=[],
       returns=Constant(value='bool')),
     FunctionDef(
@@ -128,8 +128,8 @@ Module(
             func=Call(
               func=Name(id='__midas_in_range__'),
               args=[
-                Constant(value=100.0),
-                Constant(value=200.0)],
+                Constant(value=100),
+                Constant(value=200)],
               keywords=[]),
             args=[
               Name(id='_')],
@@ -142,8 +142,8 @@ Module(
       value=Call(
         func=Name(id='__midas_in_range__'),
         args=[
-          Constant(value=0.0),
-          Constant(value=18.0)],
+          Constant(value=0),
+          Constant(value=18)],
         keywords=[])),
     FunctionDef(
       name='__midas_p3__',
@@ -208,7 +208,7 @@ Module(
         args=[
           Name(id='__midas_a0__')],
         keywords=[]),
-      msg=Constant(value="02_constraints.py:L5:10: ConstraintError: Value does not fit constraint '_ > 0.0'")),
+      msg=Constant(value="02_constraints.py:L5:10: ConstraintError: Value does not fit constraint '_ > 0'")),
     Assign(
       targets=[
         Name(id='t1')],
@@ -284,7 +284,7 @@ Module(
         args=[
           Name(id='__midas_a2__')],
         keywords=[]),
-      msg=Constant(value="02_constraints.py:L7:10: ConstraintError: Value does not fit constraint 'in_range(100.0, 200.0)(_)'")),
+      msg=Constant(value="02_constraints.py:L7:10: ConstraintError: Value does not fit constraint 'in_range(100, 200)(_)'")),
     Assign(
       targets=[
         Name(id='t3')],

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

@@ -9,26 +9,22 @@ type Longitude = float where (-180 <= _ <= 180)
 type Difference[T] = T
 
 // Complex custom type, containing two values accessible through properties
-type GeoLocation = {
+type GeoLocation = object
+
+extend GeoLocation {
     prop lat: Latitude
     prop lon: Longitude
 }
 
-// Define operations on our custom type
-extend GeoLocation {
-    // This type is compatible with the `-` operation with another GeoLocation
-    // i.e. you can subtract a GeoLocation from another GeoLocation, resulting
-    // in a Difference of GeoLocations
-    def __sub__: fn(GeoLocation, /) -> Difference[GeoLocation]
-}
+type GeoLocationDifference = object
 
-// For complex generics, you need to specify how the genericity the properties
-// are handled
-type Difference[GeoLocation] = {
+extend GeoLocationDifference {
     prop lat: Difference[Latitude]
     prop lon: Difference[Longitude]
 }
 
+// Define operations on our custom type
+
 // Simple operation defined on our custom types
 extend Latitude {
     def __sub__: fn(Latitude, /) -> Difference[Latitude]
@@ -38,13 +34,23 @@ extend Longitude {
     def __sub__: fn(Longitude, /) -> Difference[Longitude]
 }
 
+extend GeoLocation {
+    // This type is compatible with the `-` operation with another GeoLocation
+    // i.e. you can subtract a GeoLocation from another GeoLocation, resulting
+    // in a GeoLocationDifference
+    def __sub__: fn(GeoLocation, /) -> GeoLocationDifference
+}
+
+
 // Predefined custom predicates that can be referenced in other definitions
 predicate Positive(v: float) = v >= 0
 predicate StrictlyPositive(v: float) = v > 0
 predicate Equatorial(loc: GeoLocation) = (-10 <= loc.lat <= 10)
 predicate Arctic(loc: GeoLocation) = (loc.lat >= 66)
 
-type Person = {
+type Person = object
+
+extend Person {
     prop name: str
 
     // Property with an inline constraint

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

@@ -1,8 +1,16 @@
 {
   "stmts": [
     {
-      "_type": "RawStmt",
-      "stmt": "from __future__ import annotations"
+      "_type": "FromImportStmt",
+      "module": "__future__",
+      "imports": [
+        {
+          "_type": "ImportAlias",
+          "name": "annotations",
+          "alias": null
+        }
+      ],
+      "level": 0
     },
     {
       "_type": "TypeAssign",
@@ -16,7 +24,7 @@
             "type": {
               "_type": "BaseType",
               "base": "bool",
-              "param": null
+              "args": []
             }
           },
           {
@@ -25,7 +33,7 @@
             "type": {
               "_type": "BaseType",
               "base": "int",
-              "param": null
+              "args": []
             }
           },
           {
@@ -36,7 +44,7 @@
               "type": {
                 "_type": "BaseType",
                 "base": "float",
-                "param": null
+                "args": []
               },
               "constraint": "(_ > 0) + (_ < 250)"
             }
@@ -47,7 +55,7 @@
             "type": {
               "_type": "BaseType",
               "base": "str",
-              "param": null
+              "args": []
             }
           },
           {
@@ -56,7 +64,7 @@
             "type": {
               "_type": "BaseType",
               "base": "datetime",
-              "param": null
+              "args": []
             }
           },
           {
@@ -65,7 +73,7 @@
             "type": {
               "_type": "BaseType",
               "base": "float",
-              "param": null
+              "args": []
             }
           },
           {
@@ -79,7 +87,7 @@
             "type": {
               "_type": "BaseType",
               "base": "_",
-              "param": null
+              "args": []
             }
           }
         ]

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

@@ -1,8 +1,16 @@
 {
   "stmts": [
     {
-      "_type": "RawStmt",
-      "stmt": "from __future__ import annotations"
+      "_type": "FromImportStmt",
+      "module": "__future__",
+      "imports": [
+        {
+          "_type": "ImportAlias",
+          "name": "annotations",
+          "alias": null
+        }
+      ],
+      "level": 0
     },
     {
       "_type": "TypeAssign",
@@ -16,7 +24,7 @@
             "type": {
               "_type": "BaseType",
               "base": "GeoLocation",
-              "param": null
+              "args": []
             }
           }
         ]
@@ -28,11 +36,13 @@
       "type": {
         "_type": "BaseType",
         "base": "Column",
-        "param": {
-          "_type": "BaseType",
-          "base": "GeoLocation",
-          "param": null
-        }
+        "args": [
+          {
+            "_type": "BaseType",
+            "base": "GeoLocation",
+            "args": []
+          }
+        ]
       }
     },
     {
@@ -65,11 +75,13 @@
       "type": {
         "_type": "BaseType",
         "base": "Column",
-        "param": {
-          "_type": "BaseType",
-          "base": "GeoLocation",
-          "param": null
-        }
+        "args": [
+          {
+            "_type": "BaseType",
+            "base": "GeoLocation",
+            "args": []
+          }
+        ]
       }
     },
     {
@@ -117,7 +129,7 @@
       "type": {
         "_type": "BaseType",
         "base": "Latitude",
-        "param": null
+        "args": []
       }
     },
     {
@@ -146,7 +158,7 @@
       "type": {
         "_type": "BaseType",
         "base": "Latitude",
-        "param": null
+        "args": []
       }
     },
     {
@@ -175,11 +187,13 @@
       "type": {
         "_type": "BaseType",
         "base": "Difference",
-        "param": {
-          "_type": "BaseType",
-          "base": "Latitude",
-          "param": null
-        }
+        "args": [
+          {
+            "_type": "BaseType",
+            "base": "Latitude",
+            "args": []
+          }
+        ]
       }
     },
     {
@@ -217,7 +231,7 @@
               "type": {
                 "_type": "BaseType",
                 "base": "int",
-                "param": null
+                "args": []
               },
               "constraint": "_ >= 0"
             }
@@ -230,7 +244,7 @@
               "type": {
                 "_type": "BaseType",
                 "base": "float",
-                "param": null
+                "args": []
               },
               "constraint": "_ >= 0"
             }
@@ -252,7 +266,7 @@
               "type": {
                 "_type": "BaseType",
                 "base": "int",
-                "param": null
+                "args": []
               },
               "constraint": "Positive"
             }
@@ -265,7 +279,7 @@
               "type": {
                 "_type": "BaseType",
                 "base": "float",
-                "param": null
+                "args": []
               },
               "constraint": "Positive"
             }

tests/cases/python-parser/03_functions.py.ref.json

@@ -1,64 +1,79 @@
 {
   "stmts": [
     {
-      "_type": "RawStmt",
-      "stmt": "from __future__ import annotations"
+      "_type": "FromImportStmt",
+      "module": "__future__",
+      "imports": [
+        {
+          "_type": "ImportAlias",
+          "name": "annotations",
+          "alias": null
+        }
+      ],
+      "level": 0
     },
     {
       "_type": "Function",
       "name": "func",
-      "posonlyargs": [],
-      "args": [
-        {
-          "name": "col1",
-          "type": {
-            "_type": "BaseType",
-            "base": "Column",
-            "param": {
-              "_type": "ConstraintType",
-              "type": {
-                "_type": "BaseType",
-                "base": "float",
-                "param": null
-              },
-              "constraint": "0 <= _ <= 1"
-            }
+      "params": {
+        "_type": "ParamSpec",
+        "pos": [],
+        "mixed": [
+          {
+            "name": "col1",
+            "type": {
+              "_type": "BaseType",
+              "base": "Column",
+              "args": [
+                {
+                  "_type": "ConstraintType",
+                  "type": {
+                    "_type": "BaseType",
+                    "base": "float",
+                    "args": []
+                  },
+                  "constraint": "0 <= _ <= 1"
+                }
+              ]
+            },
+            "default": null
           },
-          "default": null
-        },
-        {
-          "name": "col2",
-          "type": {
-            "_type": "BaseType",
-            "base": "Column",
-            "param": {
-              "_type": "ConstraintType",
-              "type": {
-                "_type": "BaseType",
-                "base": "float",
-                "param": null
-              },
-              "constraint": "0 <= _ <= 1"
-            }
-          },
-          "default": null
-        }
-      ],
-      "sink": null,
-      "kwonlyargs": [],
-      "kw_sink": null,
+          {
+            "name": "col2",
+            "type": {
+              "_type": "BaseType",
+              "base": "Column",
+              "args": [
+                {
+                  "_type": "ConstraintType",
+                  "type": {
+                    "_type": "BaseType",
+                    "base": "float",
+                    "args": []
+                  },
+                  "constraint": "0 <= _ <= 1"
+                }
+              ]
+            },
+            "default": null
+          }
+        ],
+        "kw": []
+      },
       "returns": {
         "_type": "BaseType",
         "base": "Column",
-        "param": {
-          "_type": "ConstraintType",
-          "type": {
-            "_type": "BaseType",
-            "base": "float",
-            "param": null
-          },
-          "constraint": "0 <= _ <= 2"
-        }
+        "args": [
+          {
+            "_type": "ConstraintType",
+            "type": {
+              "_type": "BaseType",
+              "base": "float",
+              "args": []
+            },
+            "constraint": "0 <= _ <= 2"
+          }
+        ]
       },
       "body": [
         {
@@ -67,15 +82,17 @@
           "type": {
             "_type": "BaseType",
             "base": "Column",
-            "param": {
-              "_type": "ConstraintType",
-              "type": {
-                "_type": "BaseType",
-                "base": "float",
-                "param": null
-              },
-              "constraint": "0 <= _ <= 2"
-            }
+            "args": [
+              {
+                "_type": "ConstraintType",
+                "type": {
+                  "_type": "BaseType",
+                  "base": "float",
+                  "args": []
+                },
+                "constraint": "0 <= _ <= 2"
+              }
+            ]
           }
         },
         {
@@ -111,41 +128,42 @@
     {
       "_type": "Function",
       "name": "func2",
-      "posonlyargs": [
-        {
-          "name": "a",
-          "type": {
-            "_type": "BaseType",
-            "base": "int",
-            "param": null
-          },
-          "default": null
-        }
-      ],
-      "args": [
-        {
-          "name": "b",
-          "type": {
-            "_type": "BaseType",
-            "base": "float",
-            "param": null
-          },
-          "default": null
-        }
-      ],
-      "sink": null,
-      "kwonlyargs": [
-        {
-          "name": "c",
-          "type": {
-            "_type": "BaseType",
-            "base": "str",
-            "param": null
-          },
-          "default": null
-        }
-      ],
-      "kw_sink": null,
+      "params": {
+        "_type": "ParamSpec",
+        "pos": [
+          {
+            "name": "a",
+            "type": {
+              "_type": "BaseType",
+              "base": "int",
+              "args": []
+            },
+            "default": null
+          }
+        ],
+        "mixed": [
+          {
+            "name": "b",
+            "type": {
+              "_type": "BaseType",
+              "base": "float",
+              "args": []
+            },
+            "default": null
+          }
+        ],
+        "kw": [
+          {
+            "name": "c",
+            "type": {
+              "_type": "BaseType",
+              "base": "str",
+              "args": []
+            },
+            "default": null
+          }
+        ]
+      },
       "returns": null,
       "body": []
     }

tests/generator.py

@@ -46,7 +46,8 @@ class GeneratorTester(Tester):
 
         if not any(d.type == DiagnosticType.ERROR for d in checker.diagnostics):
             generator = Generator(workdir=path.parent, types=checker.types)
-            result.compiled_ast = generator.generate_ast(typed_ast, path)
+            generator.set_src_path(path)
+            result.compiled_ast = generator.generate_ast(typed_ast)
 
         return result
 

tests/serializer/midas.py

@@ -1,13 +1,12 @@
 from typing import Optional, Sequence
 
 from midas.ast.midas import (
+    AliasStmt,
     BinaryExpr,
     CallExpr,
-    ComplexType,
     ConstraintType,
     Expr,
     ExtendStmt,
-    ExtensionType,
     FrameType,
     FunctionType,
     GenericType,
@@ -61,6 +60,13 @@ class MidasAstJsonSerializer(
             "bound": self._serialize_optional(param.bound),
         }
 
+    def visit_alias_stmt(self, stmt: AliasStmt) -> dict:
+        return {
+            "_type": "AliasStmt",
+            "name": stmt.name.lexeme,
+            "type": stmt.type.accept(self),
+        }
+
     def visit_member_stmt(self, stmt: MemberStmt) -> dict:
         return {
             "_type": "MemberStmt",
@@ -164,12 +170,6 @@ class MidasAstJsonSerializer(
             "constraint": type.constraint.accept(self),
         }
 
-    def visit_complex_type(self, type: ComplexType) -> dict:
-        return {
-            "_type": "ComplexType",
-            "members": self._serialize_list(type.members),
-        }
-
     def visit_function_type(self, type: FunctionType) -> dict:
         return {
             "_type": "FunctionType",
@@ -180,23 +180,16 @@ class MidasAstJsonSerializer(
     def _serialize_param_spec(self, spec: ParamSpec) -> dict:
         return {
             "_type": "ParamSpec",
-            "pos": [self._serialize_func_arg(arg) for arg in spec.pos],
-            "mixed": [self._serialize_func_arg(arg) for arg in spec.mixed],
-            "kw": [self._serialize_func_arg(arg) for arg in spec.kw],
+            "pos": [self._serialize_func_param(arg) for arg in spec.pos],
+            "mixed": [self._serialize_func_param(arg) for arg in spec.mixed],
+            "kw": [self._serialize_func_param(arg) for arg in spec.kw],
         }
 
-    def _serialize_func_arg(self, arg: FunctionType.Argument) -> dict:
+    def _serialize_func_param(self, param: FunctionType.Parameter) -> dict:
         return {
-            "name": arg.name.lexeme if arg.name is not None else None,
-            "type": arg.type.accept(self),
-            "required": arg.required,
-        }
-
-    def visit_extension_type(self, type: ExtensionType) -> dict:
-        return {
-            "_type": "ExtensionType",
-            "base": type.base.accept(self),
-            "extension": type.extension.accept(self),
+            "name": param.name.lexeme if param.name is not None else None,
+            "type": param.type.accept(self),
+            "required": param.required,
         }
 
     def visit_frame_type(self, type: FrameType) -> dict:

tests/serializer/python.py

@@ -15,13 +15,17 @@ from midas.ast.python import (
     ForStmt,
     FrameColumn,
     FrameType,
+    FromImportStmt,
     Function,
     GetExpr,
     IfStmt,
+    ImportAlias,
+    ImportStmt,
     ListExpr,
     LiteralExpr,
     LogicalExpr,
     MidasType,
+    ParamSpec,
     Pass,
     RawExpr,
     RawStmt,
@@ -30,6 +34,7 @@ from midas.ast.python import (
     Stmt,
     SubscriptExpr,
     TernaryExpr,
+    TupleExpr,
     TypeAssign,
     UnaryExpr,
     VariableExpr,
@@ -98,7 +103,7 @@ class PythonAstJsonSerializer(
         return {
             "_type": "BaseType",
             "base": node.base,
-            "param": self._serialize_optional(node.param),
+            "args": self._serialize_list(node.args),
         }
 
     def visit_constraint_type(self, node: ConstraintType) -> dict:
@@ -127,32 +132,30 @@ class PythonAstJsonSerializer(
             "expr": stmt.expr.accept(self),
         }
 
-    def _serialize_argument(self, arg: Function.Argument) -> dict:
-        return {
-            "name": arg.name,
-            "type": self._serialize_optional(arg.type),
-            "default": self._serialize_optional(arg.default),
-        }
-
     def visit_function(self, stmt: Function) -> dict:
         return {
             "_type": "Function",
             "name": stmt.name,
-            "posonlyargs": [self._serialize_argument(arg) for arg in stmt.posonlyargs],
-            "args": [self._serialize_argument(arg) for arg in stmt.args],
-            "sink": (
-                self._serialize_argument(stmt.sink) if stmt.sink is not None else None
-            ),
-            "kwonlyargs": [self._serialize_argument(arg) for arg in stmt.kwonlyargs],
-            "kw_sink": (
-                self._serialize_argument(stmt.kw_sink)
-                if stmt.kw_sink is not None
-                else None
-            ),
+            "params": self._serialize_param_spec(stmt.params),
             "returns": self._serialize_optional(stmt.returns),
             "body": self._serialize_list(stmt.body),
         }
 
+    def _serialize_param_spec(self, spec: ParamSpec) -> dict:
+        return {
+            "_type": "ParamSpec",
+            "pos": [self._serialize_func_param(arg) for arg in spec.pos],
+            "mixed": [self._serialize_func_param(arg) for arg in spec.mixed],
+            "kw": [self._serialize_func_param(arg) for arg in spec.kw],
+        }
+
+    def _serialize_func_param(self, param: Function.Parameter) -> dict:
+        return {
+            "name": param.name,
+            "type": self._serialize_optional(param.type),
+            "default": self._serialize_optional(param.default),
+        }
+
     def visit_type_assign(self, stmt: TypeAssign) -> dict:
         return {
             "_type": "TypeAssign",
@@ -194,6 +197,30 @@ class PythonAstJsonSerializer(
             "body": self._serialize_list(stmt.body),
         }
 
+    def visit_import_stmt(self, stmt: ImportStmt) -> dict:
+        return {
+            "_type": "ImportStmt",
+            "imports": self._serialize_imports(stmt.imports),
+        }
+
+    def visit_from_import_stmt(self, stmt: FromImportStmt) -> dict:
+        return {
+            "_type": "FromImportStmt",
+            "module": stmt.module,
+            "imports": self._serialize_imports(stmt.imports),
+            "level": stmt.level,
+        }
+
+    def _serialize_imports(self, imports: list[ImportAlias]) -> list:
+        return [
+            {
+                "_type": "ImportAlias",
+                "name": import_.name,
+                "alias": import_.alias,
+            }
+            for import_ in imports
+        ]
+
     def visit_raw_stmt(self, stmt: RawStmt) -> dict:
         return {
             "_type": "RawStmt",
@@ -302,6 +329,12 @@ class PythonAstJsonSerializer(
             "step": self._serialize_optional(expr.step),
         }
 
+    def visit_tuple_expr(self, expr: TupleExpr) -> dict:
+        return {
+            "_type": "TupleExpr",
+            "items": [item.accept(self) for item in expr.items],
+        }
+
     def visit_raw_expr(self, expr: RawExpr) -> dict:
         return {
             "_type": "RawExpr",

vscode-ext/language-configurations.json

@@ -1,19 +1,16 @@
 {
     "brackets": [
         ["{", "}"],
-        ["[", "]"],
-        ["<", ">"]
+        ["[", "]"]
     ],
     "autoClosingPairs": [
         { "open": "{", "close": "}" },
         { "open": "[", "close": "]" },
-        { "open": "(", "close": ")" },
-        { "open": "<", "close": ">" }
+        { "open": "(", "close": ")" }
     ],
     "surroundingPairs": [
         ["{", "}"],
         ["[", "]"],
-        ["(", ")"],
-        ["<", ">"]
+        ["(", ")"]
     ]
 }

vscode-ext/package.json

@@ -10,7 +10,6 @@
             {
                 "id": "midas",
                 "extensions": [
-                    ".mpy",
                     ".midas"
                 ],
                 "aliases": [
@@ -23,10 +22,7 @@
             {
                 "language": "midas",
                 "scopeName": "source.midas",
-                "path": "./syntaxes/midas.tmLanguage.json",
-                "embeddedLanguages": {
-                    "meta.embedded.block.python": "python"
-                }
+                "path": "./syntaxes/midas.tmLanguage.json"
             }
         ]
     }

vscode-ext/syntaxes/midas.tmLanguage.json

@@ -2,167 +2,558 @@
     "$schema": "https://raw.githubusercontent.com/martinring/tmlanguage/master/tmlanguage.json",
     "name": "Midas",
     "scopeName": "source.midas",
-    "patterns": [{ "include": "#statement" }],
+    "fileTypes": [
+        "midas"
+    ],
+    "patterns": [
+        {
+            "include": "#comments"
+        },
+        {
+            "include": "#alias-stmt"
+        },
+        {
+            "include": "#type-stmt"
+        },
+        {
+            "include": "#extend-stmt"
+        },
+        {
+            "include": "#extend-body"
+        },
+        {
+            "include": "#predicate-stmt"
+        }
+    ],
     "repository": {
-        "comment": {
-            "begin": "(//)",
-            "end": "($)",
-            "name": "comment.line",
-            "beginCaptures": {
-                "1": {
-                    "name": "comment.line.double-dash"
-                }
-            }
-        },
-        "type-def": {
-            "begin": "\\b(type)\\s+([a-zA-Z_][a-zA-Z_\\d]*)",
-            "end": "$",
-            "beginCaptures": {
-                "1": {
-                    "name": "keyword.control.type.midas"
-                },
-                "2": {
-                    "name" : "variable.name"
-                }
-            },
+        "comments": {
             "patterns": [
-                { "include": "#type-base" },
-                { "include": "#type-body" }
+                {
+                    "begin": "//",
+                    "beginCaptures": {
+                        "0": {
+                            "name": "punctuation.definition.comment.midas"
+                        }
+                    },
+                    "end": "$",
+                    "name": "comment.line.double-slash.midas"
+                },
+                {
+                    "begin": "/\\*",
+                    "beginCaptures": {
+                        "0": {
+                            "name": "punctuation.definition.comment.midas"
+                        }
+                    },
+                    "end": "\\*/",
+                    "endCaptures": {
+                        "0": {
+                            "name": "punctuation.definition.comment.midas"
+                        }
+                    },
+                    "name": "comment.block.midas"
+                }
             ]
         },
-        "type-base": {
-            "begin": "(\\()([a-zA-Z_][a-zA-Z_\\d]*)(\\))",
-            "end": "$",
-            "beginCaptures": {
-                "1": {
-                    "name": "punctuation.definition.base.begin.midas"
-                },
-                "2": {
-                    "name": "variable.name"
-                },
-                "3": {
-                    "name": "punctuation.definition.base.end.midas"
-                }
-            },
-            "patterns": [
-                { "include": "#type-cond" }
-            ]
-        },
-        "type-cond": {
-            "begin": "where",
-            "end": "$",
+        "string": {
+            "begin": "\"",
             "beginCaptures": {
                 "0": {
-                    "name": "keyword.control.where.midas"
+                    "name": "punctuation.definition.string.begin.midas"
                 }
-            }
+            },
+            "end": "\"",
+            "endCaptures": {
+                "0": {
+                    "name": "punctuation.definition.string.end.midas"
+                }
+            },
+            "name": "string.quoted.double.c"
         },
-        "type-body": {
+        "alias-stmt": {
+            "begin": "\\b(alias)\\b",
+            "beginCaptures": {
+                "1": {
+                    "name": "keyword.declaration.midas"
+                }
+            },
+            "end": "$",
+            "name": "meta.declaration.alias.midas",
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "match": "[A-Za-z_][A-Za-z0-9_]*",
+                    "name": "entity.name.type.midas"
+                },
+                {
+                    "begin": "=",
+                    "beginCaptures": {
+                        "0": {
+                            "name": "keyword.operator.equal.midas"
+                        }
+                    },
+                    "end": "$",
+                    "patterns": [
+                        {
+                            "include": "#type-expr"
+                        }
+                    ]
+                }
+            ]
+        },
+        "type-stmt": {
+            "begin": "\\b(type)\\b",
+            "beginCaptures": {
+                "1": {
+                    "name": "keyword.declaration.midas"
+                }
+            },
+            "end": "$",
+            "name": "meta.declaration.type.midas",
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "match": "[A-Za-z_][A-Za-z0-9_]*",
+                    "name": "entity.name.type.midas"
+                },
+                {
+                    "begin": "\\[",
+                    "beginCaptures": {
+                        "0": {
+                            "name": "punctuation.section.brackets.begin.midas"
+                        }
+                    },
+                    "end": "\\]",
+                    "endCaptures": {
+                        "0": {
+                            "name": "punctuation.section.brackets.end.midas"
+                        }
+                    },
+                    "patterns": [
+                        {
+                            "include": "#comments"
+                        },
+                        {
+                            "include": "#type-params"
+                        }
+                    ]
+                },
+                {
+                    "begin": "=",
+                    "beginCaptures": {
+                        "0": {
+                            "name": "keyword.operator.equal.midas"
+                        }
+                    },
+                    "end": "$",
+                    "patterns": [
+                        {
+                            "include": "#type-expr"
+                        }
+                    ]
+                }
+            ]
+        },
+        "type-params": {
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "match": "<:",
+                    "name": "keyword.operator.subtype.midas"
+                },
+                {
+                    "match": "[A-Za-z][A-Za-z0-9_]*",
+                    "name": "entity.name.type.midas"
+                }
+            ]
+        },
+        "extend-stmt": {
+            "begin": "\\b(extend)\\b",
+            "beginCaptures": {
+                "1": {
+                    "name": "keyword.declaration.midas"
+                }
+            },
+            "end": "(?=\\{)|$",
+            "name": "meta.declaration.extend.midas",
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "match": "[A-Za-z_][A-Za-z0-9_]*",
+                    "name": "entity.name.type.midas"
+                },
+                {
+                    "begin": "\\[",
+                    "beginCaptures": {
+                        "0": {
+                            "name": "punctuation.section.brackets.begin.midas"
+                        }
+                    },
+                    "end": "\\]",
+                    "endCaptures": {
+                        "0": {
+                            "name": "punctuation.section.brackets.end.midas"
+                        }
+                    },
+                    "patterns": [
+                        {
+                            "include": "#comments"
+                        },
+                        {
+                            "include": "#type-params"
+                        }
+                    ]
+                }
+            ]
+        },
+        "extend-body": {
             "begin": "\\{",
-            "end": "\\}",
             "beginCaptures": {
                 "0": {
-                    "name": "punctuation.definition.type-body.begin.midas"
+                    "name": "punctuation.section.block.begin.midas"
                 }
             },
+            "end": "\\}",
             "endCaptures": {
                 "0": {
-                    "name": "punctuation.definition.type-body.end.midas"
+                    "name": "punctuation.section.block.end.midas"
                 }
             },
             "patterns": [
-                {"include": "#type-prop"},
-                {"include": "#comment"}
+                {
+                    "include": "#comments"
+                },
+                {
+                    "include": "#member-stmt"
+                }
             ]
         },
-        "type-prop": {
-            "match": "([a-zA-Z_][a-zA-Z_\\d]*)(:)\\s*([a-zA-Z_][a-zA-Z_\\d]*)",
-            "captures": {
-                "1": {
-                    "name": "variable.name"
+        "member-stmt": {
+            "patterns": [
+                {
+                    "include": "#comments"
                 },
-                "2": {
-                    "name": "punctuation.separator.annotation.midas"
-                },
-                "3": {
-                    "name": "meta.type.name"
+                {
+                    "begin": "\\b(prop|def)\\b",
+                    "beginCaptures": {
+                        "1": {
+                            "name": "keyword.other.midas"
+                        }
+                    },
+                    "end": "$",
+                    "patterns": [
+                        {
+                            "include": "#comments"
+                        },
+                        {
+                            "match": "[A-Za-z_][A-Za-z0-9_]*",
+                            "name": "variable.other.member.midas"
+                        },
+                        {
+                            "begin": ":",
+                            "beginCaptures": {
+                                "0": {
+                                    "name": "punctuation.separator.annotation.midas"
+                                }
+                            },
+                            "end": "$",
+                            "patterns": [
+                                {
+                                    "include": "#type-expr"
+                                }
+                            ]
+                        }
+                    ]
                 }
-            }
+            ]
         },
-        "extend-def": {
-            "begin": "\\b(extend)\\s*([a-zA-Z_][a-zA-Z_\\d]*)\\s+(\\{)",
-            "end": "\\}",
+        "predicate-stmt": {
+            "begin": "\\b(predicate)\\b",
             "beginCaptures": {
                 "1": {
-                    "name": "keyword.control.extend.midas"
-                },
-                "2": {
-                    "name": "variable.name"
-                },
-                "3": {
-                    "name": "punctuation.definition.extend-body.begin.midas"
+                    "name": "keyword.declaration.midas"
                 }
             },
-            "endCaptures": {
-                "0": {
-                    "name": "punctuation.definition.extend-body.end.midas"
-                }
-            },
-            "patterns": [
-                {"include": "#op-def"},
-                {"include": "#comment"}
-            ]
-        },
-        "op-def": {
-            "match": "\\b(op)\\s+(\\S+)\\s*\\(\\s*([a-zA-Z_][a-zA-Z_\\d]*)\\s*\\)\\s*(->)\\s*([a-zA-Z_][a-zA-Z_\\d]*)",
-            "captures": {
-                "1": {
-                    "name": "keyword.control.op.midas"
-                },
-                "2": {
-                    "name" : "keyword.operator"
-                },
-                "3": {
-                    "name" : "variable.name"
-                },
-                "4": {
-                    "name" : "keyword.operator.assignment"
-                },
-                "5": {
-                    "name" : "variable.name"
-                }
-            }
-        },
-        "pred-def": {
-            "begin": "(predicate)\\s+([a-zA-Z_][a-zA-Z_\\d]*)\\(([a-zA-Z_][a-zA-Z_\\d]*):\\s*([a-zA-Z_][a-zA-Z_\\d]*)\\)\\s*(=)",
             "end": "$",
-            "beginCaptures": {
-                "1": {
-                    "name": "keyword.control.pred.midas"
-                },
-                "2": {
-                    "name": "variable.name"
-                },
-                "3": {
-                    "name": "variable.name"
-                },
-                "4": {
-                    "name": "variable.name"
-                },
-                "5": {
-                    "name": "keyword.operator.assignment"
-                }
-            },
+            "name": "meta.declaration.predicate.midas",
             "patterns": [
-                { "include": "source.python" }
+                {
+                    "include": "#comments"
+                },
+                {
+                    "match": "[A-Za-z_][A-Za-z0-9_]*",
+                    "name": "entity.name.function.midas"
+                },
+                {
+                    "begin": "\\(",
+                    "beginCaptures": {
+                        "0": {
+                            "name": "punctuation.section.group.begin.midas"
+                        }
+                    },
+                    "end": "\\)",
+                    "endCaptures": {
+                        "0": {
+                            "name": "punctuation.section.group.end.midas"
+                        }
+                    },
+                    "patterns": [
+                        {
+                            "include": "#predicate-params"
+                        }
+                    ]
+                },
+                {
+                    "begin": "=",
+                    "beginCaptures": {
+                        "0": {
+                            "name": "keyword.operator.equal.midas"
+                        }
+                    },
+                    "end": "$",
+                    "patterns": [
+                        {
+                            "include": "#constraint"
+                        }
+                    ]
+                }
             ]
         },
-        "statement": {
+        "predicate-params": {
             "patterns": [
-                { "include": "#comment" },
-                { "include": "#type-def" },
-                { "include": "#extend-def" },
-                { "include": "#pred-def" }
+                {
+                    "include": "#comments"
+                },
+                {
+                    "begin": "([A-Za-z_][A-Za-z0-9_]*)(:)",
+                    "beginCaptures": {
+                        "1": {
+                            "name": "variable.parameter.midas"
+                        },
+                        "2": {
+                            "name": "punctuation.separator.annotation.midas"
+                        }
+                    },
+                    "end": "(?=,)|(?=\\))",
+                    "patterns": [
+                        {
+                            "include": "#type-expr"
+                        }
+                    ]
+                },
+                {
+                    "match": ",",
+                    "name": "punctuation.separator.midas"
+                }
+            ]
+        },
+        "type-expr": {
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "begin": "\\b(fn)\\s*(\\()",
+                    "beginCaptures": {
+                        "1": {
+                            "name": "keyword.other.midas"
+                        },
+                        "2": {
+                            "name": "punctuation.section.group.begin.midas"
+                        }
+                    },
+                    "end": "\\)",
+                    "endCaptures": {
+                        "0": {
+                            "name": "punctuation.section.group.end.midas"
+                        }
+                    },
+                    "patterns": [
+                        {
+                            "include": "#fn-params"
+                        }
+                    ]
+                },
+                {
+                    "match": "->",
+                    "name": "keyword.operator.arrow.midas"
+                },
+                {
+                    "begin": "\\b(where)\\b",
+                    "beginCaptures": {
+                        "1": {
+                            "name": "keyword.other.midas"
+                        }
+                    },
+                    "end": "$",
+                    "patterns": [
+                        {
+                            "include": "#constraint"
+                        }
+                    ]
+                },
+                {
+                    "begin": "(\\bFrame\\b)(\\s*)(\\[)",
+                    "beginCaptures": {
+                        "1": {
+                            "name": "entity.name.type.midas"
+                        },
+                        "3": {
+                            "name": "punctuation.section.brackets.begin.midas"
+                        }
+                    },
+                    "end": "\\]",
+                    "endCaptures": {
+                        "0": {
+                            "name": "punctuation.section.brackets.end.midas"
+                        }
+                    },
+                    "patterns": [
+                        {
+                            "include": "#frame-schema"
+                        }
+                    ]
+                },
+                {
+                    "match": "\\bFrame\\b",
+                    "name": "entity.name.type.midas"
+                },
+                {
+                    "match": "[A-Za-z_][A-Za-z0-9_]*",
+                    "name": "entity.name.type.midas"
+                },
+                {
+                    "begin": "\\[",
+                    "beginCaptures": {
+                        "0": {
+                            "name": "punctuation.section.brackets.begin.midas"
+                        }
+                    },
+                    "end": "\\]",
+                    "endCaptures": {
+                        "0": {
+                            "name": "punctuation.section.brackets.end.midas"
+                        }
+                    },
+                    "patterns": [
+                        {
+                            "include": "#comments"
+                        },
+                        {
+                            "include": "#type-expr"
+                        },
+                        {
+                            "match": ",",
+                            "name": "punctuation.separator.midas"
+                        }
+                    ]
+                }
+            ]
+        },
+        "fn-params": {
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "begin": "([A-Za-z_][A-Za-z0-9_]*)(:)",
+                    "beginCaptures": {
+                        "1": {
+                            "name": "variable.parameter.midas"
+                        },
+                        "2": {
+                            "name": "punctuation.separator.annotation.midas"
+                        }
+                    },
+                    "end": "(?=,)|(?=\\))",
+                    "patterns": [
+                        {
+                            "include": "#type-expr"
+                        },
+                        {
+                            "match": "\\?",
+                            "name": "keyword.operator.qmark.midas"
+                        }
+                    ]
+                },
+                {
+                    "match": ",",
+                    "name": "punctuation.separator.midas"
+                },
+                {
+                    "include": "#type-expr"
+                }
+            ]
+        },
+        "constraint": {
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "match": "\\d+(\\.\\d+)?",
+                    "name": "constant.numeric.midas"
+                },
+                {
+                    "match": "\\b(true|false|none)\\b",
+                    "name": "constant.language.midas"
+                },
+                {
+                    "include": "#string"
+                },
+                {
+                    "match": "(<=|>=|<|>|==|!=|&)",
+                    "name": "keyword.operator.midas"
+                },
+                {
+                    "match": "_",
+                    "name": "variable.language.midas"
+                },
+                {
+                    "match": "[A-Za-z_][A-Za-z0-9_]*(?=\\s*\\()",
+                    "name": "variable.function.midas"
+                },
+                {
+                    "match": "[A-Za-z_][A-Za-z0-9_]*",
+                    "name": "variable.other.readwrite.midas"
+                }
+            ]
+        },
+        "frame-schema": {
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "match": "[A-Za-z_][A-Za-z0-9_]*",
+                    "name": "variable.other.member.midas"
+                },
+                {
+                    "begin": ":",
+                    "beginCaptures": {
+                        "0": {
+                            "name": "punctuation.separator.annotation.midas"
+                        }
+                    },
+                    "end": "(?=,)|(?=\\])",
+                    "patterns": [
+                        {
+                            "include": "#type-expr"
+                        }
+                    ]
+                },
+                {
+                    "match": ",",
+                    "name": "punctuation.separator.midas"
+                }
             ]
         }
     }