Merge pull request 'Various features for example pipeline' (#31) from feat/complete-example into main

Reviewed-on: #31

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: object,
+    temperature: Temperature,
+    pressure: Pressure,
+    humidity: Humidity,
+]
+
+alias DataWithHI = Frame[
+    station_id: StationID,
+    timestamp: object,
+    temperature: Temperature,
+    pressure: Pressure,
+    humidity: Humidity,
+    heat_index: HeatIndex,
+]
+
+alias DailyAverages = Frame[
+    timestamp: object,
+    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, 1.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=[
+                df["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/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/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.
-        The function doesn't take the raw expression as a parameter to accommodate
-        for desugared calls such as for operators.
+        Argument types are matched with the defined parameters.
+        This function doesn't take the raw expression as a parameter to
+        accommodate for desugared calls such as for operators.
 
         Args:
             location (Location): the call location
             callee (Type): the called function
-            positional (list[TypedExpr]): the list 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
-            the call is invalid or no overload matched the arguments uniquely
+            CallResult: the call result, either a type or an error
         """
         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
-            determined unambiguously, or `None`.
+            Union[tuple[Function, None], tuple[None, str]]: a tuple containing
+                the resolved function signature if it can be determined
+                unambiguously, or `None`, and an error message, or `None`
         """
         candidates: list[OverloadCandidate] = []
         errors: list[CallError] = []
@@ -345,7 +391,7 @@ class CallDispatcher(Generic[E]):
 
         Returns:
             tuple[bool, list[MappedArgument]]: a boolean reporting whether
-            the call is valid and the list of mapped arguments
+                the call is valid and the list of mapped arguments
         """
         set_params: set[str] = set()
 
@@ -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
-        a more specific function/ a subtype of another.
+        This is used to check whether a given overload is a more specific
+        function / a subtype of another.
 
         Args:
             mapped1 (list[MappedArgument]): the first argument mappings (subtype)

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=(
-                call.groupby.column
-                if preserve_inner_type
-                else ColumnType(type=TopType())
-            ),
+            returns=returns,
         )
 
         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())
-        if len(call.positional) != 0:
-            other: Type = call.positional[0][1]
-            unfolded_other: Type = unfold_type(other)
-            if isinstance(unfolded_other, ColumnType):
-                column2 = unfolded_other
-                col_type2: Type = column2.type
+        operand: TypedExpr = call.positional[0]
+        unfolded_operand: Type = unfold_type(operand[1])
+        col_type2: Type
 
-                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),
-                    method=method,
-                )
-                new_column = ColumnType(type=new_inner_type)
-        return new_column
+        column_operand: bool = isinstance(unfolded_operand, ColumnType)
+
+        # Operand is a column -> get the inner type
+        if column_operand:
+            col_type2 = unfolded_operand.type
+        # Otherwise use the operand type itself
+        else:
+            col_type2 = operand[1]
+
+        new_inner_type = self.typer.result_of_binary_op(
+            location=call.location,
+            expr=call.call_expr,
+            left=(call.column_expr, col_type1),
+            right=(operand[0], col_type2),
+            method=method,
+        )
+        return ColumnType(type=new_inner_type), column_operand
 
     def _element_wise(self, call: Call, method: str) -> Type:
-        # 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)
-        signature = GenericType(
-            name="add",
-            params=[param_type],
-            body=Function(
-                params=ParamSpec(
-                    mixed=[
-                        Function.Parameter(
-                            pos=0,
-                            name="other",
-                            type=ColumnType(type=param_type),
-                            required=True,
-                        ),
-                    ],
-                ),
-                returns=self._element_binary_op(call, method),
+        returns, column_operand = self._element_binary_op(call, method)
+        signature = Function(
+            params=ParamSpec(
+                mixed=[
+                    Function.Parameter(
+                        pos=0,
+                        name="other",
+                        type=TopType(),
+                        required=True,
+                    ),
+                ],
             ),
+            returns=returns,
         )
 
         # Map arguments and compute result type
@@ -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,14 +36,18 @@ class Call:
 
 
 class FrameGroupByMethodRegistry(MethodRegistry[Call]):
-    NAMED_ARGS: dict[str, str] = {
-        "numeric_only": "bool",
-        "skipna": "bool",
-        "engine": "str",
-        "engine_kwargs": "dict",
-    }
+    """The method registry for frame group-by types"""
 
     def _aggregate(self, call: Call, method: str) -> Type:
+        """Compute the result type of an aggregate method call
+
+        Args:
+            call (Call): the call object
+            method (str): the method's name
+
+        Returns:
+            Type: the result type
+        """
         new_columns: list[DataFrameType.Column] = []
 
         for column in call.groupby.frame.columns:

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
-        if len(call.positional) != 0:
-            operand: TypedExpr = call.positional[0]
-            unfolded_other: Type = unfold_type(operand[1])
-            if isinstance(unfolded_other, DataFrameType):
-                frame2 = unfolded_other
-                by_name = {
-                    col.name: col for col in frame2.columns if col.name is not None
-                }
+        # Get map of operand's columns by name, if the operand is a dataframe
+        unfolded_other: Type = unfold_type(operand[1])
+        frame_operand: bool = isinstance(unfolded_other, DataFrameType)
+        if frame_operand:
+            frame2 = unfolded_other
+            by_name = {col.name: col for col in frame2.columns if col.name is not None}
 
         # Compute new schema:
         # Step 1: for all columns in frame1:
@@ -118,10 +145,20 @@ class FrameMethodRegistry(MethodRegistry[Call]):
 
             col_type1: ColumnType = column.type
             col_type: ColumnType = ColumnType(type=UnknownType())
-            if column.name in by_name:
-                column2 = by_name[column.name]
-                col_type2: ColumnType = column2.type
 
+            col_type2: Optional[ColumnType] = None
+
+            # Operand is a frame -> lookup column with the same name
+            if frame2 is not None:
+                if column.name in by_name:
+                    column2 = by_name[column.name]
+                    col_type2 = column2.type
+
+            # Operand is not a frame -> scalar operation -> ad-hoc column
+            else:
+                col_type2 = ColumnType(type=operand[1])
+
+            if col_type2 is not None:
                 col_type = self._get_method_result(call, col_type1, col_type2, method)
 
             new_column = DataFrameType.Column(
@@ -145,10 +182,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 +207,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 +222,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 +352,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,6 +513,88 @@ 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
+
     @method()
     def groupby(self, call: Call) -> Type:
         bool_: Type = self.types.get_type("bool")
@@ -425,6 +638,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
 
@@ -29,27 +28,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 +58,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 +85,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 +113,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 +155,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 +236,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 +271,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 +297,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 +388,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:

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,10 +83,24 @@ 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)
 
@@ -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,6 +237,24 @@ 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 DataFrameType():
@@ -283,6 +314,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 +335,23 @@ class PythonTyper(
         kw: list[Function.Parameter] = []
 
         def eval_param_type(param: p.Function.Parameter) -> Type:
-            if param.type is not None:
-                return self.resolve_type_expr(param.type)
+            default_type: Optional[Type] = None
             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 +461,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 +491,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 +516,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 +542,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
@@ -533,6 +616,23 @@ class PythonTyper(
         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 +699,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 +924,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 +956,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 +1003,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 +1031,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 +1083,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 +1108,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 +1174,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 +1320,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

@@ -29,11 +29,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 +85,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 +132,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
@@ -218,6 +259,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 +384,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 +461,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:
@@ -459,18 +534,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 +589,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

@@ -7,6 +7,8 @@ from midas.checker.diagnostic import Diagnostic, DiagnosticType
 
 
 class Reporter:
+    """Helper class to store diagnostics"""
+
     def __init__(self):
         self.diagnostics: list[Diagnostic] = []
 
@@ -17,6 +19,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 +37,52 @@ 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
 
     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)
 
     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
+        """
         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 +90,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 +103,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 +116,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

@@ -78,6 +78,14 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
                 return
 
     def is_defined(self, name: str) -> bool:
+        """Check whether the given variable is defined in any scope
+
+        Args:
+            name (str): the name of the variable
+
+        Returns:
+            bool: `True` if the variable is defined in a scope, `False` otherwise
+        """
         for scope in self.scopes:
             if name in scope:
                 return True
@@ -93,6 +101,10 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
             function (p.Function): the function to resolve
         """
         self.begin_scope()
+        for param in function.params.all:
+            if param.default is not None:
+                self.resolve(param.default)
+
         for param in function.params.all:
             self.declare(param.name)
             self.define(param.name)
@@ -167,6 +179,18 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
         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(name)
+            self.define(name)
+
     def visit_raw_stmt(self, stmt: p.RawStmt) -> None:
         pass
 

midas/checker/types.py

@@ -10,12 +10,16 @@ from midas.ast.printer import MidasPrinter
 
 @dataclass(frozen=True, kw_only=True)
 class TopType:
+    """The top type (`Any`)"""
+
     def __str__(self) -> str:
         return "Any"
 
 
 @dataclass(frozen=True, kw_only=True)
 class BaseType:
+    """A base / builtin type"""
+
     name: str
 
     def __str__(self) -> str:
@@ -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,6 +105,8 @@ class ParamSpec:
 
 @dataclass(frozen=True, kw_only=True)
 class OverloadedFunction:
+    """A list of method overloads"""
+
     overloads: list[Type]
 
     def __str__(self) -> str:
@@ -95,6 +115,8 @@ class OverloadedFunction:
 
 @dataclass(frozen=True, kw_only=True)
 class ComplexType:
+    """A type with inline members"""
+
     members: dict[str, Type]
 
     def __str__(self) -> str:
@@ -104,6 +126,8 @@ class ComplexType:
 
 @dataclass(frozen=True, kw_only=True)
 class ExtensionType:
+    """An extension of a type, adding members through a `ComplexType`"""
+
     base: Type
     extension: ComplexType
 
@@ -112,6 +136,8 @@ class ExtensionType:
 
 
 class Variance(StrEnum):
+    """The variance of a :class:`TypeVar`"""
+
     INVARIANT = "INVARIANT"
     COVARIANT = "COVARIANT"
     CONTRAVARIANT = "CONTRAVARIANT"
@@ -119,6 +145,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 +164,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 +176,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 +188,8 @@ class AppliedType:
 
 @dataclass(frozen=True, kw_only=True)
 class ConstraintType:
+    """A type with a constraint expression"""
+
     type: Type
     constraint: m.Expr
 
@@ -166,6 +200,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 +210,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 +220,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 +237,8 @@ class DataFrameType:
 
 @dataclass(frozen=True, kw_only=True)
 class FrameGroupBy:
+    """A frame group-by object"""
+
     frame: DataFrameType
 
     def __str__(self) -> str:
@@ -205,6 +247,8 @@ class FrameGroupBy:
 
 @dataclass(frozen=True, kw_only=True)
 class ColumnGroupBy:
+    """A column group-by object"""
+
     column: ColumnType
 
     def __str__(self) -> str:
@@ -212,6 +256,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,
@@ -354,6 +411,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 +427,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 "..."
@@ -430,6 +504,8 @@ def to_annotation(type: Type) -> str:
 
 @dataclass(frozen=True, kw_only=True)
 class Predicate:
+    """A predicate"""
+
     type: Type
     body: m.Expr
     alias: bool

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

@@ -40,11 +40,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 +84,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 +115,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 +128,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 +179,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 +194,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 +322,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 +386,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 +458,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 +486,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 +503,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 +530,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 []
@@ -522,6 +695,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 +736,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 +754,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 +779,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 +796,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 +840,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 +886,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 +911,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 +934,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

@@ -32,6 +32,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 +45,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 +99,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 +142,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 +197,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 +228,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
@@ -273,14 +315,11 @@ class StubsGenerator:
                     ),
                 )
 
-            case ColumnType(type=inner):
+            case ColumnType():
                 self.import_pandas = True
-                return ast.Subscript(
-                    value=ast.Attribute(
-                        value=ast.Name(id="pd"),
-                        attr="Series",
-                    ),
-                    slice=self.dump_type(inner),
+                return ast.Attribute(
+                    value=ast.Name(id="pd"),
+                    attr="Series",
                 )
 
             case DataFrameType():
@@ -322,6 +361,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 +400,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 +452,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 +481,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

@@ -35,10 +35,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 +74,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 +96,14 @@ class MidasParser(Parser):
     def type_declaration(self) -> TypeStmt:
         """Parse a type declaration
 
-        A type declaration can either be a simple type alias or a new complex type.
-        In either case, it can have an optional template expression after its name, wrapped in brackets.
-        A simple type alias is derived from a base type expression, and can have a optional constraint expression preceded by the `where` keyword.
-        A full simple type alias is thus written:
-        ```
-        type Name[Template](TypeExpr) where Condition
-        ```
+        A type declaration creates a named subtype of a type expression.
+        It can have an optional template expression after its name, wrapped in brackets, to handle type parameters.
 
-        A new complex type has a set of properties which are named, have a type and an optional constraint expression (also preceded by the `where` keyword).
-        A full complex type definition is thus written:
-        ```
-        type Name[Template] {
-            prop1: TypeExpr1 where Condition1
-            prop2: TypeExpr2 where Condition2
-            ...
-        }
-        ```
+        A type statement consists of:
+        - the `type` keyword
+        - a name (identifier)
+        - (optional) type parameters
+        - a body, a type expression (see :func:`type_expr`)
 
         Returns:
             TypeStmt: the parsed type declaration statement
@@ -165,11 +157,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,8 +181,8 @@ class MidasParser(Parser):
     def type_expr(self) -> Type:
         """Parse a type expression
 
-        A type is an identifier, optionally followed by a template expression.
-        It can also optionally be followed by a '?' to indicate a nullable type
+        A type expression can either be a function type (see :func:`function`)
+        or a constraint type (see :func:`constraint_type`)
 
         Returns:
             TypeExpr: the parsed type expression
@@ -205,6 +202,15 @@ class MidasParser(Parser):
         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,6 +222,14 @@ 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")
@@ -227,6 +241,17 @@ class MidasParser(Parser):
         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 +265,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,6 +282,13 @@ 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(),
@@ -257,13 +296,13 @@ class MidasParser(Parser):
         )
 
     def complex_type(self) -> ComplexType:
-        """Parse a type definition body
+        """Parse a complex type expression
 
-        A type definition body is a set of whitespace-separated
-        property statements enclosed in curly braces
+        A complex type consists of zero or more member statements enclosed in
+        curly braces
 
         Returns:
-            ComplexType: the parsed complex type
+            ComplexType: the parsed complex type expression
         """
         left: Token = self.consume(
             TokenType.LEFT_BRACE, "Expected '{' to start type body"
@@ -285,6 +324,20 @@ class MidasParser(Parser):
         )
 
     def frame_type(self) -> FrameType:
+        """Parse a frame type expression
+
+        A frame type consists of:
+        - the `Frame` identifier
+        - an opening bracket `[`
+        - a list of comma-separated column expression consisting of:
+            - a name (token)
+            - a colon `:`
+            - a type expression (see :func:`type_expr`)
+        - a closing bracket `]`
+
+        Returns:
+            FrameType: the parsed frame type
+        """
         keyword: Token = self.previous()
         self.consume(TokenType.LEFT_BRACKET, "Expected '[' to start frame schema")
 
@@ -311,9 +364,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 +374,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 +403,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 +422,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 +446,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 +465,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 +484,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 +500,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 +557,14 @@ class MidasParser(Parser):
             else:
                 value = self.expression()
                 if self.check(TokenType.EQUAL):
+                    error_msg: str
                     if keywords:
-                        raise self.error(self.peek(), "Invalid keyword argument name")
+                        error_msg = "Invalid keyword argument name"
                     else:
-                        raise self.error(
-                            self.peek(),
-                            "Cannot pass positional arguments after a keyword argument",
+                        error_msg = (
+                            "Cannot pass positional arguments after a keyword argument"
                         )
+                    raise self.error(self.peek(), error_msg)
                 pos_args.append(value)
 
             if not self.match(TokenType.COMMA):
@@ -460,7 +581,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 +601,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
@@ -508,14 +639,41 @@ class MidasParser(Parser):
         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 +682,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 +716,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 +747,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 +775,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 +798,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)

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

tests/__main__.py

@@ -31,7 +31,7 @@ def main():
         GeneratorTester,
     ]
 
-    success: bool = all(map(run_tests, testers))
+    success: bool = all(list(map(run_tests, testers)))  # list to avoid early stop
 
     if success:
         print(Ansi.FG(Ansi.BRIGHT_GREEN) + "All tests passed!" + Ansi.RESET)

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.ref.json

@@ -318,7 +318,8 @@
               "pos": 0,
               "name": "object",
               "type": {},
-              "required": false
+              "required": false,
+              "unsupported": false
             }
           ],
           "mixed": [],

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.ref.json

@@ -2296,14 +2296,18 @@
             "index": 0,
             "name": "a",
             "type": {
-              "type": {}
+              "type": {
+                "name": "float"
+              }
             }
           },
           {
             "index": 1,
             "name": "b",
             "type": {
-              "type": {}
+              "type": {
+                "name": "float"
+              }
             }
           }
         ]
@@ -2517,14 +2521,18 @@
             "index": 0,
             "name": "a",
             "type": {
-              "type": {}
+              "type": {
+                "name": "int"
+              }
             }
           },
           {
             "index": 1,
             "name": "b",
             "type": {
-              "type": {}
+              "type": {
+                "name": "float"
+              }
             }
           }
         ]
@@ -2659,14 +2667,18 @@
             "index": 0,
             "name": "a",
             "type": {
-              "type": {}
+              "type": {
+                "name": "int"
+              }
             }
           },
           {
             "index": 1,
             "name": "b",
             "type": {
-              "type": {}
+              "type": {
+                "name": "float"
+              }
             }
           }
         ]
@@ -3687,7 +3699,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "float"
+        }
       }
     },
     {
@@ -3841,7 +3855,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "int"
+        }
       }
     },
     {
@@ -3878,7 +3894,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "int"
+        }
       }
     },
     {
@@ -3952,7 +3970,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "int"
+        }
       }
     },
     {
@@ -4167,7 +4187,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "float"
+        }
       }
     },
     {
@@ -4288,7 +4310,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "int"
+        }
       }
     },
     {
@@ -4366,7 +4390,9 @@
         "keywords": {}
       },
       "type": {
-        "type": {}
+        "type": {
+          "name": "int"
+        }
       }
     },
     {

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

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

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

@@ -2474,7 +2474,7 @@
                 "operator": "-",
                 "right": {
                   "_type": "LiteralExpr",
-                  "value": 90.0
+                  "value": 90
                 }
               },
               "operator": "<=",
@@ -2485,7 +2485,7 @@
             "operator": "<=",
             "right": {
               "_type": "LiteralExpr",
-              "value": 90.0
+              "value": 90
             }
           }
         }
@@ -2512,7 +2512,7 @@
                 "operator": "-",
                 "right": {
                   "_type": "LiteralExpr",
-                  "value": 180.0
+                  "value": 180
                 }
               },
               "operator": "<=",
@@ -2523,7 +2523,7 @@
             "operator": "<=",
             "right": {
               "_type": "LiteralExpr",
-              "value": 180.0
+              "value": 180
             }
           }
         }
@@ -2780,7 +2780,7 @@
         "operator": ">=",
         "right": {
           "_type": "LiteralExpr",
-          "value": 0.0
+          "value": 0
         }
       }
     },
@@ -2813,7 +2813,7 @@
         "operator": ">",
         "right": {
           "_type": "LiteralExpr",
-          "value": 0.0
+          "value": 0
         }
       }
     },
@@ -2848,7 +2848,7 @@
               "operator": "-",
               "right": {
                 "_type": "LiteralExpr",
-                "value": 10.0
+                "value": 10
               }
             },
             "operator": "<=",
@@ -2864,7 +2864,7 @@
           "operator": "<=",
           "right": {
             "_type": "LiteralExpr",
-            "value": 10.0
+            "value": 10
           }
         }
       }
@@ -2904,7 +2904,7 @@
           "operator": ">=",
           "right": {
             "_type": "LiteralExpr",
-            "value": 66.0
+            "value": 66
           }
         }
       }
@@ -2950,7 +2950,7 @@
                         "_type": "BinaryExpr",
                         "left": {
                           "_type": "LiteralExpr",
-                          "value": 0.0
+                          "value": 0
                         },
                         "operator": "<=",
                         "right": {
@@ -2960,7 +2960,7 @@
                       "operator": "<",
                       "right": {
                         "_type": "LiteralExpr",
-                        "value": 150.0
+                        "value": 150
                       }
                     }
                   }

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

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

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

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

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

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

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",