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

@@ -258,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(

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
@@ -99,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:
@@ -123,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
 
@@ -137,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 _:
@@ -175,6 +182,7 @@ class Evaluator(m.Expr.Visitor[Any]):
 
     def _evaluate_predicate(
         self,
+        location: Location,
         predicate: Predicate,
         args: list[Any],
         kwargs: dict[str, Any],
@@ -196,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,
@@ -205,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 _:
@@ -213,7 +221,13 @@ 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
@@ -232,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:
@@ -49,21 +49,18 @@ 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 [].
-            preserve_inner_type (bool, optional): If `True`, the result type
-                will preserve the column's inner type (e.g. for `min`/`max`),
-                otherwise the inner type is widened to `TopType`. Defaults to False.
 
         Returns:
             Type: the result type
@@ -87,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(
@@ -108,6 +113,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def kurt(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "kurt",
             ["skipna", "numeric_only"],
         )
 
@@ -115,6 +121,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def max(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "max",
             [
                 "numeric_only",
                 (
@@ -126,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"],
         )
 
@@ -140,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",
                 (
@@ -159,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",
                 (
@@ -181,6 +189,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def std(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "std",
             [
                 (
                     "ddof",
@@ -198,6 +207,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def sum(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "sum",
             [
                 "numeric_only",
                 (
@@ -215,6 +225,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def var(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "var",
             [
                 (
                     "var",

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,6 +24,31 @@ 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:
@@ -44,7 +69,72 @@ class Call:
 class ColumnMethodRegistry(MethodRegistry[Call]):
     """The method registry for column types"""
 
-    def _element_binary_op(self, call: Call, method: str) -> ColumnType:
+    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
@@ -55,28 +145,34 @@ 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
+
+        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=(call.positional[0][0], col_type2),
+            right=(operand[0], col_type2),
             method=method,
         )
-                new_column = ColumnType(type=new_inner_type)
-        return new_column
+        return ColumnType(type=new_inner_type), column_operand
 
     def _element_wise(self, call: Call, method: str) -> Type:
         """Compute the result of an element-wise method call
@@ -91,26 +187,21 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         Returns:
             Type: the result type
         """
-        # TODO: support add with scalar
 
         # 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(
+        returns, column_operand = self._element_binary_op(call, method)
+        signature = Function(
             params=ParamSpec(
                 mixed=[
                     Function.Parameter(
                         pos=0,
                         name="other",
-                            type=ColumnType(type=param_type),
+                        type=TopType(),
                         required=True,
                     ),
                 ],
             ),
-                returns=self._element_binary_op(call, method),
-            ),
+            returns=returns,
         )
 
         # Map arguments and compute result type
@@ -120,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__")
@@ -184,7 +345,7 @@ 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
 
@@ -192,13 +353,23 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
             call (Call): the call object
             kwargs (list[Function.Parameter], optional): a list of extra
                 keyword-only parameters. Defaults to [].
-            preserve_inner_type (bool, optional): If `True`, the result type
-                will preserve the column's inner type (e.g. for `min`/`max`),
-                otherwise the inner type is widened to `TopType`. Defaults to False.
+            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=[
@@ -211,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(
@@ -228,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:
@@ -266,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:
@@ -330,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")

midas/checker/frames/frame_manager.py

@@ -132,7 +132,7 @@ class FrameManager:
                 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(

midas/checker/frames/frame_methods.py

@@ -17,6 +17,7 @@ from midas.checker.types import (
     ParamSpec,
     TopType,
     Type,
+    UnitType,
     UnknownType,
     unfold_type,
 )
@@ -44,6 +45,25 @@ 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,
@@ -82,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
@@ -95,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]
+        # Get map of operand's columns by name, if the operand is a dataframe
         unfolded_other: Type = unfold_type(operand[1])
-            if isinstance(unfolded_other, DataFrameType):
+        frame_operand: bool = isinstance(unfolded_other, DataFrameType)
+        if frame_operand:
             frame2 = unfolded_other
-                by_name = {
-                    col.name: col for col in frame2.columns if col.name is not None
-                }
+            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:
@@ -122,10 +145,20 @@ class FrameMethodRegistry(MethodRegistry[Call]):
 
             col_type1: ColumnType = column.type
             col_type: ColumnType = ColumnType(type=UnknownType())
+
+            col_type2: Optional[ColumnType] = None
+
+            # Operand is a frame -> lookup column with the same name
+            if frame2 is not None:
                 if column.name in by_name:
                     column2 = by_name[column.name]
-                col_type2: ColumnType = column2.type
+                    col_type2 = column2.type
 
+            # 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(
@@ -149,7 +182,7 @@ 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:
         """Compute the result of an element-wise method call
@@ -164,7 +197,9 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         Returns:
             Type: the result type
         """
-        # TODO: support scalar, sequence, Series, dict operand
+        # 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(
@@ -172,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
@@ -187,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__")
@@ -408,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")

midas/checker/midas.py

@@ -28,7 +28,7 @@ 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
 
 
@@ -297,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}")
@@ -383,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

@@ -30,6 +30,7 @@ class Preamble(Environment):
         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

@@ -335,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
@@ -603,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
 
@@ -669,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}")
@@ -921,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)

midas/checker/resolver.py

@@ -101,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)
@@ -175,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

@@ -69,10 +69,14 @@ 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)

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(
+                return ast.Attribute(
                     value=ast.Name(id="pd"),
                     attr="Series",
-                    ),
-                    slice=self.dump_type(inner),
                 )
 
             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/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 ".":
@@ -106,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):

midas/lexer/token.py

@@ -36,6 +36,7 @@ class TokenType(Enum):
     EQUAL = auto()
     EQUAL_EQUAL = auto()
     BANG_EQUAL = auto()
+    BANG = auto()
 
     # Literals
     IDENTIFIER = auto()

midas/parser/midas.py

@@ -487,12 +487,12 @@ class MidasParser(Parser[list[Stmt]]):
         """Parse a unary expression
 
         A unary consists of a call expression (see :func:`call`) optionally
-        preceded by zero or more unary operators (`+`, `-`)
+        preceded by zero or more unary operators (`+`, `-`, `!`)
 
         Returns:
             Expr: the parsed expression
         """
-        if self.match(TokenType.PLUS, 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)

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