Merge pull request 'Cleanup and minor improvements' (#37) from feat/tidying into main

Reviewed-on: #37

docs/function_subtyping.typ

@@ -0,0 +1,97 @@
+#import "@preview/curryst:0.6.0": prooftree, rule, rule-set
+#import "@preview/gentle-clues:1.3.1" as gc
+#import "@preview/lovelace:0.3.1": pseudocode-list
+
+#set text(font: "Source Sans 3")
+#show link: set text(fill: blue)
+#set document(
+  title: [Function subtyping rule],
+)
+
+#let req = math.op("req")
+
+#align(center, title())
+
+
+This document formalizes the logic used when checking whether a function is a subtype of another.
+
+= Definitions
+
+A *Parameter specification* has a list of positional-only parameters $P$, mixed parameters $M$ and keyword-only parameters $K$:
+$ S = (P, M, K) $
+
+A *Parameter* has an index $i$, a name $n$ and a required flag $r$:
+$ p = (i, n, r) $
+
+A *Function* has a param spec $S$ and a return type $R$:
+$ F = (S, R) $
+
+= Main rules
+
+We want to define a rule for checking structural subtyping of functions, i.e. to check when $F_1 <: F_2$
+
+There are two conditions to check:
+#align(center, prooftree(
+  rule(
+    $Gamma tack S_2 <: S_1$,
+    $Gamma tack R_1 <: R_2$,
+    $Gamma tack F_1 <: F_2$,
+  ),
+))
+
+The second condition is trivial to check.
+The first condition is a bit more tricky.
+
+For a parameter specification $S_2$ to be a subtype of another $S_1$, the latter needs to be fully compatible with any usages of the former. What this means is that #link(<cond-1>)[(1)] any argument that can be passed to $S_2$ must be accepted by $S_1$, and #link(<cond-2>)[(2)] $S_1$ must not have additional required arguments that are not in $S_2$. However, $S_1$ can have additional optional parameters not present in $S_2$.
+
+After mapping parameters of $S_1$ and $S_2$, types must be checked such that if a parameter $p_i: T in S_1$ is mapped to a parameter $q_j: U in S_2$, $U <: T$.
+
+= Detailed rules for *ParamSpec* subtyping
+
+#gc.info(title: [Notation])[
+  In the following equations:
+  - the notation $S in.rev a^i: T$ will be used to denote that a parameter spec $S$ accepts an argument named $a$ at index $i$ of type $T$
+  - the special name $alpha$ will be used to denote any argument without a specific name
+  - the special name $phi$ will be used to denote any parameter without a specific name
+  - $AA$ denotes the group of all arguments, i.e. $alpha in AA$
+  - $PP$ denotes the group of all parameters, i.e. $p in PP$
+  - $req(S, alpha)$ is a predicate checking whether $alpha$ is required in $S$
+]
+
+== Arguments of $S_2$ are compatible with $S_1$ <cond-1>
+
+Formally, the condition is:
+$
+  forall alpha in AA, S_2 in.rev alpha => S_1 in.rev alpha
+$
+
+Let $S_1 = (P_1, M_1, K_1)$ and $S_2 = (P_2, M_2, K_2)$.
+
+For each positional-only parameter $phi_i in P_2$, $phi_i in P_1 or phi_i in M_1$. A positional-only parameter of $S_2$ can either be positional-only or mixed in $S_1$. Additionally, $not req(S_2, phi_i) => not req(S_1, phi_i)$. If $phi_i$ is optional in $S_2$, it must also be optional in $S_1$ so that a call omitting it is valid.
+
+Similarly, for each keyword-only parameter $p in K_2$, $p in K_1 or p in M_1$. A keyword-only parameter of $S_2$ can either be keyword-only or mixed in $S_1$. Additionally, $not req(S_2, p) => not req(S_1, p)$. If $p$ is optional in $S_2$, it must also be optional in $S_1$ so that a call omitting it is valid.
+
+Finally, for mixed parameters, the rule is slightly more complex. Either there is a corresponding mixed parameter in $S_1$, or the parameter is covered by both a positional/mixed and a keyword/mixed parameter. In the second case, we must keep in mind that only one of the match will made at runtime, or maybe even none if the parameter is optional in $S_2$, meaning the parameters in $S_1$ _must_ be optional.
+
+We can thus split the rule in two. $forall p_i in M_2$:
+$
+  cases(
+    p_i in M_1 and not req(S_2, p_i) => not req(S_1, p_i),
+    "or",
+    underbrace((phi_i in P_1 or phi_i in M_1), "Positional") and
+    underbrace((p in K_1 or p in M_1), "Keyword") and
+    underbrace((not req(S_1, phi_i) and not req(S_1, p)), "Optional")
+  )
+$
+
+== No additional required arguments in $S_1$ <cond-2>
+
+Formally, the condition is:
+$
+  forall alpha in AA, S_1 in.rev alpha and req(S_1, alpha) => S_2 in.rev alpha and req(S_2, alpha)
+$
+
+Each parameter in $S_1$ that is not matched by a parameter in $S_2$ must be optional:
+$
+  forall p_i in S_1, p_i in.not S_2 => not req(S_1, p_i)
+$

docs/manual.typ

@@ -624,7 +624,7 @@ For example:
 
 == Control flow
 
-Some control flow features are supported. For the limited code of this project, not all constructs are supported. The following are those currently handled and typ checked by Midas.
+Some control flow features are supported. For the limited code of this project, not all constructs are supported. The following are those currently handled and type checked by Midas.
 
 === `if` / `elif` / `else` <if-else>
 
@@ -757,7 +757,7 @@ If the value passed to `cast` or `unsafe_cast` is a literal (e.g. an integer, a
 
 Vanilla Python already lets you use type hints to specify the type of variables and function parameters.
 
-Midas use them to type check your code. Additionally, it allows you to use a special syntax to define a `Frame` types directly in these annotations.
+Midas use them to type check your code. Additionally, it allows you to use a special syntax to define a `Frame` type directly in these annotations.
 
 Because these annotations are not interpretable by Python, your integrated type checker might complain loudly about them being invalid.
 A workaround is to silence it by adding a type comment at the end of the line, as shown in @silence-errors.

docs/template.typ

@@ -61,6 +61,7 @@
   set document(
     title: title,
     author: author,
+    date: none,
   )
   set text(
     font: "Source Sans 3",

gen/python.py

@@ -44,11 +44,6 @@ class BaseType:
     args: tuple[MidasType, ...]
 
 
-class ConstraintType:
-    type: MidasType
-    constraint: ast.expr
-
-
 class FrameColumn:
     name: Optional[str]
     type: Optional[MidasType]

midas/ast/printer/python_ast.py

@@ -18,14 +18,6 @@ class PythonAstPrinter(
             self._write_line(f"base: {node.base}")
             self._write_sequence("args", node.args, last=True)
 
-    def visit_constraint_type(self, node: p.ConstraintType) -> None:
-        self._write_line("ConstraintType")
-        with self._child_level():
-            self._write_line("type")
-            with self._child_level(single=True):
-                node.type.accept(self)
-            self._write_line(f"constraint: {ast.unparse(node.constraint)}", last=True)
-
     def visit_frame_column(self, node: p.FrameColumn) -> None:
         self._write_line("FrameColumn")
         with self._child_level():

midas/ast/python.py

@@ -52,9 +52,6 @@ class MidasType(ABC):
         @abstractmethod
         def visit_base_type(self, node: BaseType) -> T: ...
 
-        @abstractmethod
-        def visit_constraint_type(self, node: ConstraintType) -> T: ...
-
         @abstractmethod
         def visit_frame_column(self, node: FrameColumn) -> T: ...
 
@@ -71,15 +68,6 @@ class BaseType(MidasType):
         return visitor.visit_base_type(self)
 
 
-@dataclass(frozen=True)
-class ConstraintType(MidasType):
-    type: MidasType
-    constraint: ast.expr
-
-    def accept(self, visitor: MidasType.Visitor[T]) -> T:
-        return visitor.visit_constraint_type(self)
-
-
 @dataclass(frozen=True)
 class FrameColumn(MidasType):
     name: Optional[str]

midas/checker/builtins.py

@@ -26,7 +26,11 @@ Circular dependencies and diamond inheritance MUST be avoided
 
 
 def define_builtins(reg: TypesRegistry):
-    """Define builtin types and operations"""
+    """Define builtin types and operations
+
+    Args:
+        reg (TypesRegistry): the types registry
+    """
     any = reg.define_type("Any", TopType())
     unit = reg.define_type("None", UnitType())
     object = reg.define_type("object", BaseType(name="object"))

midas/checker/dispatcher.py

@@ -102,6 +102,11 @@ class CallDispatcher(Generic[E]):
         self.logger: logging.Logger = logging.getLogger("CallDispatcher")
 
     def set_reporter(self, reporter: FileReporter):
+        """Set the current reporter
+
+        Args:
+            reporter (FileReporter): the new file reporter
+        """
         self.reporter = reporter
 
     def get_result(
@@ -123,8 +128,8 @@ class CallDispatcher(Generic[E]):
         Args:
             location (Location): the call location
             callee (Type): the called function
-            positional (list[TypedExpr]): the list of positional arguments
+            positional (list[TypedExpr[E]]): the list of positional arguments
-            keywords (dict[str, TypedExpr]): the map of keyword arguments
+            keywords (dict[str, TypedExpr[E]]): the map of keyword arguments
             report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
 
         Returns:
@@ -250,7 +255,7 @@ class CallDispatcher(Generic[E]):
         """Check whether the passed argument types correspond to their matched parameter definitions
 
         Args:
-            arguments (list[MappedArgument]): the list of argument/parameter pairs
+            arguments (list[MappedArgument[E]]): the list of argument/parameter pairs
             report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
 
         Returns:
@@ -286,8 +291,8 @@ class CallDispatcher(Generic[E]):
         Args:
             overloads (list[Type]): the list of possible overloads
             location (Location): the call location
-            positional (list[TypedExpr]): the list of positional arguments
+            positional (list[TypedExpr[E]]): the list of positional arguments
-            keywords (dict[str, TypedExpr]): the map of keywords arguments
+            keywords (dict[str, TypedExpr[E]]): the map of keywords arguments
             report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
 
         Returns:
@@ -385,8 +390,8 @@ class CallDispatcher(Generic[E]):
         Args:
             function (Function): the function definition
             location (Location): the call location
-            positional (list[TypedExpr]): the list of positional arguments
+            positional (list[TypedExpr[E]]): the list of positional arguments
-            keywords (dict[str, TypedExpr]): the map of keyword arguments
+            keywords (dict[str, TypedExpr[E]]): the map of keyword arguments
             report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
 
         Returns:
@@ -514,8 +519,8 @@ class CallDispatcher(Generic[E]):
         function / a subtype of another.
 
         Args:
-            mapped1 (list[MappedArgument]): the first argument mappings (subtype)
+            mapped1 (list[MappedArgument[E]]): the first argument mappings (subtype)
-            mapped2 (list[MappedArgument]): the second argument mappings (supertype)
+            mapped2 (list[MappedArgument[E]]): the second argument mappings (supertype)
 
         Returns:
             bool: `True` if `mapped1` is a subtype of `mapped2`, `False` otherwise

midas/checker/evaluator.py

@@ -190,6 +190,7 @@ class Evaluator(m.Expr.Visitor[Any]):
         """Evaluate a predicate function call
 
         Args:
+            location (Location): the location of the call expression
             predicate (Predicate): the predicate to evaluate
             args (list[Any]): a list of positional arguments
             kwargs (dict[str, Any]): a map of keyword arguments
@@ -234,6 +235,7 @@ class Evaluator(m.Expr.Visitor[Any]):
         is set in the context using :func:`set_value` with the parameter's name
 
         Args:
+            location (Location): the location of the call expression
             function (Function): the called function
             args (list[Any]): a list of positional arguments
             kwargs (dict[str, Any]): a map of keyword arguments

midas/checker/frames/column_manager.py

@@ -45,7 +45,7 @@ class ColumnManager:
         Args:
             reporter (FileReporter): the file reporter to use for diagnostics
             location (Location): the subscript's location
-            column (DataFrameType): the column type
+            column (ColumnType): the column type
             index (TypedExpr): the index
 
         Returns:

midas/checker/frames/column_methods.py

@@ -353,10 +353,12 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
             call (Call): the call object
             kwargs (list[Function.Parameter], optional): a list of extra
                 keyword-only parameters. Defaults to [].
-            formula (Callable[[Type], Formula], optional): optional formula
+            formula (Optional[Callable[[Type], Formula]], optional):
-                builder function to compute the return type. If set, the function
+                optional formula builder function to compute the return type.<br>
-                should accept the inner column type and return a formula.
+                If set, the function should accept the inner column type and
-                If `None`, the result is typed as `Column[Any]`. Defaults to None.
+                return a formula.<br>
+                If `None`, the result is typed as `Column[Any]`.
+                Defaults to None.
 
         Returns:
             Type: the result type

midas/checker/frames/frame_manager.py

@@ -232,7 +232,7 @@ class FrameManager:
             if col.name == name:
                 index = i
                 replace = True
-                # TODO: check column type here to prevent changing it
+                # TODO: might want to check column type here to disallow changing it
             new_columns.append(col)
 
         new_col: DataFrameType.Column = DataFrameType.Column(

midas/checker/midas.py

@@ -24,7 +24,7 @@ from midas.checker.types import (
     TypeVar,
     UnknownType,
 )
-from midas.checker.variance import VarianceInferrer
+from midas.checker.variance import VarianceManager
 from midas.lexer.midas import MidasLexer
 from midas.lexer.token import Token, TokenType
 from midas.parser.midas import MidasParser
@@ -147,10 +147,8 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         for stmt in stmts:
             stmt.accept(self)
 
-        for name, type in self.types._types.items():
+        manager: VarianceManager = VarianceManager(self.types)
-            if isinstance(type, GenericType):
+        manager.infer_all()
-                inferrer = VarianceInferrer(self.types)
-                self.types._types[name] = inferrer.infer(type)
 
     def assert_bool(self, expr: m.Expr):
         """Check that the given expression is a subtype of `bool` or report an error

midas/checker/python.py

@@ -454,7 +454,6 @@ class PythonTyper(
         self.env.define(stmt.name, function)
 
     def visit_type_assign(self, stmt: p.TypeAssign) -> None:
-        # TODO check not yet defined locally
         type: Type = self.resolve_type_expr(stmt.type)
         self.env.define(stmt.name, type)
 
@@ -487,11 +486,10 @@ class PythonTyper(
                 self._assign_sub(location, var, index, value_type)
 
             case _:
-                if not isinstance(target, p.VariableExpr):
+                self.logger.warning(f"Unsupported assignment to {target}")
-                    self.logger.warning(f"Unsupported assignment to {target}")
+                self.reporter.warning(
-                    self.reporter.warning(
+                    target.location, f"Unsupported assignment to {target}"
-                        target.location, f"Unsupported assignment to {target}"
+                )
-                    )
 
     def _assign_var(self, location: Location, target: p.VariableExpr, value_type: Type):
         """Type check assignment to the given target
@@ -519,11 +517,12 @@ class PythonTyper(
     def _assign_attr(
         self, location: Location, object: p.Expr, name: str, value_type: Type
     ):
-        """Type check assignment to the given target
+        """Type check assignment to the given attribute target
 
         Args:
             location (Location): the location of the assignment
-            target (p.VariableExpr): the assignment's target
+            object (p.Expr): the target attribute's owner object
+            name (str): the target attribute's name
             value_type (Type): the value to be assigned
         """
         object_type: Type = self.type_of(object)
@@ -545,11 +544,15 @@ class PythonTyper(
         index: p.Expr,
         value_type: Type,
     ):
-        """Type check assignment to the given target
+        """Type check assignment to the given subscript target
 
         Args:
             location (Location): the location of the assignment
-            target (p.VariableExpr): the assignment's target
+            var (p.VariableExpr): the target subscript's owner. We only allow
+                a variable expression here because we might modify its type (for
+                example when assigning a column to a dataframe) and reference
+                types are not implemented
+            index (p.Expr): the target subscript's index expression
             value_type (Type): the value to be assigned
         """
         var_type: Type = self.type_of(var)
@@ -691,6 +694,20 @@ class PythonTyper(
         right: TypedExpr,
         method: str,
     ) -> Type:
+        """Compute the result type of a binary operation method call
+
+        This method is called for dunder methods called by binary operators
+
+        Args:
+            location (Location): the location of the operation
+            expr (p.Expr): the expression which triggered this resolution
+            left (TypedExpr): the left operand
+            right (TypedExpr): the right operand
+            method (str): the method name
+
+        Returns:
+            Type: the result type
+        """
         try:
             return self.call_method(
                 location=location,
@@ -844,7 +861,7 @@ class PythonTyper(
     def visit_ternary_expr(self, expr: p.TernaryExpr) -> Type:
         test_type: Type = self.type_of(expr.test)
 
-        # TODO Allow subtypes or any type
+        # Strict: test must be a subtype of bool, or UnknownType
         if (
             not self.is_subtype(test_type, self.types.get_type("bool"))
             and test_type != UnknownType()
@@ -986,10 +1003,6 @@ class PythonTyper(
             return self.types.apply_generic(base, args)
         return base
 
-    def visit_constraint_type(self, node: p.ConstraintType) -> Type:
-        self.reporter.warning(node.location, "ConstraintType not yet supported")
-        return UnknownType()
-
     def visit_frame_column(self, node: p.FrameColumn) -> ColumnType:
         return ColumnType(
             type=(
@@ -1154,8 +1167,9 @@ class PythonTyper(
                         return False, None
 
                     if key is None:
-                        # TODO: check that value is always a dict
+                        # If literal value is not a dict, invalid Python -> abort
-                        assert isinstance(value_val, dict)
+                        if not isinstance(value_val, dict):
+                            return False, None
                         pairs.extend(value_val.items())
                     else:
                         pairs.append((key_val, value_val))
@@ -1281,9 +1295,7 @@ class PythonTyper(
 
             case BaseType():
                 # TODO: do we want to allow cast(float, int)? would require runtime conversion
-                if not self.types.is_subtype(
+                if not self.types.are_equivalent(subject_type, target_type):
-                    subject_type, target_type
-                ) or not self.types.is_subtype(target_type, subject_type):
                     self.reporter.error(
                         expr.location,
                         f"Value {lit_value!r} of type {subject_type} cannot be cast as {target_type}",

midas/checker/registry.py

@@ -1,6 +1,6 @@
 import logging
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, TypeAlias
 
 from midas.ast.midas import MemberKind
 from midas.checker.builtins import BUILTIN_SUBTYPES
@@ -24,6 +24,8 @@ from midas.checker.types import (
     substitute_typevars,
 )
 
+Match: TypeAlias = tuple[Function.Parameter, Function.Parameter]
+
 
 @dataclass
 class Member:
@@ -213,7 +215,6 @@ class TypesRegistry:
                 return True
 
             case (ColumnType(type=inner1), ColumnType(type=inner2)):
-                # TODO: invariant, replace ColumnType with simple GenericType
                 if not self.are_equivalent(inner1, inner2):
                     return False
                 return True
@@ -260,7 +261,6 @@ class TypesRegistry:
         """
         return self.is_subtype(type1, type2) and self.is_subtype(type2, type1)
 
-    # TODO: verify the logic in here
     def is_func_subtype(self, func1: Function, func2: Function) -> bool:
         """Check whether a function is a subtype of another
 
@@ -271,14 +271,24 @@ class TypesRegistry:
         Returns:
             bool: whether `func1` is a subtype of `func2`
         """
+
+        # Let func1 = (S1, R1) where S1 = (P1, M1, K1)
+        # Let func2 = (S2, R2) where S2 = (P2, M2, K2)
+        # We want to check that func1 <: func2
+        # i.e. R1 <: R2 and S2 <: S1
+
+        # R1 <: R2
         if not self.is_subtype(func1.returns, func2.returns):
             return False
 
+        # Extract P1, M1, K1
         pos1: list[Function.Parameter] = func1.params.pos
         mixed1: list[Function.Parameter] = func1.params.mixed
         kw1: dict[str, Function.Parameter] = {
             param.name: param for param in func1.params.kw
         }
+
+        # Extract P2, M2, K2
         pos2: list[Function.Parameter] = func2.params.pos
         mixed2: list[Function.Parameter] = func2.params.mixed
         kw2: dict[str, Function.Parameter] = {
@@ -286,89 +296,118 @@ class TypesRegistry:
         }
 
         mixed_by_pos: dict[int, Function.Parameter] = {
-            param.pos: param for param in mixed2
+            param.pos: param for param in mixed1
         }
         mixed_by_name: dict[str, Function.Parameter] = {
-            param.name: param for param in mixed2
+            param.name: param for param in mixed1
         }
 
-        def is_arg_subtype(sub: Function.Parameter, sup: Function.Parameter) -> bool:
+        matches: list[Match] = []
-            if not self.is_subtype(sub.type, sup.type):
-                return False
-            if not sup.required and sub.required:
-                return False
-            return True
 
-        for param1 in pos1:
+        # Each parameter at position i in P2 must be valid at position i in S1
-            param2: Function.Parameter
+        # either as a positional-only parameter in P1
-            if param1.pos < len(pos2):
+        # or a mixed parameter in M1
-                param2 = pos2[param1.pos]
-            elif param1.pos in mixed_by_pos:
-                param2 = mixed_by_pos[param1.pos]
-            elif not param1.required:
-                continue
-            else:
-                return False
-            if not is_arg_subtype(param2, param1):
-                return False
-
-        for name, param1 in kw1.items():
-            param2: Function.Parameter
-            if name in kw2:
-                param2 = kw2[name]
-            elif name in mixed_by_name:
-                param2 = mixed_by_name[name]
-            elif not param1.required:
-                continue
-            else:
-                return False
-            if not is_arg_subtype(param2, param1):
-                return False
-
-        for param1 in mixed1:
-            pos_param2: Optional[Function.Parameter] = None
-            kw_param2: Optional[Function.Parameter] = None
-            if param1.name in kw2:
-                kw_param2 = kw2[param1.name]
-            elif param1.name in mixed_by_name:
-                kw_param2 = mixed_by_name[param1.name]
-            if param1.pos < len(pos2):
-                pos_param2 = pos2[param1.pos]
-            elif param1.pos in mixed_by_pos:
-                pos_param2 = mixed_by_pos[param1.pos]
-
-            # No match in func2 and arg is required
-            if pos_param2 is None and kw_param2 is None and param1.required:
-                return False
-
-            # Matching keyword argument
-            if kw_param2 is not None and not is_arg_subtype(kw_param2, param1):
-                return False
-
-            # Matching positional argument
-            if pos_param2 is not None and not is_arg_subtype(pos_param2, param1):
-                return False
-
-        mixed_positions: set[int] = {param.pos for param in mixed1}
-        mixed_names: set[str] = {param.name for param in mixed1}
         for param2 in pos2:
-            if not param2.required:
+            param1: Function.Parameter
-                continue
+
-            if param2.pos >= len(pos1) and param2.pos not in mixed_positions:
+            # In P1
+            if param2.pos < len(pos1):
+                param1 = pos1[param2.pos]
+
+            # In M1
+            elif param2.pos in mixed_by_pos:
+                param1 = mixed_by_pos[param2.pos]
+
+            else:
                 return False
 
+            # not req(p2) => not req(p1)
+            if not param2.required and param1.required:
+                return False
+            matches.append((param1, param2))
+
+        # Each parameter named p in K2 must be valid with name p in S1
+        # either as a keyword-only parameter in K1
+        # or a mixed parameter in M1
         for name, param2 in kw2.items():
-            if not param2.required:
+            param1: Function.Parameter
-                continue
+
-            if name not in kw1 and name not in mixed_names:
+            # In K1
+            if name in kw1:
+                param1 = kw1[name]
+
+            # in M1
+            elif name in mixed_by_name:
+                param1 = mixed_by_name[name]
+            else:
                 return False
 
+            # not req(p2) => not req(p1)
+            if not param2.required and param1.required:
+                return False
+            matches.append((param1, param2))
+
+        # Each parameter named p at position i in M2 must be valid with name p
+        # in S1 *and* at position i in S1
+        # either as a single mixed parameter in M1
+        # or split into a positional parameter in P1/M1
+        # and a keyword parameter in K1/M1
         for param2 in mixed2:
-            if param2.required:
+            pos_param1: Optional[Function.Parameter] = None
-                continue
+            kw_param1: Optional[Function.Parameter] = None
-            pos_match: bool = param2.pos < len(pos1) or param2.pos in mixed_positions
+
-            kw_match: bool = param2.name in kw1 or param2.name in mixed_names
+            # By name in K1
-            if not pos_match or not kw_match:
+            if param2.name in kw1:
+                kw_param1 = kw1[param2.name]
+
+            # By name in M1
+            elif param2.name in mixed_by_name:
+                kw_param1 = mixed_by_name[param2.name]
+
+            # By pos in P1
+            if param2.pos < len(pos1):
+                pos_param1 = pos1[param2.pos]
+
+            # By pos in M1
+            elif param2.pos in mixed_by_pos:
+                pos_param1 = mixed_by_pos[param2.pos]
+
+            # Not fully covered
+            if pos_param1 is None or kw_param1 is None:
+                return False
+
+            # Covered by unique mixed parameter in M1
+            if pos_param1 == kw_param1:
+                param1: Function.Parameter = pos_param1
+                # not req(p2) => not req(p1)
+                if not param2.required and param1.required:
+                    return False
+
+                matches.append((param1, param2))
+
+            else:
+                # not req(p1)
+                if pos_param1.required or kw_param1.required:
+                    return False
+
+                matches.append((pos_param1, param2))
+                matches.append((kw_param1, param2))
+
+        def is_matched(param: Function.Parameter) -> bool:
+            for p1, _ in matches:
+                if p1 == param:
+                    return True
+            return False
+
+        all_params1: list[Function.Parameter] = pos1 + mixed1 + list(kw1.values())
+
+        for param1 in all_params1:
+            # No new required parameters
+            if not is_matched(param1) and param1.required:
+                return False
+
+        for param1, param2 in matches:
+            if not self.is_subtype(param2.type, param1.type):
                 return False
 
         return True

midas/checker/reporter.py

@@ -70,6 +70,11 @@ class FileReporter:
 
     @contextmanager
     def with_context(self, ctx: str):
+        """Push given context for reports inside this manager and pop it on exit
+
+        Args:
+            ctx (str): the context to temporarily push on the stack
+        """
         self._context.append(ctx)
         try:
             yield

midas/checker/types.py

@@ -437,8 +437,8 @@ def to_annotation(type: Type) -> str:
         case AppliedType(name=name, args=args):
             return f"{name}[{', '.join(map(to_annotation, args))}]"
 
-        case ConstraintType():
+        case ConstraintType(type=base):
-            return str(type)
+            return to_annotation(base)
 
         case TupleType(items=items):
             return f"Tuple[{', '.join(map(to_annotation, items))}]"

midas/checker/variance.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Literal, Optional, cast
 
 from midas.checker.registry import Member, TypesRegistry
@@ -73,10 +75,14 @@ class Tracker:
 class VarianceInferrer:
     """Helper class to compute type parameter variance"""
 
-    def __init__(self, types: TypesRegistry) -> None:
+    def __init__(self, manager: VarianceManager) -> None:
-        self.types: TypesRegistry = types
+        self.manager: VarianceManager = manager
         self.tracker: Tracker = Tracker([])
 
+    @property
+    def types(self) -> TypesRegistry:
+        return self.manager.types
+
     def infer(self, type: GenericType) -> GenericType:
         """Infer the variance of a generic type's parameters
 
@@ -150,11 +156,13 @@ class VarianceInferrer:
             # Get inferred variance of parameters and multiply with current
             # polarity to recurse through arguments
             case AppliedType(name=name, args=args):
-                # TODO: handle mutually recursive types
+                if self.manager.is_in_queue(name):
-                if name == base_name:
                     return
+
                 generic: Type = self.types.get_type(name)
                 assert isinstance(generic, GenericType)
+                generic = self.manager.infer(name, generic)
+
                 params: list[TypeVar] = generic.params
                 polarities: dict[Variance, Polarity] = {
                     Variance.INVARIANT: 0,
@@ -179,3 +187,66 @@ class VarianceInferrer:
             case TypeVar():
                 if type in self.tracker:
                     self.tracker.record(type, polarity)
+
+
+class VarianceManager:
+    """Coordinator for VarianceInferrer to handle recursive types"""
+
+    def __init__(self, types: TypesRegistry) -> None:
+        self.types: TypesRegistry = types
+        self._queue: list[str] = []
+        self._inferred: set[str] = set()
+
+    def infer_all(self):
+        """Infer variance on all generic types defined in the registry"""
+
+        for name, type in self.types._types.items():
+            if isinstance(type, GenericType):
+                self.infer(name, type)
+
+    def infer(self, name: str, type: GenericType) -> GenericType:
+        """Infer variance of parameters of the given type
+
+        Args:
+            name (str): the type's name
+            type (GenericType): the type
+
+        Returns:
+            GenericType: a new generic type with its parameters updated with
+                their inferred variance
+        """
+        if self.is_inferred(name):
+            return type
+
+        self._queue.append(name)
+
+        inferrer: VarianceInferrer = VarianceInferrer(self)
+        inferred: GenericType = inferrer.infer(type)
+        self.types._types[name] = inferred
+
+        self._queue.pop()
+        self._inferred.add(name)
+
+        return inferred
+
+    def is_in_queue(self, name: str) -> bool:
+        """Whether the given type's variance is currently being inferred
+
+        Args:
+            name (str): the type's name
+
+        Returns:
+            bool: whether the type is in the queue
+        """
+        return name in self._queue
+
+    def is_inferred(self, name: str) -> bool:
+        """Whether the given type's variance has already been inferred
+
+        Args:
+            name (str): the type's name
+
+        Returns:
+            bool: whether the type has been processed
+        """
+        return name in self._inferred

midas/cli/highlighter.py

@@ -138,10 +138,6 @@ class PythonHighlighter(
             self.wrap(arg, "arg")
             arg.accept(self)
 
-    def visit_constraint_type(self, node: p.ConstraintType) -> None:
-        self.wrap(node, "constraint-type")
-        node.type.accept(self)
-
     def visit_frame_column(self, node: p.FrameColumn) -> None:
         self.wrap(node, "frame-column")
         if node.type is not None:

midas/cli/hl_python.css

@@ -7,10 +7,6 @@ span {
         --col: 103, 192, 224;
     }
 
-    &.constraint-type {
-        --col: 174, 200, 195;
-    }
-
     &.frame-column {
         --col: 216, 231, 81;
     }

midas/generator/generator.py

@@ -269,7 +269,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         alias: ast.expr = self._make_alias(expr.expr, expr2)
 
         type: Type = self._get_expr_type(expr)
-        asserts: list[ast.stmt] = self._make_cast_asserts(expr.location, alias, type)
+        asserts: list[ast.stmt] = self._make_cast_asserts(
+            expr.location, alias, type, context=[]
+        )
         for assert_ in asserts:
             self._add_assert(assert_)
 
@@ -352,7 +354,6 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         )
 
     def visit_type_assign(self, stmt: p.TypeAssign) -> ast.stmt:
-        # TODO: is that ok?
         return ast.Pass()
 
     def visit_assign_stmt(self, stmt: p.AssignStmt) -> ast.stmt:
@@ -526,7 +527,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         raise RuntimeError(f"Cannot get type judgement for {query}")
 
     def _make_cast_asserts(
-        self, src_location: Location, expr: ast.expr, type: Type
+        self,
+        src_location: Location,
+        expr: ast.expr,
+        type: Type,
+        *,
+        context: list[str],
     ) -> list[ast.stmt]:
         """Generate assertions for the given cast expression
 
@@ -535,6 +541,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 the source file
             expr (ast.expr): the expression being cast
             type (Type): the target type
+            context (list[str]): the current context
 
         Returns:
             list[ast.stmt]: the generated assertion statements
@@ -551,12 +558,16 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                             args=[expr, ast.Name(id=name)],
                             keywords=[],
                         ),
-                        self._make_cast_assert_message(src_location, expr, type),
+                        self._make_cast_assert_message(
+                            src_location, expr, type, context=context
+                        ),
                     )
                 ]
 
             case DerivedType(type=base):
-                return self._make_cast_asserts(src_location, expr, base)
+                return self._make_cast_asserts(
+                    src_location, expr, base, context=context
+                )
 
             case UnitType():
                 return [
@@ -568,19 +579,25 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                                 ast.Constant(value=None),
                             ],
                         ),
-                        self._make_cast_assert_message(src_location, expr, type),
+                        self._make_cast_assert_message(
+                            src_location, expr, type, context=context
+                        ),
                     ),
                 ]
 
             case AppliedType(body=body):
-                return self._make_cast_asserts(src_location, expr, body)
+                return self._make_cast_asserts(
+                    src_location, expr, body, context=context
+                )
 
             case ConstraintType(type=base, constraint=constraint):
                 asserts: list[ast.stmt] = self._make_cast_asserts(
-                    src_location, expr, base
+                    src_location, expr, base, context=context
                 )
                 asserts.append(
-                    self._make_constraint_assert(src_location, expr, constraint)
+                    self._make_constraint_assert(
+                        src_location, expr, constraint, context=context
+                    )
                 )
                 return asserts
 
@@ -588,7 +605,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 # TODO: check with type from arguments / use call-site context
                 if bound is None:
                     return []
-                return self._make_cast_asserts(src_location, expr, bound)
+                return self._make_cast_asserts(
+                    src_location, expr, bound, context=context
+                )
 
             case TupleType(items=items):
                 asserts: list[ast.stmt] = [
@@ -598,13 +617,17 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                             args=[expr, ast.Name(id="tuple")],
                             keywords=[],
                         ),
-                        self._make_cast_assert_message(src_location, expr, type),
+                        self._make_cast_assert_message(
+                            src_location, expr, type, context=context
+                        ),
                     ),
                 ]
                 assert isinstance(expr, ast.Tuple)
                 for item, item_type in zip(expr.elts, items):
                     asserts.extend(
-                        self._make_cast_asserts(src_location, item, item_type)
+                        self._make_cast_asserts(
+                            src_location, item, item_type, context=context
+                        )
                     )
                 return asserts
 
@@ -618,7 +641,11 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                             keywords=[],
                         ),
                         self._make_cast_assert_message(
-                            src_location, expr, type, ": Not a dataframe"
+                            src_location,
+                            expr,
+                            type,
+                            context=context,
+                            extra=": Not a dataframe",
                         ),
                     ),
                 ]
@@ -634,7 +661,8 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                                 src_location,
                                 expr,
                                 type,
-                                f": Missing column {column.name}",
+                                context=context,
+                                extra=f": Missing column '{column.name}'",
                             ),
                         )
                     )
@@ -645,6 +673,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                                 value=expr, slice=ast.Constant(value=column.name)
                             ),
                             column.type,
+                            context=context + [f"in column '{column.name}'"],
                         )
                     )
                 return asserts
@@ -659,12 +688,19 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                             keywords=[],
                         ),
                         self._make_cast_assert_message(
-                            src_location, expr, type, ": Not a column"
+                            src_location,
+                            expr,
+                            type,
+                            context=context,
+                            extra=": Not a column",
                         ),
                     ),
                 ]
                 inner_assert: Optional[ast.stmt] = self._make_column_inner_assert(
-                    src_location, expr, type
+                    src_location,
+                    expr,
+                    type,
+                    context,
                 )
                 if inner_assert is not None:
                     asserts.append(inner_assert)
@@ -689,7 +725,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         location: Location,
         expr: ast.expr,
         type: Type,
-        extra: Optional[str] = None,
+        *,
+        context: list[str],
+        extra: str = "",
     ) -> ast.expr:
         """Build an AST node for a cast assertion message
 
@@ -703,12 +741,14 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 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
+            context (list[str]): the current context
-                the message. Defaults to None.
+            extra (str, optional): extra text to append at the end of
+                the message. Defaults to "".
 
         Returns:
             ast.expr: the generated message (as an f-string)
         """
+        context_str: str = "".join(map(lambda c: f", {c}", context))
         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(
@@ -725,12 +765,17 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                     ),
                     conversion=-1,
                 ),
-                ast.Constant(f" to {type}{extra or ''}"),
+                ast.Constant(f" to {type}{context_str}{extra}"),
             ]
         )
 
     def _make_constraint_assert(
-        self, src_location: Location, expr: ast.expr, constraint: m.Expr
+        self,
+        src_location: Location,
+        expr: ast.expr,
+        constraint: m.Expr,
+        *,
+        context: list[str],
     ) -> ast.stmt:
         """Build an assertion for the given constraint on the given expression
 
@@ -739,6 +784,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 source file
             expr (ast.expr): the expression subject to `constraint`
             constraint (m.Expr): the constraint applied on `expr`
+            context (list[str]): the current context
 
         Returns:
             ast.stmt: the assert statement checking the constraint
@@ -750,11 +796,13 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 args=[expr],
                 keywords=[],
             ),
-            self._make_constraint_assert_message(src_location, constraint),
+            self._make_constraint_assert_message(
+                src_location, constraint, context=context
+            ),
         )
 
     def _make_constraint_assert_message(
-        self, location: Location, constraint: m.Expr
+        self, location: Location, constraint: m.Expr, *, context: list[str]
     ) -> ast.expr:
         """Build an assert message for the given constraint
 
@@ -762,16 +810,18 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
             location (Location): the location of the cast expression in the
                 source file
             constraint (m.Expr): the constraint
+            context (list[str]): the current context
 
         Returns:
             ast.expr: the assert message
         """
         printer = MidasPrinter()
         constraint_str: str = printer.print(constraint)
+        context_str: str = "".join(map(lambda c: f", {c}", context))
         loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
         # f"file.py:L1:1: ConstraintError: Value does not fit constraint 'v > 0'"
         return ast.Constant(
-            f"{loc_str}: ConstraintError: Value does not fit constraint '{constraint_str}'"
+            f"{loc_str}: ConstraintError: Value does not fit constraint '{constraint_str}'{context_str}"
         )
 
     def _get_constraint(self, expr: m.Expr) -> ast.expr:
@@ -880,7 +930,11 @@ 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
+        self,
+        src_location: Location,
+        column: ast.expr,
+        type: ColumnType,
+        context: list[str],
     ) -> Optional[ast.stmt]:
         """Build a for-loop checking the type of values inside a column
 
@@ -889,18 +943,21 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 source file
             column (ast.expr): the column being cast
             type (ColumnType): the type of the column
+            context (list[str]): the current context
 
         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")
+        value: ast.expr = ast.Name(id="value")
-        body: list[ast.stmt] = self._make_cast_asserts(src_location, col, type.type)
+        body: list[ast.stmt] = self._make_cast_asserts(
+            src_location, value, type.type, context=context
+        )
         if len(body) == 0:
             return None
         return ast.For(
-            target=col,
+            target=value,
             iter=column,
             body=body,
             orelse=[],

midas/generator/stubs.py

@@ -53,7 +53,7 @@ class StubsGenerator:
         self.import_pandas = False
         for name, type in self.types._types.items():
             # Skip builtin types, not just based on name so the user can override
-            # TODO: check if added members on builtin type
+            # TODO: check if added members on builtin type, or prevent it
             match type:
                 case BaseType(name=name_) if name == name_:
                     continue
@@ -105,7 +105,9 @@ class StubsGenerator:
         """
         base_type: Type = type
 
-        # TODO: improve
+        # Generate simple assignment for type aliases
+        # A type alias will have a name that is different from the type represents
+        # or will neither be a `DeriveType` nor a `GenericType`
         match type:
             case DerivedType(name=name_) | GenericType(name=name_) if name_ == name:
                 pass

midas/parser/python.py

@@ -9,7 +9,6 @@ from midas.ast.python import (
     CallExpr,
     CastExpr,
     CompareExpr,
-    ConstraintType,
     DictExpr,
     Expr,
     ExpressionStmt,
@@ -337,30 +336,6 @@ class PythonParser:
                     args=(),
                 )
 
-            case ast.BinOp(left=left_expr, op=ast.Add(), right=right_expr):
-                left = self._parse_type(left_expr)
-                match left:
-                    # If chained constraints, separate base type and rebuild constraint
-                    case ConstraintType(type=left_type, constraint=left_constraint):
-                        constraint = ast.BinOp(
-                            left=left_constraint,
-                            op=ast.Add(),
-                            right=right_expr,
-                        )
-                        ast.copy_location(constraint, type_expr)
-                        return ConstraintType(
-                            location=loc,
-                            type=left_type,
-                            constraint=constraint,
-                        )
-
-                    case _:
-                        return ConstraintType(
-                            location=loc,
-                            type=left,
-                            constraint=right_expr,
-                        )
-
             case ast.Constant(value=None):
                 return BaseType(
                     location=loc,

scripts/docstring_checker.py

@@ -0,0 +1,174 @@
+import ast
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional
+
+
+@dataclass
+class ArgDoc:
+    name: str
+    type: str
+    optional: bool
+
+
+@dataclass
+class Param:
+    name: str
+    annotation: Optional[str]
+    optional: bool
+
+
+class Checker(ast.NodeVisitor):
+    def _get_args(self, docstring: str) -> list[ArgDoc]:
+        args: list[ArgDoc] = []
+
+        in_args: bool = False
+        for line in docstring.splitlines():
+            if not in_args:
+                if line == "Args:":
+                    in_args = True
+                continue
+
+            # End of args
+            if not line.startswith("    "):
+                break
+
+            # Continuation line
+            if line.startswith("     "):
+                continue
+
+            line = line.strip()
+            m = re.match(r"(?P<name>\w+) \((?P<type>.*?)(?P<opt>, optional)?\):", line)
+            if m is None:
+                continue
+
+            args.append(
+                ArgDoc(
+                    name=m.group("name"),
+                    type=m.group("type"),
+                    optional=m.group("opt") is not None,
+                )
+            )
+
+        return args
+
+    def log(self, node: ast.FunctionDef, msg: str):
+        loc: str = f"{node.name} L{node.lineno}:{node.col_offset+1}"
+        print(f"    ({loc}) {msg}")
+
+    def _is_ignored(self, node: ast.FunctionDef) -> bool:
+        name: str = node.name
+        if name.startswith("visit_") or name.startswith("_visit_"):
+            return True
+        if name.startswith("parse_") or name.startswith("_parse_"):
+            return True
+        if name.startswith("_print"):
+            return True
+        if name.startswith("_write"):
+            return True
+        if name.startswith("__") and name.endswith("__"):
+            return True
+        if name == "accept":
+            return True
+        node.decorator_list
+        match node:
+            case ast.FunctionDef(
+                decorator_list=[
+                    ast.Call(
+                        func=ast.Name(id="method"),
+                    ),
+                ],
+            ):
+                return True
+        return False
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
+        docstring: Optional[str] = ast.get_docstring(node)
+        func_name: str = node.name
+        if docstring is None:
+            if not self._is_ignored(node):
+                self.log(node, f"Missing docstring for function {func_name}")
+            return
+
+        args_doc: list[ArgDoc] = self._get_args(docstring)
+        by_name: dict[str, ArgDoc] = {}
+        for doc in args_doc:
+            if doc.name in by_name:
+                self.log(node, f"Multiple documentation lines for argument {doc.name}")
+            by_name[doc.name] = doc
+
+        all_params: list[Param] = []
+
+        pos_args: list[ast.arg] = node.args.posonlyargs
+        mixed_args: list[ast.arg] = node.args.args
+        kw_args: list[ast.arg] = node.args.kwonlyargs
+
+        def add_param(arg: ast.arg, optional: bool):
+            all_params.append(
+                Param(
+                    name=arg.arg,
+                    annotation=(
+                        ast.unparse(arg.annotation)
+                        if arg.annotation is not None
+                        else None
+                    ),
+                    optional=optional,
+                )
+            )
+
+        n_pos: int = len(pos_args) + len(mixed_args)
+        for i, arg in enumerate(pos_args):
+            j: int = n_pos - i - 1
+            optional: bool = j < len(node.args.defaults)
+            add_param(arg, optional)
+
+        for i, arg in enumerate(mixed_args):
+            j: int = len(mixed_args) - i - 1
+            optional: bool = j < len(node.args.defaults)
+            add_param(arg, optional)
+
+        for arg, default in zip(kw_args, node.args.kw_defaults):
+            optional: bool = default is not None
+            add_param(arg, optional)
+
+        for param in all_params:
+            doc: Optional[ArgDoc] = by_name.get(param.name, None)
+            if doc is None:
+                if param.name not in {"self", "cls"}:
+                    self.log(
+                        node, f"Missing documentation for parameter '{param.name}'"
+                    )
+                continue
+
+            if doc.name != param.name:
+                self.log(node, f"Documentation mismatch for '{param.name}': wrong name")
+
+            if doc.type != param.annotation:
+                self.log(node, f"Documentation mismatch for '{param.name}': wrong type")
+
+            if doc.optional != param.optional:
+                self.log(
+                    node,
+                    f"Documentation mismatch for '{param.name}': wrong optionality",
+                )
+
+
+def check_file(path: Path):
+    source: str = path.read_text()
+    tree = ast.parse(source)
+    checker = Checker()
+    checker.visit(tree)
+
+
+def main():
+    folder: Path = (Path(__file__).parent.parent / "midas").resolve()
+    all_files = folder.rglob("*.py")
+    for f in all_files:
+        print(f.relative_to(folder))
+        check_file(f)
+        print()
+
+
+if __name__ == "__main__":
+    main()

tests/cases/checker/01_simple_types.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 df: Frame[
     verified: bool,
     birth_year: int,
-    height: float + ( _ > 0 ) + ( _ < 250 ),
+    height: float,
     name: str,
     date: datetime,
     float,

tests/cases/checker/01_simple_types.py.ref.json

@@ -1,19 +1,5 @@
 {
   "diagnostics": [
-    {
-      "type": "Warning",
-      "location": {
-        "start": [
-          8,
-          12
-        ],
-        "end": [
-          8,
-          43
-        ]
-      },
-      "message": "ConstraintType not yet supported"
-    },
     {
       "type": "Warning",
       "location": {

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

@@ -5,7 +5,7 @@ from __future__ import annotations
 df: Frame[
     verified: bool,
     birth_year: int,
-    height: float + ( _ > 0 ) + ( _ < 250 ),
+    height: float,
     name: str,
     date: datetime,
     float,

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

@@ -40,13 +40,9 @@
             "_type": "FrameColumn",
             "name": "height",
             "type": {
-              "_type": "ConstraintType",
+              "_type": "BaseType",
-              "type": {
+              "base": "float",
-                "_type": "BaseType",
+              "args": []
-                "base": "float",
-                "args": []
-              },
-              "constraint": "(_ > 0) + (_ < 250)"
             }
           },
           {

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

@@ -16,10 +16,6 @@ lat2: Latitude = lat[1]
 lat_diff: Difference[Latitude] = lat2 - lat1
 
 df2: Frame[
-    age: int + (_ >= 0),
+    age: int,
-    height: float + (_ >= 0),
+    height: float,
-]
-df2_bis: Frame[
-    age: int + Positive,
-    height: float + Positive,
 ]

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

@@ -227,61 +227,18 @@
             "_type": "FrameColumn",
             "name": "age",
             "type": {
-              "_type": "ConstraintType",
+              "_type": "BaseType",
-              "type": {
+              "base": "int",
-                "_type": "BaseType",
+              "args": []
-                "base": "int",
-                "args": []
-              },
-              "constraint": "_ >= 0"
             }
           },
           {
             "_type": "FrameColumn",
             "name": "height",
             "type": {
-              "_type": "ConstraintType",
+              "_type": "BaseType",
-              "type": {
+              "base": "float",
-                "_type": "BaseType",
+              "args": []
-                "base": "float",
-                "args": []
-              },
-              "constraint": "_ >= 0"
-            }
-          }
-        ]
-      }
-    },
-    {
-      "_type": "TypeAssign",
-      "name": "df2_bis",
-      "type": {
-        "_type": "FrameType",
-        "columns": [
-          {
-            "_type": "FrameColumn",
-            "name": "age",
-            "type": {
-              "_type": "ConstraintType",
-              "type": {
-                "_type": "BaseType",
-                "base": "int",
-                "args": []
-              },
-              "constraint": "Positive"
-            }
-          },
-          {
-            "_type": "FrameColumn",
-            "name": "height",
-            "type": {
-              "_type": "ConstraintType",
-              "type": {
-                "_type": "BaseType",
-                "base": "float",
-                "args": []
-              },
-              "constraint": "Positive"
             }
           }
         ]

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

@@ -4,11 +4,10 @@ from __future__ import annotations
 
 
 def func(
-    col1: Column[float + (0 <= _ <= 1)],
+    col1: Column[float],
-    col2: Column[float + (0 <= _ <= 1)],
+    col2: Column[float],
-) -> Column[float + (0 <= _ <= 2)]:
+) -> Column[float]:
-    result: Column[float + (0 <= _ <= 2)] = col1 + col2
+    return col1 + col2
-    return result
 
 
 def func2(a: int, /, b: float, *, c: str):

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

@@ -26,13 +26,9 @@
               "base": "Column",
               "args": [
                 {
-                  "_type": "ConstraintType",
+                  "_type": "BaseType",
-                  "type": {
+                  "base": "float",
-                    "_type": "BaseType",
+                  "args": []
-                    "base": "float",
-                    "args": []
-                  },
-                  "constraint": "0 <= _ <= 1"
                 }
               ]
             },
@@ -45,13 +41,9 @@
               "base": "Column",
               "args": [
                 {
-                  "_type": "ConstraintType",
+                  "_type": "BaseType",
-                  "type": {
+                  "base": "float",
-                    "_type": "BaseType",
+                  "args": []
-                    "base": "float",
-                    "args": []
-                  },
-                  "constraint": "0 <= _ <= 1"
                 }
               ]
             },
@@ -65,44 +57,15 @@
         "base": "Column",
         "args": [
           {
-            "_type": "ConstraintType",
+            "_type": "BaseType",
-            "type": {
+            "base": "float",
-              "_type": "BaseType",
+            "args": []
-              "base": "float",
-              "args": []
-            },
-            "constraint": "0 <= _ <= 2"
           }
         ]
       },
       "body": [
         {
-          "_type": "TypeAssign",
+          "_type": "ReturnStmt",
-          "name": "result",
-          "type": {
-            "_type": "BaseType",
-            "base": "Column",
-            "args": [
-              {
-                "_type": "ConstraintType",
-                "type": {
-                  "_type": "BaseType",
-                  "base": "float",
-                  "args": []
-                },
-                "constraint": "0 <= _ <= 2"
-              }
-            ]
-          }
-        },
-        {
-          "_type": "AssignStmt",
-          "targets": [
-            {
-              "_type": "VariableExpr",
-              "name": "result"
-            }
-          ],
           "value": {
             "_type": "BinaryExpr",
             "left": {
@@ -115,13 +78,6 @@
               "name": "col2"
             }
           }
-        },
-        {
-          "_type": "ReturnStmt",
-          "value": {
-            "_type": "VariableExpr",
-            "name": "result"
-          }
         }
       ]
     },

tests/serializer/python.py

@@ -8,7 +8,6 @@ from midas.ast.python import (
     CallExpr,
     CastExpr,
     CompareExpr,
-    ConstraintType,
     DictExpr,
     Expr,
     ExpressionStmt,
@@ -106,13 +105,6 @@ class PythonAstJsonSerializer(
             "args": self._serialize_list(node.args),
         }
 
-    def visit_constraint_type(self, node: ConstraintType) -> dict:
-        return {
-            "_type": "ConstraintType",
-            "type": node.type.accept(self),
-            "constraint": ast.unparse(node.constraint),
-        }
-
     def visit_frame_column(self, node: FrameColumn) -> dict:
         return {
             "_type": "FrameColumn",

vscode-ext/language-configuration.json

@@ -0,0 +1,83 @@
+{
+    "comments": {
+        "lineComment": {
+            "comment": "//",
+        },
+        "blockComment": [
+            "/*",
+            "*/"
+        ]
+    },
+    "brackets": [
+        [
+            "{",
+            "}"
+        ],
+        [
+            "[",
+            "]"
+        ],
+        [
+            "(",
+            ")"
+        ]
+    ],
+    "autoClosingPairs": [
+        {
+            "open": "{",
+            "close": "}"
+        },
+        {
+            "open": "[",
+            "close": "]"
+        },
+        {
+            "open": "(",
+            "close": ")"
+        },
+        {
+            "open": "'",
+            "close": "'",
+            "notIn": [
+                "string",
+                "comment"
+            ]
+        },
+        {
+            "open": "\"",
+            "close": "\"",
+            "notIn": [
+                "string"
+            ]
+        },
+        {
+            "open": "/**",
+            "close": " */",
+            "notIn": [
+                "string"
+            ]
+        }
+    ],
+    "surroundingPairs": [
+        [
+            "{",
+            "}"
+        ],
+        [
+            "[",
+            "]"
+        ],
+        [
+            "(",
+            ")"
+        ],
+        [
+            "'",
+            "'"
+        ],
+        [
+            "\"",
+            "\""
+        ]
+    ]
+}

vscode-ext/language-configurations.json

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

vscode-ext/package.json

@@ -4,7 +4,9 @@
     "engines": {
         "vscode": "*"
     },
-    "categories": ["Programming Languages"],
+    "categories": [
+        "Programming Languages"
+    ],
     "contributes": {
         "languages": [
             {
@@ -24,6 +26,12 @@
                 "scopeName": "source.midas",
                 "path": "./syntaxes/midas.tmLanguage.json"
             }
+        ],
+        "snippets": [
+            {
+                "language": "midas",
+                "path": "./snippets.json"
+            }
         ]
     }
 }

vscode-ext/snippets.json

@@ -0,0 +1,50 @@
+{
+    "Type alias": {
+        "prefix": "alias",
+        "body": "alias ${1:name} = $0",
+        "description": "Declare a type alias"
+    },
+    "Derived type": {
+        "prefix": "type",
+        "body": "type ${1:name} = $0",
+        "description": "Declare a derived type"
+    },
+    "Predicate": {
+        "prefix": "predicate",
+        "body": "predicate ${1:signature} = $0",
+        "description": "Declare a predicate"
+    },
+    "Extend": {
+        "prefix": "extend",
+        "body": [
+            "extend ${1:type} {",
+            "\t$0",
+            "}"
+        ],
+        "description": "Extend a type to add members"
+    },
+    "Property": {
+        "prefix": "prop",
+        "body": "prop ${1:name}: $0",
+        "description": "Declare a property"
+    },
+    "Method": {
+        "prefix": "def",
+        "body": "def ${1:name}: $0",
+        "description": "Declare a method"
+    },
+    "Function type": {
+        "prefix": "fn",
+        "body": "fn(${1:parameters}) -> ${2:returns}",
+        "description": "A function type"
+    },
+    "Frame type": {
+        "prefix": "frame",
+        "body": [
+            "Frame[",
+            "\t$0",
+            "]"
+        ],
+        "description": "A frame type"
+    }
+}