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

Reviewed-on: #36

.gitea/workflows/tests.yaml

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

.gitignore

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

README.md

@@ -51,16 +51,7 @@ This framework is being developed as part of a Bachelor's Thesis by Louis Herede
 
 ## Commands
 
-<!--
+Hereafter is a description of the commands you can use with Midas. For a full description, refer to [the manual](./docs/manual.pdf) or `midas <subcommand> --help`.
-check
-compile
-format
-highlight
-parse
-dump_registry
-types
-validate
--->
 
 ### 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`

docs/fonts.yaml

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

docs/manual.typ

@@ -6,10 +6,20 @@
 #import "@preview/codly-languages:0.1.10": codly-languages
 #import "template.typ": TODO, project
 #import "@preview/gentle-clues:1.3.1" as gc
+#import "@preview/treet:1.0.0": tree-list
 
 #let midas-version = toml("../pyproject.toml").project.version
-#let head-ref = read("../.git/HEAD").split(":").at(1).trim()
+
-#let commit-hash = read("../.git/" + head-ref).slice(0, 8)
+#let commit-hash = if "hash" in sys.inputs {
+  sys.inputs.hash.slice(0, 8)
+} else {
+  let head-ref = read("../.git/HEAD").split(":").at(1, default: "").trim()
+  if head-ref.len() != 0 {
+    read("../.git/" + head-ref).slice(0, 8)
+  } else {
+    none
+  }
+}
 
 #show: project.with(
   title: [Midas User Manual],
@@ -371,6 +381,7 @@ Notice that you don't need to specify `Column` types.
     height: float where _ >= 0
   ]
   ```,
+  caption: [Simple frame type definition],
 ) <simple-frame>
 
 #pagebreak()
@@ -773,18 +784,213 @@ The example of @python-frame-type shows how you can annotate a dataframe with so
 
 = Commands <commands>
 
-#TODO
+Midas offers several commands to parse, check and compile your code.
+
+All commands presented in this chapter are subcommands of the `midas` command.
+For example:
+#figure(
+  ```sh
+  midas compile script.py -t types.midas
+  ```,
+  caption: [Example midas command],
+)
 
 == Type Checking (`check`) <cmd-check>
+
+#figure(
+  ```sh
+  midas check -t types.midas source.py
+  ```,
+  caption: [Example usage of `midas check`],
+)
+
+This command parses the given files and run the type checkers against the Midas definitions and Python program. Diagnostics are then printed showing warnings and errors.
+
+=== Parameters
+
+- `FILE`: the Python file to type check
+
+=== Options
+
+- `-t` / `--types` `FILENAME`: Midas files defining type definitions. This option can be used multiple times. Files should end with `.midas`
+- `-l` / `--highlight` `FILENAME`: if set, a highlighted version of the Python file showing inline diagnostics will be written to the given file
+
 == Compiling (`compile`) <cmd-compile>
+
+
+#figure(
+  ```sh
+  midas compile -t types.midas source.py
+  ```,
+  caption: [Example usage of `midas compile`],
+)
+
+With the `compile` command, you can process a source Python file, with any number of custom type definition files, and the type checker will verify the coherence of your program and generate the runnable code with valid syntax and runtime assertions.
+
+=== Parameters
+
+- `FILE`: the Python file to compile
+
+=== Options
+
+#let output = {
+  set text(size: .8em)
+  tree-list[
+    - build/
+      - midas/
+        - stubs1.py
+        - misc.py
+        - source.py
+  ]
+}
+
+- `-t` / `--types` `FILENAME`: Midas files defining type definitions. This option can be used multiple times. Files should end with `.midas`
+- `-s` / `--stubs` `TEXT`: the file name of generated stub files. Each instance of this option maps to an instance of `-t`.
+  For example:
+  ```sh
+  midas -t types1.midas -s stubs1 -t types2.midas -s misc source.py
+  ```
+  will compile to the following files:
+
+  #output
+- `--ignore-errors`: if set, error diagnostics will not prevent the generator from running
+
 == Formatting (`format`) <cmd-format>
+
+#figure(
+  ```sh
+  midas format types.midas
+  midas format types.midas -o formatted.midas
+  ```,
+  caption: [Example usage of `midas format`],
+)
+
+This command parses the given Midas file and outputs a pretty printed file from the AST.
+
+=== Parameters
+
+- `FILE`: the Midas file to format
+
+=== Options
+
+- `-o` / `--output` `FILENAME`: if set, the output is written to the given file instead of STDOUT
+
 == Highlighting (`highlight`) <cmd-highlight>
+
+#figure(
+  ```sh
+  midas highlight source.py
+  midas highlight source.py -o highlighted.html
+  midas highlight types.midas
+  midas highlight types.midas -o highlighted.html
+  ```,
+  caption: [Example usage of `midas highlight`],
+)
+
+The `highlight` command takes in a source file (Python or Midas), runs the appropriate parser and outputs an HTML file containing the source code with added highlighting. This highlighting takes the form of hoverable annotations showing some of the parsed structures (e.g. a function definition, an assignment, a generic type, etc.)
+
+=== Parameters
+
+- `FILE`: the Python or Midas file to highlight
+
+=== Options
+
+- `-o` / `--output` `FILENAME`: if set, the output is written to the given file instead of STDOUT
+
 == Dumping the AST (`parse`) <cmd-parse>
+
+#figure(
+  ```sh
+  midas parse source.py
+  midas parse types.midas
+  ```,
+  caption: [Example usage of `midas parse`],
+)
+
+For debugging purposes, you can use this command to output the AST parsed from a Python or Midas file
+
+=== Parameters
+
+- `FILE`: the Python or Midas file to parse
+
+=== Options
+
+- `--raw`: if `FILE` is a Python file, the raw AST is returned, as produced by the builtin `ast` module, instead of the custom AST nodes used by Midas internally. This flag has no effect on Midas files
+
 == Dumping the Registry (`dump-registry`) <cmd-registry>
+
+#figure(
+  ```sh
+  midas dump-registry -t types.midas
+  ```,
+  caption: [Example usage of `midas dump-registry`],
+)
+
+This command processes the given Midas definitions and dumps the contents of the types registry.
+
+=== Options
+
+- `-t` / `--types` `FILENAME`: Midas files defining type definitions. This option can be used multiple times. Files should end with `.midas`
+
 == Generating Stubs (`stubs`) <cmd-stubs>
+
+#figure(
+  ```sh
+  midas stubs types.midas
+  ```,
+  caption: [Example usage of `midas stubs`],
+)
+
+This command generates Python stubs from a Midas definition file
+
+=== Parameters
+
+- `FILE`: the source Midas file
+
+=== Options
+
+- `-o` / `--output` `FILENAME`: if set, the stubs are written to the given file instead of the default `FILE.pyi`
+- `-w` / `--watch`: if set, the input file will be watched and any changes to it will regenerate the stubs, otherwise the command exits after generating the stubs once
+
 == Showing Type Judgements (`types`) <cmd-types>
+
+#figure(
+  ```sh
+  midas types -t types.midas source.py
+  ```,
+  caption: [Example usage of `midas types`],
+)
+
+This command type checks the given Python source file and logs all typing judgements made by the type checker.
+
+=== Parameters
+
+- `FILE`: the Python file to type check
+
+=== Options
+
+- `-t` / `--types` `FILENAME`: Midas files defining type definitions. This option can be used multiple times. Files should end with `.midas`
+- `-l` / `--highlight` `FILENAME`: if set, a highlighted version of the Python file showing inline diagnostics will be written to the given file
+
 == Validating Definitions (`validate`) <cmd-validate>
 
+#figure(
+  ```sh
+  midas validate types.midas
+  ```,
+  caption: [Example usage of `midas validate`],
+)
+
+This command lets you validate a Midas definition file by running the parser and type checker, verifying syntax and references.
+
+=== Parameters
+
+- `FILE`: the Midas file to validate
+
+=== Options
+
+- `-l` / `--highlight` `FILENAME`: if set, a highlighted version of the Midas file showing inline diagnostics will be written to the given file
+
 = Known limitations <limitations>
 
 == Eager evaluation in runtime assertions <eager-eval>

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

@@ -52,7 +52,7 @@
   title: none,
   author: none,
   version: "0.0.1",
-  hash: "abcdefgh",
+  hash: none,
   icon-path: none,
   doc,
 ) = {
@@ -69,13 +69,17 @@
   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()
       }
 
-      v#version - #hash
+      #version-name
 
       #if icon-path != none {
         v(1cm)

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/midas.py

@@ -136,15 +136,6 @@ class ConstraintType:
     constraint: Expr
 
 
-class ComplexType:
-    members: list[MemberStmt]
-
-
-class ExtensionType:
-    base: Type
-    extension: ComplexType
-
-
 class FunctionType:
     params: ParamSpec
     returns: Type

gen/python.py

@@ -24,6 +24,17 @@ class ParamSpec:
         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
+
+
 ###<
 
 
@@ -99,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
 

midas/ast/midas.py

@@ -256,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: ...
 
@@ -295,23 +289,6 @@ 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

midas/ast/printer/midas.py

@@ -124,19 +124,6 @@ class MidasPrinter(
         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)}"

midas/ast/printer/midas_ast.py

@@ -177,21 +177,6 @@ class MidasAstPrinter(
             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_sequence("members", type.members, last=True)
-
-    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():

midas/ast/python.py

@@ -25,6 +25,17 @@ class ParamSpec:
         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 #
 ####################
@@ -123,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: ...
 
@@ -207,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

midas/checker/diagnostic.py

@@ -18,9 +18,6 @@ class Diagnostic:
 
     Holds a location, a diagnostic type and a message.
     Optionally bound to a file path
-
-    Returns:
-        _type_: _description_
     """
 
     file_path: Optional[str]
@@ -30,7 +27,7 @@ class Diagnostic:
 
     @property
     def location_str(self) -> str:
-        """The diagnostic type and location as a human readable string
+        """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"
@@ -39,7 +36,7 @@ class Diagnostic:
         If the location's end is not specified, the formulation "at L<start_line>:<start_col>" is used.
 
         Returns:
-            str: _description_
+            str: the formatted type and location string
         """
 
         start_loc: str = f"L{self.location.lineno}:{self.location.col_offset+1}"

midas/checker/dispatcher.py

@@ -26,10 +26,13 @@ class HasLocation(Protocol):
 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
@@ -37,11 +40,15 @@ class MappedArgument(Generic[E]):
 
 @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"
@@ -50,16 +57,28 @@ class CallError(StrEnum):
 
 @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:
@@ -68,6 +87,15 @@ class CallResult:
 
 
 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
@@ -86,22 +114,21 @@ class CallDispatcher(Generic[E]):
     ) -> CallResult:
         """Get the result type of a function call
 
-        If the function has overloads, the function will try to resolve the
+        If the callee has overloads, this function will try to resolve the
         appropriate signature.
-        Argument types are matched to the defined parameters.
+        Argument types are matched with the defined parameters.
-        The function doesn't take the raw expression as a parameter to accommodate
+        This function doesn't take the raw expression as a parameter to
-        for desugared calls such as for operators.
+        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
+            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:
-            Type: the return type of the call, or `None` if either
+            CallResult: the call result, either a type or an error
-            the call is invalid or no overload matched the arguments uniquely
         """
         match callee:
             case Function() as function:
@@ -179,6 +206,18 @@ class CallDispatcher(Generic[E]):
         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)
@@ -219,6 +258,12 @@ class CallDispatcher(Generic[E]):
         """
         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(
@@ -246,8 +291,9 @@ class CallDispatcher(Generic[E]):
             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
+            Union[tuple[Function, None], tuple[None, str]]: a tuple containing
-            determined unambiguously, or `None`.
+                the resolved function signature if it can be determined
+                unambiguously, or `None`, and an error message, or `None`
         """
         candidates: list[OverloadCandidate] = []
         errors: list[CallError] = []
@@ -464,8 +510,8 @@ class CallDispatcher(Generic[E]):
         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
+        This is used to check whether a given overload is a more specific
-        a more specific function/ a subtype of another.
+        function / a subtype of another.
 
         Args:
             mapped1 (list[MappedArgument]): the first argument mappings (subtype)

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,7 +221,23 @@ class Evaluator(m.Expr.Visitor[Any]):
         self.scopes.pop()
         return res
 
-    def _map_args(self, function: Function, args: list[Any], kwargs: dict[str, Any]):
+    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
         )
@@ -166,9 +246,20 @@ class Evaluator(m.Expr.Visitor[Any]):
         }
 
         for i, arg in enumerate(args):
+            if i >= len(positional):
+                if self.reporter is not None:
+                    self.reporter.error(
+                        location,
+                        f"Too many positional arguments, expected at most {len(positional)}, got {len(args)}",
+                    )
+                break
             param: Function.Parameter = positional[i]
             self.set_value(param.name, arg)
 
         for name, arg in kwargs.items():
+            if name not in keywords:
+                if self.reporter is not None:
+                    self.reporter.error(location, f"Unknown keyword argument '{name}'")
+                break
             param: Function.Parameter = keywords[name]
             self.set_value(param.name, arg)

midas/checker/frames/column_groupby_methods.py

@@ -12,8 +12,8 @@ from midas.checker.types import (
     ColumnType,
     Function,
     ParamSpec,
-    TopType,
     Type,
+    UnknownType,
 )
 
 if TYPE_CHECKING:
@@ -22,6 +22,8 @@ if TYPE_CHECKING:
 
 @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
@@ -35,6 +37,8 @@ class Call:
 
 
 class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
+    """The method registry for column group-by types"""
+
     NAMED_ARGS: dict[str, str] = {
         "numeric_only": "bool",
         "skipna": "bool",
@@ -45,10 +49,22 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def _aggregate(
         self,
         call: Call,
+        method: str,
         params: list[str | tuple[str, str, bool]] = [],
-        *,
-        preserve_inner_type: bool = False,
     ) -> 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:
@@ -68,13 +84,21 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
                     )
             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=returns,
-                call.groupby.column
-                if preserve_inner_type
-                else ColumnType(type=TopType())
-            ),
         )
 
         result: CallResult = self.dispatcher.get_result(
@@ -89,6 +113,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def kurt(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "kurt",
             ["skipna", "numeric_only"],
         )
 
@@ -96,6 +121,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def max(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "max",
             [
                 "numeric_only",
                 (
@@ -107,13 +133,13 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
                 "engine",
                 "engine_kwargs",
             ],
-            preserve_inner_type=True,
         )
 
     @method()
     def mean(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "mean",
             ["numeric_only", "skipna", "engine", "engine_kwargs"],
         )
 
@@ -121,14 +147,15 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def median(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "median",
             ["numeric_only", "skipna"],
-            preserve_inner_type=True,
         )
 
     @method()
     def min(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "min",
             [
                 "numeric_only",
                 (
@@ -140,13 +167,13 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
                 "engine",
                 "engine_kwargs",
             ],
-            preserve_inner_type=True,
         )
 
     @method()
     def prod(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "prod",
             [
                 "numeric_only",
                 (
@@ -162,6 +189,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def std(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "std",
             [
                 (
                     "ddof",
@@ -179,6 +207,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def sum(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "sum",
             [
                 "numeric_only",
                 (
@@ -196,6 +225,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def var(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "var",
             [
                 (
                     "var",

midas/checker/frames/column_manager.py

@@ -4,17 +4,28 @@ 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.types import ColumnGroupBy, ColumnType, Type
+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)
@@ -22,6 +33,62 @@ class ColumnManager:
             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,
@@ -32,6 +99,20 @@ class ColumnManager:
         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,
@@ -52,6 +133,20 @@ class ColumnManager:
         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,
@@ -63,6 +158,15 @@ class ColumnManager:
         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":

midas/checker/frames/column_methods.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import ast
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Callable, Optional, TypeAlias, Union
 
 import midas.ast.python as p
 from midas.ast.location import Location
@@ -12,11 +12,11 @@ from midas.checker.types import (
     ColumnGroupBy,
     ColumnType,
     Function,
-    GenericType,
+    OverloadedFunction,
     ParamSpec,
     TopType,
     Type,
-    TypeVar,
+    UnitType,
     UnknownType,
     unfold_type,
 )
@@ -24,9 +24,36 @@ from midas.checker.types import (
 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
@@ -40,7 +67,74 @@ class Call:
 
 
 class ColumnMethodRegistry(MethodRegistry[Call]):
-    def _element_binary_op(self, call: Call, method: str) -> ColumnType:
+    """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
@@ -51,50 +145,63 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
             method (str): the method name
 
         Returns:
-            ColumnType: the resulting column type
+            tuple[Type, bool]: the resulting type and a boolean indicating
+                whether the operand is a column
         """
-        column2: Optional[ColumnType] = None
+        if len(call.positional) == 0:
+            return UnknownType(), False
 
         col_type1: Type = call.column.type
-        new_column: Type = ColumnType(type=UnknownType())
+        operand: TypedExpr = call.positional[0]
-        if len(call.positional) != 0:
+        unfolded_operand: Type = unfold_type(operand[1])
-            other: Type = call.positional[0][1]
+        col_type2: Type
-            unfolded_other: Type = unfold_type(other)
+
-            if isinstance(unfolded_other, ColumnType):
+        column_operand: bool = isinstance(unfolded_operand, ColumnType)
-                column2 = unfolded_other
+
-                col_type2: Type = column2.type
+        # 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=(call.positional[0][0], col_type2),
+            right=(operand[0], col_type2),
             method=method,
         )
-                new_column = ColumnType(type=new_inner_type)
+        return ColumnType(type=new_inner_type), column_operand
-        return new_column
 
     def _element_wise(self, call: Call, method: str) -> Type:
-        # TODO: support add with scalar
+        """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
-        param_type: TypeVar = TypeVar(name="T", bound=None)
+        returns, column_operand = self._element_binary_op(call, method)
-        signature = GenericType(
+        signature = Function(
-            name="add",
-            params=[param_type],
-            body=Function(
             params=ParamSpec(
                 mixed=[
                     Function.Parameter(
                         pos=0,
                         name="other",
-                            type=ColumnType(type=param_type),
+                        type=TopType(),
                         required=True,
                     ),
                 ],
             ),
-                returns=self._element_binary_op(call, method),
+            returns=returns,
-            ),
         )
 
         # Map arguments and compute result type
@@ -104,13 +211,83 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
             positional=call.positional,
             keywords=call.keywords,
         )
-        if result.is_valid:
+        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__")
@@ -168,8 +345,31 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         call: Call,
         kwargs: list[Function.Parameter] = [],
         *,
-        preserve_inner_type: bool = False,
+        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=[
@@ -182,7 +382,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
                     *kwargs,
                 ],
             ),
-            returns=call.column if preserve_inner_type else ColumnType(type=TopType()),
+            returns=returns,
         )
 
         result: CallResult = self.dispatcher.get_result(
@@ -199,27 +399,29 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
 
     @method()
     def max(self, call: Call) -> Type:
-        return self._aggregate(call, preserve_inner_type=True)
+        return self._aggregate(call, formula=lambda t: t)
 
     @method()
     def mean(self, call: Call) -> Type:
-        return self._aggregate(call)
+        return self._aggregate(
+            call, formula=lambda t: ((t, "__add__", t), "__truediv__", "int")
+        )
 
     @method()
     def median(self, call: Call) -> Type:
-        return self._aggregate(call, preserve_inner_type=True)
+        return self._aggregate(call, formula=lambda t: t)
 
     @method()
     def min(self, call: Call) -> Type:
-        return self._aggregate(call, preserve_inner_type=True)
+        return self._aggregate(call, formula=lambda t: t)
 
     @method()
     def mode(self, call: Call) -> Type:
-        return self._aggregate(call, preserve_inner_type=True)
+        return self._aggregate(call, formula=lambda t: t)
 
     @method("product", "prod")
     def product(self, call: Call) -> Type:
-        return self._aggregate(call)
+        return self._aggregate(call, formula=lambda t: (t, "__mul__", t))
 
     @method()
     def std(self, call: Call) -> Type:
@@ -237,7 +439,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
 
     @method()
     def sum(self, call: Call) -> Type:
-        return self._aggregate(call)
+        return self._aggregate(call, formula=lambda t: (t, "__add__", t))
 
     @method()
     def var(self, call: Call) -> Type:
@@ -301,6 +503,79 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         )
         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")
@@ -344,6 +619,14 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         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

midas/checker/frames/frame_groupby_methods.py

@@ -21,6 +21,8 @@ if TYPE_CHECKING:
 
 @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
@@ -34,17 +36,22 @@ class Call:
 
 
 class FrameGroupByMethodRegistry(MethodRegistry[Call]):
-    NAMED_ARGS: dict[str, str] = {
+    """The method registry for frame group-by types"""
-        "numeric_only": "bool",
-        "skipna": "bool",
-        "engine": "str",
-        "engine_kwargs": "dict",
-    }
 
     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,

midas/checker/frames/frame_manager.py

@@ -24,10 +24,20 @@ if TYPE_CHECKING:
 
 
 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)
@@ -43,6 +53,18 @@ class FrameManager:
         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)
@@ -93,12 +115,24 @@ class FrameManager:
         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 frame
+            return self._set_column(frame, name, ColumnType(type=UnknownType()))
         return self._set_column(frame, name, type)
 
     def get(
@@ -108,6 +142,17 @@ class FrameManager:
         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)
@@ -142,6 +187,17 @@ class FrameManager:
         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():
@@ -159,6 +215,16 @@ class FrameManager:
     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
@@ -185,12 +251,31 @@ class FrameManager:
     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
@@ -200,6 +285,15 @@ class FrameManager:
     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(
@@ -212,6 +306,20 @@ class FrameManager:
         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,
@@ -232,6 +340,20 @@ class FrameManager:
         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,
@@ -243,6 +365,15 @@ class FrameManager:
         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":

midas/checker/frames/frame_methods.py

@@ -17,6 +17,7 @@ from midas.checker.types import (
     ParamSpec,
     TopType,
     Type,
+    UnitType,
     UnknownType,
     unfold_type,
 )
@@ -27,6 +28,8 @@ if TYPE_CHECKING:
 
 @dataclass(frozen=True, kw_only=True)
 class Call:
+    """A frame method call, implements :class:`utils.MethodCall`"""
+
     location: Location
     call_expr: p.Expr
     frame: DataFrameType
@@ -40,6 +43,27 @@ class Call:
 
 
 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,
@@ -78,7 +102,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
             return ColumnType(type=UnknownType())
         return result
 
-    def _element_binary_op(self, call: Call, method: str) -> DataFrameType:
+    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
@@ -91,21 +115,24 @@ class FrameMethodRegistry(MethodRegistry[Call]):
             method (str): the method name
 
         Returns:
-            DataFrameType: the resulting frame type
+            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 there is at least 1 operand, which is a dataframe
+        # Get map of operand's columns by name, if the operand is a dataframe
-        if len(call.positional) != 0:
-            operand: TypedExpr = call.positional[0]
         unfolded_other: Type = unfold_type(operand[1])
-            if isinstance(unfolded_other, DataFrameType):
+        frame_operand: bool = isinstance(unfolded_other, DataFrameType)
+        if frame_operand:
             frame2 = unfolded_other
-                by_name = {
+            by_name = {col.name: col for col in frame2.columns if col.name is not None}
-                    col.name: col for col in frame2.columns if col.name is not None
-                }
 
         # Compute new schema:
         # Step 1: for all columns in frame1:
@@ -118,11 +145,24 @@ class FrameMethodRegistry(MethodRegistry[Call]):
 
             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: ColumnType = column2.type
+                    col_type2 = column2.type
 
-                col_type = self._get_method_result(call, col_type1, col_type2, method)
+            # 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,
@@ -145,10 +185,24 @@ class FrameMethodRegistry(MethodRegistry[Call]):
                     )
                 )
 
-        return DataFrameType(columns=new_columns)
+        return DataFrameType(columns=new_columns), frame_operand
 
     def _element_wise(self, call: Call, method: str) -> Type:
-        # TODO: support scalar, sequence, Series, dict operand
+        """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(
@@ -156,12 +210,12 @@ class FrameMethodRegistry(MethodRegistry[Call]):
                     Function.Parameter(
                         pos=0,
                         name="other",
-                        type=DataFrameType(columns=[]),
+                        type=TopType(),
                         required=True,
                     ),
                 ],
             ),
-            returns=self._element_binary_op(call, method),
+            returns=returns,
         )
 
         # Map arguments and compute result type
@@ -171,13 +225,83 @@ class FrameMethodRegistry(MethodRegistry[Call]):
             positional=call.positional,
             keywords=call.keywords,
         )
-        if result.is_valid:
+        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__")
@@ -231,6 +355,16 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         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=[
@@ -382,8 +516,144 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         )
         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(
@@ -413,7 +683,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
                     )
                 ],
             ),
-            returns=FrameGroupBy(frame=call.frame),
+            returns=FrameGroupBy(frame=new_frame),
         )
 
         result: CallResult = self.dispatcher.get_result(
@@ -425,6 +695,14 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         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

midas/checker/frames/utils.py

@@ -24,6 +24,12 @@ if TYPE_CHECKING:
 
 
 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__(
@@ -42,6 +48,11 @@ class _MethodRegistryMeta(type):
 
 
 class MethodCall(Protocol):
+    """A method call object
+
+    Must have at least `location`, `call_expr` and `subject` properties
+    """
+
     @property
     def location(self) -> Location: ...
 
@@ -56,6 +67,8 @@ 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
 
@@ -76,6 +89,15 @@ class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta):
         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(
@@ -90,6 +112,13 @@ 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:

midas/checker/midas.py

@@ -1,5 +1,4 @@
 import logging
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
@@ -14,11 +13,9 @@ from midas.checker.registry import TypesRegistry
 from midas.checker.reporter import FileReporter, Reporter
 from midas.checker.types import (
     ColumnType,
-    ComplexType,
     ConstraintType,
     DataFrameType,
     DerivedType,
-    ExtensionType,
     Function,
     GenericType,
     ParamSpec,
@@ -29,27 +26,10 @@ from midas.checker.types import (
 )
 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
 
 
-class ReturnException(Exception):
-    pass
-
-
-@dataclass(frozen=True, kw_only=True)
-class MappedArgument:
-    expr: m.Expr
-    type: Type
-    argument: Function.Parameter
-
-
-@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"""
 
@@ -76,10 +56,21 @@ 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]):
+        """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)
 
@@ -92,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
 
@@ -112,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)
@@ -139,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}")
@@ -215,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)
@@ -240,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)
@@ -262,6 +295,11 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         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}")
@@ -348,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:
@@ -357,24 +405,21 @@ 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:
         return Function(
             params=self._visit_param_spec(type.params),

midas/checker/preamble.py

@@ -23,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[..., 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)

midas/checker/python.py

@@ -1,6 +1,5 @@
 import ast
 import logging
-from dataclasses import dataclass
 from typing import Any, Optional
 
 import midas.ast.python as p
@@ -56,19 +55,6 @@ class UndefinedMethodException(Exception):
     pass
 
 
-@dataclass(frozen=True, kw_only=True)
-class MappedArgument:
-    expr: p.Expr
-    type: Type
-    argument: Function.Parameter
-
-
-@dataclass(frozen=True, kw_only=True)
-class OverloadCandidate:
-    function: Function
-    mapped: list[MappedArgument]
-
-
 class PythonTyper(
     p.Stmt.Visitor[None],
     p.Expr.Visitor[Type],
@@ -97,17 +83,31 @@ class PythonTyper(
         self.assertions: AssertionCollector = AssertionCollector()
 
     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(self.reporter)
 
     def process(self, source: str, path: Optional[str]) -> TypedAST:
+        """Process some Python source code
+
+        Args:
+            source (str): the Python source code
+            path (Optional[str]): the path of the source file, if known
+
+        Returns:
+            TypedAST: all generated typechecking information
+        """
         reporter: FileReporter = self.reporter.for_file(path)
         self.set_reporter(reporter)
 
         tree: ast.Module = ast.parse(source, filename=path or "<unknown>")
         parser = PythonParser()
         stmts: list[p.Stmt] = parser.parse_module(tree)
-        resolver = Resolver()
+        resolver = Resolver(reporter)
         resolver.resolve(*stmts)
 
         self.env = self.global_env
@@ -125,7 +125,7 @@ class PythonTyper(
         )
 
     def judge(self, expr: p.Expr, type: Type):
-        """Record a typing judgement
+        """Record a typing judgement for the given expression
 
         Args:
             expr (p.Expr): the judged expression
@@ -134,7 +134,7 @@ class PythonTyper(
         self.judgements.append((expr, type))
 
     def compute_type(self, expr: p.Expr) -> Type:
-        """Evaluate the type of an expression
+        """Evaluate the type of the given expression
 
         Args:
             expr (p.Expr): the expression to type
@@ -145,7 +145,7 @@ class PythonTyper(
         return expr.accept(self)
 
     def type_of(self, expr: p.Expr) -> Type:
-        """Evaluate the type of an expression and record the judgement
+        """Evaluate the type of the given expression and record the judgement
 
         Args:
             expr (p.Expr): the expression to evaluate
@@ -158,9 +158,22 @@ class PythonTyper(
         return type
 
     def resolve_type_expr(self, expr: p.MidasType) -> Type:
+        """Resolve the type of a type expression (annotation)
+
+        Args:
+            expr (p.MidasType): the type expression
+
+        Returns:
+            Type: the resolved type
+        """
         return expr.accept(self)
 
     def process_stmt(self, stmt: p.Stmt) -> None:
+        """Type check the given statement
+
+        Args:
+            stmt (p.Stmt): the statement to type-check
+        """
         stmt.accept(self)
 
     def process_block(self, block: list[p.Stmt], env: Environment) -> bool:
@@ -224,8 +237,29 @@ class PythonTyper(
         positional: list[TypedExpr],
         keywords: dict[str, TypedExpr],
     ) -> Type:
+        """Evaluate a method call on an object
+
+        Calls to dataframes and columns types are delegated to the appropriate manager
+
+        Args:
+            location (Location): the location of the call
+            call_expr (p.Expr): the call expression
+            obj (TypedExpr): the object on which the method is called
+            method_name (str): the method name
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Raises:
+            UndefinedMethodException: if the method is not defined
+
+        Returns:
+            Type: the return type of the call
+        """
         unfolded: Type = unfold_type(obj[1])
         match unfolded:
+            case TopType() | UnknownType():
+                return UnknownType()
+
             case DataFrameType():
                 return self.frame_mgr.call(
                     method=method_name,
@@ -283,6 +317,15 @@ class PythonTyper(
         return result.result
 
     def is_subtype(self, type1: Type, type2: Type) -> bool:
+        """Check whether `type1` is a subtype of `type2`
+
+        Args:
+            type1 (Type): the potential "subtype"
+            type2 (Type): the potential "supertype"
+
+        Returns:
+            bool: whether `type1` is a subtype of `type2`
+        """
         return self.types.is_subtype(type1, type2)
 
     def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
@@ -295,10 +338,23 @@ class PythonTyper(
         kw: list[Function.Parameter] = []
 
         def eval_param_type(param: p.Function.Parameter) -> Type:
-            if param.type is not None:
+            default_type: Optional[Type] = None
-                return self.resolve_type_expr(param.type)
             if param.default is not None:
-                return self.type_of(param.default)
+                default_type = self.type_of(param.default)
+
+            if param.type is not None:
+                param_type: Type = self.resolve_type_expr(param.type)
+                if default_type is not None:
+                    if not self.types.is_subtype(default_type, param_type):
+                        self.reporter.error(
+                            param.location or stmt.location,
+                            f"Cannot use default value of type {default_type} for parameter of type {param_type}",
+                        )
+                return param_type
+
+            if default_type is not None:
+                return default_type
+
             return UnknownType()
 
         position: int = 0
@@ -408,6 +464,15 @@ class PythonTyper(
             self._assign(stmt.location, target, value_type)
 
     def _assign(self, location: Location, target: p.Expr, value_type: Type):
+        """Handle an assignment to the given target
+
+        Delegate to the appropriate method according to the target type
+
+        Args:
+            location (Location): the location of the assignment
+            target (p.Expr): the assignment's target
+            value_type (Type): the value to be assigned
+        """
         match target:
             case p.VariableExpr():
                 self._assign_var(location, target, value_type)
@@ -429,6 +494,13 @@ class PythonTyper(
                     )
 
     def _assign_var(self, location: Location, target: p.VariableExpr, value_type: Type):
+        """Type check assignment to the given target
+
+        Args:
+            location (Location): the location of the assignment
+            target (p.VariableExpr): the assignment's target
+            value_type (Type): the value to be assigned
+        """
         name: str = target.name
         var_type: Optional[Type] = self.look_up_variable(name, target)
 
@@ -447,6 +519,13 @@ class PythonTyper(
     def _assign_attr(
         self, location: Location, object: p.Expr, name: str, value_type: Type
     ):
+        """Type check assignment to the given target
+
+        Args:
+            location (Location): the location of the assignment
+            target (p.VariableExpr): the assignment's target
+            value_type (Type): the value to be assigned
+        """
         object_type: Type = self.type_of(object)
         member: Optional[Type] = self.types.lookup_member(object_type, name)
         if member is None:
@@ -466,6 +545,13 @@ class PythonTyper(
         index: p.Expr,
         value_type: Type,
     ):
+        """Type check assignment to the given target
+
+        Args:
+            location (Location): the location of the assignment
+            target (p.VariableExpr): the assignment's target
+            value_type (Type): the value to be assigned
+        """
         var_type: Type = self.type_of(var)
         unfolded_type: Type = unfold_type(var_type)
         # TODO: what happens if type is an alias of a dataframe type
@@ -515,6 +601,10 @@ class PythonTyper(
         pass
 
     def visit_for_stmt(self, stmt: p.ForStmt) -> None:
+        outer_env: Environment = self.env
+        inner_env: Environment = Environment(self.env)
+        self.env = inner_env
+
         item_type: Type = UnknownType()
         iterator_type: Type = self.type_of(stmt.iterator)
         if iterator_type != UnknownType():
@@ -528,11 +618,30 @@ class PythonTyper(
 
         self._assign(stmt.location, stmt.target, item_type)
         self.judge(stmt.target, item_type)
-        env: Environment = Environment(self.env)
+        body_returned: bool = self.process_block(stmt.body, inner_env)
-        body_returned: bool = self.process_block(stmt.body, env)
+
+        self.env = outer_env
+
         if body_returned:
             raise ReturnException()
 
+    def visit_import_stmt(self, stmt: p.ImportStmt) -> None:
+        self._visit_imports(stmt.location, stmt.imports)
+
+    def visit_from_import_stmt(self, stmt: p.FromImportStmt) -> None:
+        self._visit_imports(stmt.location, stmt.imports)
+
+    def _visit_imports(self, location: Location, imports: list[p.ImportAlias]) -> None:
+        for import_ in imports:
+            self._assign_var(
+                location,
+                p.VariableExpr(
+                    name=import_.imported_name,
+                    location=import_.location,
+                ),
+                UnknownType(),
+            )
+
     def visit_raw_stmt(self, stmt: p.RawStmt) -> None:
         pass
 
@@ -599,6 +708,11 @@ class PythonTyper(
             return UnknownType()
 
     def visit_unary_expr(self, expr: p.UnaryExpr) -> Type:
+        # Special case because there is no __not__ dunder method
+        match expr.operator:
+            case ast.Not():
+                return self.types.get_type("bool")
+
         method: Optional[str] = PY_UNARY_METHODS.get(expr.operator.__class__)
         if method is None:
             self.logger.warning(f"Unsupported operator {expr.operator}")
@@ -819,6 +933,8 @@ class PythonTyper(
                 return self._visit_frame_subscript(unfolded, expr)
             case FrameGroupBy():
                 return self._visit_frame_groupby_subscript(unfolded, expr)
+            case ColumnType():
+                return self._visit_column_subscript(unfolded, expr)
 
         operation: Optional[Type] = self.types.lookup_member(object, "__getitem__")
         if operation is None:
@@ -849,6 +965,15 @@ class PythonTyper(
         return UnknownType()
 
     def visit_base_type(self, node: p.BaseType) -> Type:
+        if node.base == "Column":
+            if len(node.args) != 1:
+                self.reporter.error(
+                    node.location,
+                    f"Column requires 1 type argument, {len(node.args)} provided",
+                )
+                return ColumnType(type=UnknownType())
+            return ColumnType(type=self.resolve_type_expr(node.args[0]))
+
         base: Type
         try:
             base = self.types.get_type(node.base)
@@ -887,6 +1012,15 @@ class PythonTyper(
         )
 
     def _get_iterator_type(self, expr: p.Expr, type: Type) -> Optional[Type]:
+        """Get the item type of an iterator type
+
+        Args:
+            expr (p.Expr): the iterator expression
+            type (Type): the iterator type
+
+        Returns:
+            Optional[Type]: the item type, or `None` if it cannot be determined
+        """
         # TODO: lookup __iter__
         getitem: Optional[Type] = self.types.lookup_member(type, "__getitem__")
         if getitem is None:
@@ -906,6 +1040,16 @@ class PythonTyper(
         return result.result
 
     def define_typevar(self, call: p.CallExpr) -> Optional[TypeVar]:
+        """Define a type variable from a call to `typing.TypeVar`
+
+        Args:
+            call (p.CallExpr): the call to `typing.TypeVar`
+
+        Returns:
+            Optional[TypeVar]: the define type variable, or `None` if the call
+                is invalid
+        """
+
         def is_kw_true(name: str) -> bool:
             match call.keywords.get(name):
                 case p.LiteralExpr(value=True):
@@ -948,6 +1092,19 @@ class PythonTyper(
                 return None
 
     def _parse_type_from_expr(self, expr: p.Expr) -> p.MidasType:
+        """Parse a type expression from a raw expression
+
+        This is useful for expressions inside a `TypeVar`'s `bound` parameter
+
+        Args:
+            expr (p.Expr): the expression to parse
+
+        Raises:
+            NotImplementedError: if the expression is not supported
+
+        Returns:
+            p.MidasType: the parsed type node
+        """
         location: Location = expr.location
         parser = PythonParser()
         match expr:
@@ -960,6 +1117,16 @@ class PythonTyper(
                 raise NotImplementedError
 
     def _get_literal(self, expr: p.Expr) -> tuple[bool, Any]:
+        """Get the literal value of a literal-like expression
+
+        Args:
+            expr (p.Expr): the expression
+
+        Returns:
+            tuple[bool, Any]: a tuple containing a boolean indicating whether
+                the given expression is literal-like, and the literal value (or
+                `None` if the first value is `False`)
+        """
         match expr:
             case p.LiteralExpr(value=value):
                 return True, value
@@ -1016,6 +1183,17 @@ class PythonTyper(
     def _evaluate_cast_statically(
         self, expr: p.CastExpr, subject_type: Type, target_type: Type, lit_value: Any
     ) -> bool:
+        """Evaluate the given cast expression statically
+
+        Args:
+            expr (p.CastExpr): the cast expression
+            subject_type (Type): the subject type being casted
+            target_type (Type): the target type to which the expression is casted
+            lit_value (Any): the literal value of the expression
+
+        Returns:
+            bool: whether the cast expression could be evaluated successfully
+        """
         match target_type:
             case TopType():
                 return True
@@ -1151,3 +1329,14 @@ class PythonTyper(
         return self.frame_mgr.groupby_get(
             self.reporter, expr.location, groupby, expr.index
         )
+
+    def _visit_column_subscript(
+        self, column: ColumnType, expr: p.SubscriptExpr
+    ) -> Type:
+        index_type: Type = self.type_of(expr.index)
+        return self.column_mgr.get(
+            self.reporter,
+            expr.location,
+            column,
+            (expr.index, index_type),
+        )

midas/checker/registry.py

@@ -8,11 +8,9 @@ from midas.checker.types import (
     AppliedType,
     BaseType,
     ColumnType,
-    ComplexType,
     ConstraintType,
     DataFrameType,
     DerivedType,
-    ExtensionType,
     Function,
     GenericType,
     OverloadedFunction,
@@ -29,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] = {}
@@ -81,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]
@@ -109,11 +130,29 @@ 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
@@ -161,14 +200,6 @@ class TypesRegistry:
             case (BaseType(name=name1), BaseType(name=name2)):
                 return self.is_builtin_subtype(name1, name2)
 
-            case (ComplexType(properties=props1), ComplexType(properties=props2)):
-                for k, t in props2.items():
-                    if k not in props1:
-                        return False
-                    if not self.is_subtype(props1[k], t):
-                        return False
-                return True
-
             case (DataFrameType(columns=columns1), DataFrameType(columns=columns2)):
                 # TODO: check order?
                 by_name1: dict[str, DataFrameType.Column] = {
@@ -218,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
@@ -334,6 +374,18 @@ class TypesRegistry:
         return True
 
     def apply_generic(self, type: Type, args: list[Type]) -> Type:
+        """Instantiate a generic type with the given type arguments
+
+        Args:
+            type (Type): the generic
+            args (list[Type]): the type arguments
+
+        Raises:
+            ValueError: if the arguments are invalid (wrong number, bound violation)
+
+        Returns:
+            Type: the applied generic type
+        """
         match type:
             case DerivedType(name=name, type=base):
                 return DerivedType(name=name, type=self.apply_generic(base, args))
@@ -399,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:
@@ -431,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)
 
@@ -459,18 +510,54 @@ 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_,
@@ -478,6 +565,15 @@ class TypesRegistry:
         )
 
     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_,

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,28 +28,28 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
         """Begin a new scope inside the current one"""
         self.scopes.append({})
 
-    def end_scope(self):
+    def end_scope(self) -> dict[str, bool]:
-        """Close the current scope"""
+        """Close and return the current scope"""
-        self.scopes.pop()
+        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(
+            self.reporter.error(
-                f"A variable with the name {name} is already declared in this scope"
+                location,
+                f"A variable with the name '{name}' is already declared in this scope",
             )
+        else:
             scope[name] = False
 
     def define(self, name: str) -> None:
@@ -77,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
@@ -94,7 +105,11 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
         """
         self.begin_scope()
         for param in function.params.all:
-            self.declare(param.name)
+            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)
+        self.declare(stmt.location, stmt.name)
-        # NOTE: resolve type here?
-        self.define(stmt.name)
 
     def visit_assign_stmt(self, stmt: p.AssignStmt) -> None:
         self.resolve(stmt.value)
@@ -121,8 +134,8 @@ 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):
+                if not self.is_declared(name):
-                    self.declare(name)
+                    self.declare(target.location, name)
                 self.define(name)
                 target.accept(self)
 
@@ -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(
+            self.reporter.error(
-                f"Cannot use local variable '{expr.name}' in its own initializer"
+                expr.location,
+                f"Variable '{expr.name}' is declared but may not be defined",
             )  # aka. UnboundLocalError
         self.resolve_local(expr, expr.name)
 

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:
@@ -24,6 +28,8 @@ class BaseType:
 
 @dataclass(frozen=True, kw_only=True)
 class DerivedType:
+    """A derived type, i.e. a named subtype of another type"""
+
     name: str
     type: Type
 
@@ -33,18 +39,24 @@ class DerivedType:
 
 @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:
+    """A function type"""
+
     params: ParamSpec
     returns: Type
 
@@ -57,14 +69,20 @@ class Function:
         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)
@@ -87,31 +105,17 @@ class ParamSpec:
 
 @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"
@@ -119,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
@@ -136,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
@@ -146,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
@@ -156,6 +166,8 @@ class AppliedType:
 
 @dataclass(frozen=True, kw_only=True)
 class ConstraintType:
+    """A type with a constraint expression"""
+
     type: Type
     constraint: m.Expr
 
@@ -166,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:
@@ -174,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:
@@ -182,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:
@@ -197,6 +215,8 @@ class DataFrameType:
 
 @dataclass(frozen=True, kw_only=True)
 class FrameGroupBy:
+    """A frame group-by object"""
+
     frame: DataFrameType
 
     def __str__(self) -> str:
@@ -205,6 +225,8 @@ class FrameGroupBy:
 
 @dataclass(frozen=True, kw_only=True)
 class ColumnGroupBy:
+    """A column group-by object"""
+
     column: ColumnType
 
     def __str__(self) -> str:
@@ -212,6 +234,19 @@ class ColumnGroupBy:
 
 
 def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
+    """Substitute type variables in the given type
+
+    This function is called recursively on inner type structures
+
+    Args:
+        type (Type): the type in which to substitute type variables
+        substitutions (dict[str, Type]): a mapping of type variable names to
+            concrete types
+
+    Returns:
+        Type: the resulting type with substitutions applied
+    """
+
     def sub_parameter(param: Function.Parameter):
         return Function.Parameter(
             pos=param.pos,
@@ -266,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,
@@ -354,6 +371,14 @@ 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 DerivedType(type=ref_type):
             return unfold_type(ref_type)
@@ -362,6 +387,15 @@ def unfold_type(type: Type) -> Type:
 
 
 def to_annotation(type: Type) -> str:
+    """Convert the given type to a Python annotation string
+
+    Args:
+        type (Type): the type to convert
+
+    Returns:
+        str: the annotation string
+    """
+
     def _params_annotation(spec: ParamSpec) -> str:
         if len(spec.kw) != 0:
             return "..."
@@ -394,9 +428,6 @@ def to_annotation(type: Type) -> str:
         case OverloadedFunction():
             return "Callable"
 
-        case ComplexType() | ExtensionType():
-            raise NotImplementedError
-
         case TypeVar(name=name):
             return name
 
@@ -430,6 +461,8 @@ def to_annotation(type: Type) -> str:
 
 @dataclass(frozen=True, kw_only=True)
 class Predicate:
+    """A predicate"""
+
     type: Type
     body: m.Expr
     alias: bool
@@ -443,8 +476,6 @@ Type = (
     | UnitType
     | Function
     | OverloadedFunction
-    | ComplexType
-    | ExtensionType
     | TypeVar
     | GenericType
     | AppliedType

midas/checker/unifier.py

@@ -19,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")
@@ -29,6 +37,16 @@ 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(
             params=ParamSpec(
                 pos=[
@@ -60,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)
@@ -81,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):
@@ -150,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():
@@ -164,6 +222,15 @@ class Unifier:
     def map_params(
         self, func1: Function, func2: Function
     ) -> list[tuple[Function.Parameter, Function.Parameter]]:
+        """Map parameters of two functions
+
+        Args:
+            func1 (Function): the first function
+            func2 (Function): the second function
+
+        Returns:
+            list[tuple[Function.Parameter, Function.Parameter]]: the list of parameter pairs
+        """
         pos1: list[Function.Parameter] = func1.params.pos
         mixed1: list[Function.Parameter] = func1.params.mixed
         kw1: list[Function.Parameter] = func1.params.kw

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,6 +107,22 @@ 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 = []
 
@@ -109,10 +161,10 @@ class VarianceInferrer:
                     Variance.COVARIANT: 1,
                     Variance.CONTRAVARIANT: -1,
                 }
-                for param, param in zip(args, params):
+                for arg, param in zip(args, params):
                     param_polarity: Polarity = polarities[param.variance]
                     self.walk(
-                        param,
+                        arg,
                         cast(Polarity, polarity * param_polarity),
                         base_name,
                         path + [f"applied:'{name}'"],

midas/cli/highlighter.py

@@ -338,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,
+    &.constraint-type {
-    &.complex-type {
         --col: 150, 150, 150;
     }
 

midas/generator/collector.py

@@ -5,20 +5,45 @@ 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] = {}
@@ -30,6 +55,15 @@ class AssertionCollector:
         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,
@@ -40,20 +74,51 @@ class AssertionCollector:
         )
 
     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

@@ -18,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,
 }
 
@@ -40,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
@@ -47,9 +49,22 @@ 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)
@@ -75,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}"
@@ -85,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(
@@ -96,6 +128,14 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
         return self.make_func(name, [ast.Return(value=body)], predicate.type)
 
     def make_args(self, params: ParamSpec) -> ast.arguments:
+        """Translate the given parameter spec into an `ast.arguments` node
+
+        Args:
+            params (ParamSpec): the parameter spec to translate
+
+        Returns:
+            ast.arguments: the equivalent `ast.arguments`
+        """
         return ast.arguments(
             posonlyargs=[
                 ast.arg(
@@ -125,6 +165,33 @@ 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(params=params, returns=Function()):
                 inner_name: str = f"inner{level}"
@@ -152,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

@@ -16,11 +16,9 @@ from midas.checker.types import (
     BaseType,
     ColumnGroupBy,
     ColumnType,
-    ComplexType,
     ConstraintType,
     DataFrameType,
     DerivedType,
-    ExtensionType,
     FrameGroupBy,
     Function,
     GenericType,
@@ -40,11 +38,23 @@ 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__"
 
@@ -72,12 +82,28 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         self.define_is_column: bool = False
 
     def set_src_path(self, path: Path):
+        """Set the current source file path
+
+        Args:
+            path (Path): the new source file path
+        """
         self.rel_src_path = path.resolve().relative_to(self.workdir)
 
     def generate_ast(self, typed_ast: TypedAST) -> ast.AST:
+        """Translate the given type checked AST into a Python `ast.AST`
+
+        Args:
+            typed_ast (TypedAST): the type checked Python AST
+
+        Returns:
+            ast.AST: the generated raw AST
+        """
         self._typed_ast = typed_ast
         body: list[ast.stmt] = self._visit_body(typed_ast.stmts, can_be_empty=True)
         predicates: list[ast.stmt] = self._constraint_generator.get_definitions()
+        assertion_definitions: list[ast.stmt] = list(
+            typed_ast.assertions.definitions.values()
+        )
 
         body = predicates + body
 
@@ -87,6 +113,8 @@ 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
@@ -98,6 +126,29 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         out_path: Optional[Path] = None,
         type_files: Optional[list[tuple[Path, Optional[str]]]] = None,
     ) -> Path:
+        """Generate all project files for the given source file and AST
+
+        This function calls :func:`generate_ast` to generate the output AST,
+        unparses it to runnable Python code, and also generates stubs for
+        user-defined Midas types in the same output directory
+
+        Args:
+            typed_ast (TypedAST): the type-checked AST
+            src_path (Path): the source file path
+            out_path (Optional[Path], optional): the output file path. If `None`,
+                the relative path of the source file to the working directory is
+                used to compute an equivalent path in the build directory.
+                Defaults to None.
+            type_files (Optional[list[tuple[Path, Optional[str]]]], optional):
+                the list of Midas files used to type check the AST. Defaults to None.
+
+        Raises:
+            ValueError: if `out_path` is `None` and the computed path is outside
+                the build directory
+
+        Returns:
+            Path: the actual `out_path` used
+        """
         self.set_src_path(src_path)
         if out_path is None:
             if self.build_dir.exists():
@@ -126,6 +177,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         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)
@@ -135,6 +192,18 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         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
@@ -251,6 +320,14 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.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],
@@ -307,12 +384,56 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
             orelse=[],
         )
 
+    def visit_import_stmt(self, stmt: p.ImportStmt) -> ast.stmt:
+        return ast.Import(
+            names=self._convert_imports(stmt.imports),
+        )
+
+    def visit_from_import_stmt(self, stmt: p.FromImportStmt) -> ast.stmt:
+        return ast.ImportFrom(
+            module=stmt.module,
+            names=self._convert_imports(stmt.imports),
+            level=stmt.level,
+        )
+
+    def _convert_imports(self, imports: list[p.ImportAlias]) -> list[ast.alias]:
+        """Translate a list of import aliases
+
+        Args:
+            imports (list[p.ImportAlias]): the import aliases to translate
+
+        Returns:
+            list[ast.alias]: the translated aliases
+        """
+        return [
+            ast.alias(
+                name=import_.name,
+                asname=import_.alias,
+            )
+            for import_ in imports
+        ]
+
     def visit_raw_stmt(self, stmt: p.RawStmt) -> ast.stmt:
         return stmt.stmt
 
     def _visit_body(
         self, stmts: list[p.Stmt], can_be_empty: bool = False
     ) -> list[ast.stmt]:
+        """Translate a list of statements
+
+        Assertions generated while translating a statement are inserted before it,
+        and aliases are deleted after the statement they're used in.
+
+        Extraneous `pass` statements are automatically removed
+
+        Args:
+            stmts (list[p.Stmt]): the statements to translate
+            can_be_empty (bool, optional): if `False` and no statement is
+            generated, an `ast.Pass` statement is returned. Defaults to False.
+
+        Returns:
+            list[ast.stmt]: the generated statements
+        """
         generated: list[ast.stmt] = []
         for stmt in stmts:
             scope = Scope()
@@ -335,6 +456,20 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         return generated
 
     def _make_alias(self, node: p.Expr, expr: ast.expr) -> ast.expr:
+        """Generate a unique alias for the given expression
+
+        This function creates a unique name, generates an assignment statement
+        to define the alias before the current statement, adds the alias to the
+        list of aliases defined in the current statement, and returns an
+        expression that can be used in place of `expr`
+
+        Args:
+            node (p.Expr): the AST node that generated `expr`
+            expr (ast.expr): the expression to alias
+
+        Returns:
+            ast.expr: the generated alias reference
+        """
         name: str = f"__midas_a{self._alias_count}__"
         alias = ast.Name(id=name)
         self._alias_count += 1
@@ -349,6 +484,15 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         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(
@@ -357,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
@@ -368,6 +528,17 @@ 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() | TopType():
                 return []
@@ -502,8 +673,6 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
             case (
                 Function()
                 | OverloadedFunction()
-                | ComplexType()
-                | ExtensionType()
                 | GenericType()
                 | FrameGroupBy()
                 | ColumnGroupBy()
@@ -522,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(
@@ -545,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(
@@ -552,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}"
@@ -567,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
@@ -576,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(
@@ -612,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(
@@ -650,6 +882,18 @@ 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
     ) -> 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)
@@ -663,6 +907,14 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         )
 
     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:
@@ -678,6 +930,15 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         )
 
     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)

midas/generator/stubs.py

@@ -8,11 +8,9 @@ from midas.checker.types import (
     BaseType,
     ColumnGroupBy,
     ColumnType,
-    ComplexType,
     ConstraintType,
     DataFrameType,
     DerivedType,
-    ExtensionType,
     FrameGroupBy,
     Function,
     GenericType,
@@ -32,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] = []
@@ -43,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
@@ -92,6 +97,12 @@ 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
@@ -129,6 +140,17 @@ 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 DerivedType(type=base):
                 return [self.dump_type(base)], {}
@@ -173,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)]
 
@@ -194,6 +226,14 @@ 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 DerivedType(name=name) | GenericType(name=name) if (
                 name in self.substitutions
@@ -227,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)
 
@@ -273,14 +305,11 @@ class StubsGenerator:
                     ),
                 )
 
-            case ColumnType(type=inner):
+            case ColumnType():
                 self.import_pandas = True
-                return ast.Subscript(
+                return ast.Attribute(
-                    value=ast.Attribute(
                     value=ast.Name(id="pd"),
                     attr="Series",
-                    ),
-                    slice=self.dump_type(inner),
                 )
 
             case DataFrameType():
@@ -322,6 +351,17 @@ class StubsGenerator:
     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:
@@ -350,6 +390,16 @@ class StubsGenerator:
                 ]
 
     def dump_params(self, params: ParamSpec, with_self: bool = False) -> ast.arguments:
+        """Generate an `ast.arguments` node for the given parameter spec
+
+        Args:
+            params (ParamSpec): the parameter spec to translate
+            with_self (bool, optional): whether to include a `self` parameter.
+                Defaults to False.
+
+        Returns:
+            ast.arguments: the generate Python AST node
+        """
         pos: list[ast.arg] = [
             ast.arg(
                 arg=f"_{param.pos}",
@@ -392,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(
@@ -413,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 ".":
@@ -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

@@ -36,6 +36,7 @@ class TokenType(Enum):
     EQUAL = auto()
     EQUAL_EQUAL = auto()
     BANG_EQUAL = auto()
+    BANG = auto()
 
     # Literals
     IDENTIFIER = auto()
@@ -104,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

@@ -5,11 +5,9 @@ from midas.ast.midas import (
     AliasStmt,
     BinaryExpr,
     CallExpr,
-    ComplexType,
     ConstraintType,
     Expr,
     ExtendStmt,
-    ExtensionType,
     FrameType,
     FunctionType,
     GenericType,
@@ -35,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,
@@ -73,10 +72,10 @@ 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):
@@ -95,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.
+        A type declaration creates a named subtype of a type expression.
-        In either case, it can have an optional template expression after its name, wrapped in brackets.
+        It can have an optional template expression after its name, wrapped in brackets, to handle type parameters.
-        A simple type alias is derived from a base type expression, and can have a optional constraint expression preceded by the `where` keyword.
-        A full simple type alias is thus written:
-        ```
-        type Name[Template](TypeExpr) where Condition
-        ```
 
-        A new complex type has a set of properties which are named, have a type and an optional constraint expression (also preceded by the `where` keyword).
+        A type statement consists of:
-        A full complex type definition is thus written:
+        - the `type` keyword
-        ```
+        - a name (identifier)
-        type Name[Template] {
+        - (optional) type parameters
-            prop1: TypeExpr1 where Condition1
+        - a body, a type expression (see :func:`type_expr`)
-            prop2: TypeExpr2 where Condition2
-            ...
-        }
-        ```
 
         Returns:
             TypeStmt: the parsed type declaration statement
@@ -165,11 +155,16 @@ class MidasParser(Parser):
     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 type name")
+        name: Token = self.consume_identifier("Expected alias name")
 
         self.consume(TokenType.EQUAL, "Expected '=' before alias definition")
 
@@ -184,27 +179,26 @@ class MidasParser(Parser):
     def type_expr(self) -> Type:
         """Parse a type expression
 
-        A type is an identifier, optionally followed by a template expression.
+        A type expression can either be a function type (see :func:`function`)
-        It can also optionally be followed by a '?' to indicate a nullable type
+        or a constraint type (see :func:`constraint_type`)
 
         Returns:
             TypeExpr: the parsed type expression
         """
-        base: Type
         if self.match(TokenType.FUNC):
-            base = self.function()
+            return self.function()
-        else:
+        return self.constraint_type()
-            base = self.constraint_type()
-        if self.match(TokenType.AND):
-            extension: ComplexType = self.complex_type()
-            return ExtensionType(
-                location=Location.span(base.location, extension.location),
-                base=base,
-                extension=extension,
-            )
-        return base
 
     def constraint_type(self) -> Type:
+        """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()
@@ -216,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":
@@ -240,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):
@@ -250,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:
+    def frame_type(self) -> FrameType:
-        """Parse a type definition body
+        """Parse a frame type expression
 
-        A type definition body is a set of whitespace-separated
+        A frame type consists of:
-        property statements enclosed in curly braces
+        - 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")
 
@@ -311,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
@@ -321,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
@@ -340,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
@@ -356,7 +379,10 @@ 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
@@ -377,6 +403,14 @@ class MidasParser(Parser):
         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()
@@ -388,6 +422,14 @@ class MidasParser(Parser):
         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()
@@ -399,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)
@@ -412,12 +457,44 @@ 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
@@ -437,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(
+                        error_msg = (
-                            self.peek(),
+                            "Cannot pass positional arguments after a keyword argument"
-                            "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):
@@ -460,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
@@ -475,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
@@ -501,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
@@ -524,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
@@ -551,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
@@ -576,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
@@ -599,6 +732,17 @@ class MidasParser(Parser):
         )
 
     def function(self) -> FunctionType:
+        """Parse a function type expression
+
+        A function consists of:
+        - the `fn` keyword
+        - a parameter spec (see :func:`function_params`)
+        - the arrow keyword `->`
+        - a result type expression (see :func:`type_expr`)
+
+        Returns:
+            FunctionType: the parsed function type expression
+        """
         params: ParamSpec = self.function_params()
 
         self.consume(TokenType.ARROW, "Expected '->' before result type")
@@ -611,6 +755,21 @@ class MidasParser(Parser):
         )
 
     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"
         )

midas/parser/python.py

@@ -16,9 +16,12 @@ from midas.ast.python import (
     ForStmt,
     FrameColumn,
     FrameType,
+    FromImportStmt,
     Function,
     GetExpr,
     IfStmt,
+    ImportAlias,
+    ImportStmt,
     ListExpr,
     LiteralExpr,
     LogicalExpr,
@@ -50,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"
 
@@ -103,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)
@@ -351,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 _:
@@ -362,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,

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

syntax/annotations.ebnf

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

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 = ```
+#let extra-css = ```css
-{[`value` <
+svg.railroad .terminal rect {
-  [`number` 'digit' * ! <!, ["." 'digit' * !]>],
+  fill: #F7DCD4;
-  [`boolean` <"False", "True">],
+}
-  [`none` "None"]
+```
+#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,
+  type-args: type-args,
-  constraint,
+  frame-schema: frame-schema,
-  type-with-constraints,
+  type: type,
-  column-def,
+)
-  frame-def,
+
-  annotation,
+#let inline = (
+  "type-args",
+  "frame-schema",
 )
 
 #set text(font: "Source Sans 3")
 
-= Type annotation syntax
+#title[Supported Python annotation syntax]
 
-#for rule in rules {
+= Outline
-  render(rule)
-}
 
-/*
+#box(
-#let by-name = (
+  columns(
-  annotation: annotation,
+    2,
-  frame-def: frame-def,
+    outline(title: none),
-  column-def: column-def,
+  ),
-  type-with-constraints: type-with-constraints,
+  height: 9cm,
-  constraint: constraint,
+  stroke: 1pt,
-  value: value,
+  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() {
+  for name in inline {
-    new-rule = new-rule.replace("'" + key + "'", rule.text.slice(1, -1))
+    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 ")"
+PosArg ::= Expression
-Primary ::= "_" | Value | Identifier | Grouping
+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
+CallArgs ::= "(" Args ")" 
-Comparison ::= Unary (ComparisonOp Unary)*
+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)*)? "]"
+TypeArgs ::= "[" (Type ("," Type)*)? "]"
-GenericType ::= NamedType TypeParams?
+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
+MemberStatement ::= ("prop" | "def") Identifier ":" Type
-ExtendBody ::= "{" OpDefinition* "}"
+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 = ```
+#let literal = ```
-{[`value` <
+{[`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 = ```
+#let type-args = ```
-{[`type-params` "[" <!, 'type'*","> "]"]}
+{[`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 = ```
+#let member-stmt = ```
-{[`op-definition` "op" 'identifier' "(" 'type' ")" "->" 'type']}
+{[`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",
+  "call-args",
-  "type-params",
+  "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

@@ -1,7 +1,8 @@
+import sys
 from typing import Type
 
 from midas.cli.ansi import Ansi
-from tests.base import Tester
+from tests.base import Tester, TestsSummary
 from tests.checker import CheckerTester
 from tests.generator import GeneratorTester
 from tests.midas import MidasTester
@@ -15,12 +16,12 @@ def print_banner(name: str):
     print(horizontal)
 
 
-def run_tests(tester_cls: Type[Tester]) -> bool:
+def run_tests(tester_cls: Type[Tester]) -> TestsSummary:
     print_banner(tester_cls.__name__)
     tester: Tester = tester_cls()
-    success: bool = tester.run_all_tests()
+    summary: TestsSummary = tester.run_all_tests()
     print()
-    return success
+    return summary
 
 
 def main():
@@ -31,12 +32,17 @@ def main():
         GeneratorTester,
     ]
 
-    success: bool = all(map(run_tests, testers))
+    summaries: list[TestsSummary] = list(
+        map(run_tests, testers)
+    )  # list to avoid early stop
+    summary: TestsSummary = TestsSummary.concat(*summaries)
 
-    if success:
+    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__":

tests/base.py

@@ -4,6 +4,7 @@ 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
 
@@ -14,6 +15,43 @@ 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"""
 
@@ -34,40 +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):
             path: Path = test.resolve().relative_to(self.CASES_DIR)
             print(f"{Ansi.FG(Ansi.BRIGHT_CYAN)}Case {i+1}/{n}: {path}{Ansi.RESET}")
-            print(Ansi.DIM, end="")
             success: bool = self._run_test(test)
-            print(Ansi.RESET, end="")
+            summary.add(str(path), success)
-            if success:
-                successes += 1
-            else:
-                failures += 1
 
         print(rule)
-        print(f"Success: {successes}/{n}")
+        print(f"Success: {summary.successes}/{n}")
-        print(f"Failed: {failures}/{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()
 
@@ -143,16 +178,18 @@ 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)
+                    summary = tester.run_tests(args.FILE)
-                if not success:
+                if not summary.success:
+                    summary.print()
                     sys.exit(1)
             case None:
-                success: bool = tester.run_all_tests()
+                summary: TestsSummary = tester.run_all_tests()
-                if not success:
+                if not summary.success:
+                    summary.print()
                     sys.exit(1)
             case _:
                 print(f"Unknown subcommand '{args.subcommand}'")

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

@@ -285,7 +285,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -295,7 +296,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -305,7 +307,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -361,7 +364,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -371,7 +375,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -381,7 +386,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -455,7 +461,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -465,7 +472,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -475,7 +483,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -553,7 +562,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -563,7 +573,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -573,7 +584,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -665,7 +677,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -675,7 +688,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -685,7 +699,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -780,7 +795,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -790,7 +806,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -800,7 +817,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -870,7 +888,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -880,7 +899,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -890,7 +910,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -951,7 +972,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -961,7 +983,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -971,7 +994,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -1058,7 +1082,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -1068,7 +1093,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -1078,7 +1104,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -1176,7 +1203,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -1186,7 +1214,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -1196,7 +1225,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -1294,7 +1324,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -1304,7 +1335,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -1314,7 +1346,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },
@@ -1412,7 +1445,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "mixed": [
@@ -1422,7 +1456,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": [
@@ -1432,7 +1467,8 @@
               "type": {
                 "name": "str"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ]
         },

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

@@ -145,7 +145,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             },
             {
               "pos": 1,
@@ -153,7 +154,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": []

tests/cases/checker/07_variance.py

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

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

@@ -126,7 +126,8 @@
               "pos": 0,
               "name": "object",
               "type": {},
-              "required": false
+              "required": false,
+              "unsupported": false
             }
           ],
           "mixed": [],
@@ -388,7 +389,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": []
@@ -459,7 +461,8 @@
                           "bound": null,
                           "variance": "INVARIANT"
                         },
-                        "required": true
+                        "required": true,
+                        "unsupported": false
                       }
                     ],
                     "mixed": [],
@@ -471,7 +474,8 @@
                     "variance": "INVARIANT"
                   }
                 },
-                "required": true
+                "required": true,
+                "unsupported": false
               },
               {
                 "pos": 1,
@@ -489,7 +493,8 @@
                     "name": "list"
                   }
                 },
-                "required": true
+                "required": true,
+                "unsupported": false
               }
             ],
             "mixed": [],
@@ -565,7 +570,8 @@
               "type": {
                 "name": "float"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": []
@@ -636,7 +642,8 @@
                           "bound": null,
                           "variance": "INVARIANT"
                         },
-                        "required": true
+                        "required": true,
+                        "unsupported": false
                       }
                     ],
                     "mixed": [],
@@ -648,7 +655,8 @@
                     "variance": "INVARIANT"
                   }
                 },
-                "required": true
+                "required": true,
+                "unsupported": false
               },
               {
                 "pos": 1,
@@ -666,7 +674,8 @@
                     "name": "list"
                   }
                 },
-                "required": true
+                "required": true,
+                "unsupported": false
               }
             ],
             "mixed": [],
@@ -732,7 +741,8 @@
               "type": {
                 "name": "int"
               },
-              "required": true
+              "required": true,
+              "unsupported": false
             }
           ],
           "kw": []
@@ -803,7 +813,8 @@
                           "bound": null,
                           "variance": "INVARIANT"
                         },
-                        "required": true
+                        "required": true,
+                        "unsupported": false
                       }
                     ],
                     "mixed": [],
@@ -815,7 +826,8 @@
                     "variance": "INVARIANT"
                   }
                 },
-                "required": true
+                "required": true,
+                "unsupported": false
               },
               {
                 "pos": 1,
@@ -833,7 +845,8 @@
                     "name": "list"
                   }
                 },
-                "required": true
+                "required": true,
+                "unsupported": false
               }
             ],
             "mixed": [],

tests/cases/checker/09_frame_ops.py

@@ -1,8 +1,9 @@
 # type: ignore
 # ruff: disable [F821]
+import pandas as pd
 
-df1: Frame[a:int, b:float]
+df1 = cast(Frame[i:int, a:int, b:float], pd.DataFrame())
-df2: Frame[a:int, b:float]
+df2 = cast(Frame[i:int, a:int, b:float], pd.DataFrame())
 
 _: Any
 
@@ -38,7 +39,7 @@ _ = df1.sum()
 _ = df1.var()
 
 # Groupby
-df_gb = df1.groupby(by="a")
+df_gb = df1.groupby(by="i")
 
 _ = df_gb.kurt()
 _ = df_gb.max()

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=100),
-                Constant(value=200.0)],
+                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=0),
-          Constant(value=18.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
+type GeoLocationDifference = object
-extend GeoLocation {
-    // This type is compatible with the `-` operation with another GeoLocation
-    // i.e. you can subtract a GeoLocation from another GeoLocation, resulting
-    // in a Difference of GeoLocations
-    def __sub__: fn(GeoLocation, /) -> Difference[GeoLocation]
-}
 
-// For complex generics, you need to specify how the genericity the properties
+extend GeoLocationDifference {
-// are handled
-type Difference[GeoLocation] = {
     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",
+      "_type": "FromImportStmt",
-      "stmt": "from __future__ import annotations"
+      "module": "__future__",
+      "imports": [
+        {
+          "_type": "ImportAlias",
+          "name": "annotations",
+          "alias": null
+        }
+      ],
+      "level": 0
     },
     {
       "_type": "TypeAssign",

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

@@ -1,8 +1,16 @@
 {
   "stmts": [
     {
-      "_type": "RawStmt",
+      "_type": "FromImportStmt",
-      "stmt": "from __future__ import annotations"
+      "module": "__future__",
+      "imports": [
+        {
+          "_type": "ImportAlias",
+          "name": "annotations",
+          "alias": null
+        }
+      ],
+      "level": 0
     },
     {
       "_type": "TypeAssign",

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

@@ -1,8 +1,16 @@
 {
   "stmts": [
     {
-      "_type": "RawStmt",
+      "_type": "FromImportStmt",
-      "stmt": "from __future__ import annotations"
+      "module": "__future__",
+      "imports": [
+        {
+          "_type": "ImportAlias",
+          "name": "annotations",
+          "alias": null
+        }
+      ],
+      "level": 0
     },
     {
       "_type": "Function",

tests/serializer/midas.py

@@ -4,11 +4,9 @@ from midas.ast.midas import (
     AliasStmt,
     BinaryExpr,
     CallExpr,
-    ComplexType,
     ConstraintType,
     Expr,
     ExtendStmt,
-    ExtensionType,
     FrameType,
     FunctionType,
     GenericType,
@@ -172,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",
@@ -200,13 +192,6 @@ class MidasAstJsonSerializer(
             "required": param.required,
         }
 
-    def visit_extension_type(self, type: ExtensionType) -> dict:
-        return {
-            "_type": "ExtensionType",
-            "base": type.base.accept(self),
-            "extension": type.extension.accept(self),
-        }
-
     def visit_frame_type(self, type: FrameType) -> dict:
         return {
             "_type": "FrameType",

tests/serializer/python.py

@@ -15,9 +15,12 @@ from midas.ast.python import (
     ForStmt,
     FrameColumn,
     FrameType,
+    FromImportStmt,
     Function,
     GetExpr,
     IfStmt,
+    ImportAlias,
+    ImportStmt,
     ListExpr,
     LiteralExpr,
     LogicalExpr,
@@ -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",

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",
+                "path": "./syntaxes/midas.tmLanguage.json"
-                "embeddedLanguages": {
-                    "meta.embedded.block.python": "python"
-                }
             }
         ]
     }

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": {
+        "comments": {
-            "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"
-                }
-            },
             "patterns": [
-                { "include": "#type-base" },
+                {
-                { "include": "#type-body" }
+                    "begin": "//",
-            ]
-        },
-        "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": "$",
                     "beginCaptures": {
                         "0": {
-                    "name": "keyword.control.where.midas"
+                            "name": "punctuation.definition.comment.midas"
-                }
                         }
                     },
-        "type-body": {
+                    "end": "$",
-            "begin": "\\{",
+                    "name": "comment.line.double-slash.midas"
-            "end": "\\}",
+                },
+                {
+                    "begin": "/\\*",
                     "beginCaptures": {
                         "0": {
-                    "name": "punctuation.definition.type-body.begin.midas"
+                            "name": "punctuation.definition.comment.midas"
                         }
                     },
+                    "end": "\\*/",
                     "endCaptures": {
                         "0": {
-                    "name": "punctuation.definition.type-body.end.midas"
+                            "name": "punctuation.definition.comment.midas"
+                        }
+                    },
+                    "name": "comment.block.midas"
+                }
+            ]
+        },
+        "string": {
+            "begin": "\"",
+            "beginCaptures": {
+                "0": {
+                    "name": "punctuation.definition.string.begin.midas"
+                }
+            },
+            "end": "\"",
+            "endCaptures": {
+                "0": {
+                    "name": "punctuation.definition.string.end.midas"
+                }
+            },
+            "name": "string.quoted.double.c"
+        },
+        "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": "#type-prop"},
+                        {
-                {"include": "#comment"}
+                            "include": "#comments"
+                        },
+                        {
+                            "include": "#type-params"
+                        }
                     ]
                 },
-        "type-prop": {
+                {
-            "match": "([a-zA-Z_][a-zA-Z_\\d]*)(:)\\s*([a-zA-Z_][a-zA-Z_\\d]*)",
+                    "begin": "=",
-            "captures": {
+                    "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": "variable.name"
+                    "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": "\\{",
+            "beginCaptures": {
+                "0": {
+                    "name": "punctuation.section.block.begin.midas"
+                }
+            },
+            "end": "\\}",
+            "endCaptures": {
+                "0": {
+                    "name": "punctuation.section.block.end.midas"
+                }
+            },
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "include": "#member-stmt"
+                }
+            ]
+        },
+        "member-stmt": {
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "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"
+                                }
+                            ]
+                        }
+                    ]
+                }
+            ]
+        },
+        "predicate-stmt": {
+            "begin": "\\b(predicate)\\b",
+            "beginCaptures": {
+                "1": {
+                    "name": "keyword.declaration.midas"
+                }
+            },
+            "end": "$",
+            "name": "meta.declaration.predicate.midas",
+            "patterns": [
+                {
+                    "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"
+                        }
+                    ]
+                }
+            ]
+        },
+        "predicate-params": {
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "begin": "([A-Za-z_][A-Za-z0-9_]*)(:)",
+                    "beginCaptures": {
+                        "1": {
+                            "name": "variable.parameter.midas"
                         },
                         "2": {
                             "name": "punctuation.separator.annotation.midas"
-                },
-                "3": {
-                    "name": "meta.type.name"
-                }
                         }
                     },
-        "extend-def": {
+                    "end": "(?=,)|(?=\\))",
-            "begin": "\\b(extend)\\s*([a-zA-Z_][a-zA-Z_\\d]*)\\s+(\\{)",
+                    "patterns": [
-            "end": "\\}",
+                        {
+                            "include": "#type-expr"
+                        }
+                    ]
+                },
+                {
+                    "match": ",",
+                    "name": "punctuation.separator.midas"
+                }
+            ]
+        },
+        "type-expr": {
+            "patterns": [
+                {
+                    "include": "#comments"
+                },
+                {
+                    "begin": "\\b(fn)\\s*(\\()",
                     "beginCaptures": {
                         "1": {
-                    "name": "keyword.control.extend.midas"
+                            "name": "keyword.other.midas"
                         },
                         "2": {
-                    "name": "variable.name"
+                            "name": "punctuation.section.group.begin.midas"
-                },
-                "3": {
-                    "name": "punctuation.definition.extend-body.begin.midas"
                         }
                     },
+                    "end": "\\)",
                     "endCaptures": {
                         "0": {
-                    "name": "punctuation.definition.extend-body.end.midas"
+                            "name": "punctuation.section.group.end.midas"
                         }
                     },
                     "patterns": [
-                {"include": "#op-def"},
+                        {
-                {"include": "#comment"}
+                            "include": "#fn-params"
+                        }
                     ]
                 },
-        "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]*)",
+                    "match": "->",
-            "captures": {
+                    "name": "keyword.operator.arrow.midas"
-                "1": {
-                    "name": "keyword.control.op.midas"
                 },
-                "2": {
+                {
-                    "name" : "keyword.operator"
+                    "begin": "\\b(where)\\b",
-                },
-                "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"
+                            "name": "keyword.other.midas"
+                        }
                     },
-                "2": {
+                    "end": "$",
-                    "name": "variable.name"
+                    "patterns": [
+                        {
+                            "include": "#constraint"
+                        }
+                    ]
+                },
+                {
+                    "begin": "(\\bFrame\\b)(\\s*)(\\[)",
+                    "beginCaptures": {
+                        "1": {
+                            "name": "entity.name.type.midas"
                         },
                         "3": {
-                    "name": "variable.name"
+                            "name": "punctuation.section.brackets.begin.midas"
+                        }
                     },
-                "4": {
+                    "end": "\\]",
-                    "name": "variable.name"
+                    "endCaptures": {
-                },
+                        "0": {
-                "5": {
+                            "name": "punctuation.section.brackets.end.midas"
-                    "name": "keyword.operator.assignment"
                         }
                     },
                     "patterns": [
-                { "include": "source.python" }
+                        {
+                            "include": "#frame-schema"
+                        }
                     ]
                 },
-        "statement": {
+                {
+                    "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": "#comment" },
+                        {
-                { "include": "#type-def" },
+                            "include": "#comments"
-                { "include": "#extend-def" },
+                        },
-                { "include": "#pred-def" }
+                        {
+                            "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"
+                }
             ]
         }
     }