feat(gen): generate basic constraint assertion

Reviewed-on: #13

gen/midas.py

@@ -26,6 +26,14 @@ class MemberKind(Enum):
     METHOD = auto()
 
 
+@dataclass(frozen=True, kw_only=True)
+class ParamSpec:
+    l_paren: Token
+    pos: list[FunctionType.Argument]
+    mixed: list[FunctionType.Argument]
+    kw: list[FunctionType.Argument]
+
+
 ###<
 
 
@@ -50,9 +58,8 @@ class ExtendStmt:
 
 class PredicateStmt:
     name: Token
-    subject: Token
-    type: Type
-    condition: Expr
+    params: list[ParamSpec]
+    body: Expr
 
 
 ###<
@@ -78,6 +85,12 @@ class UnaryExpr:
     right: Expr
 
 
+class CallExpr:
+    callee: Expr
+    arguments: list[Expr]
+    keywords: dict[str, Expr]
+
+
 class GetExpr:
     expr: Expr
     name: Token
@@ -128,9 +141,7 @@ class ExtensionType:
 
 
 class FunctionType:
-    pos_args: list[Argument]
-    args: list[Argument]
-    kw_args: list[Argument]
+    params: ParamSpec
     returns: Type
 
     @dataclass(frozen=True, kw_only=True)

gen/python.py

@@ -157,6 +157,11 @@ class ListExpr:
     items: list[Expr]
 
 
+class DictExpr:
+    keys: list[Optional[Expr]]
+    values: list[Expr]
+
+
 class SubscriptExpr:
     object: Expr
     index: Expr

midas/ast/midas.py

@@ -27,6 +27,14 @@ class MemberKind(Enum):
     METHOD = auto()
 
 
+@dataclass(frozen=True, kw_only=True)
+class ParamSpec:
+    l_paren: Token
+    pos: list[FunctionType.Argument]
+    mixed: list[FunctionType.Argument]
+    kw: list[FunctionType.Argument]
+
+
 ##############
 # Statements #
 ##############
@@ -86,9 +94,8 @@ class ExtendStmt(Stmt):
 @dataclass(frozen=True)
 class PredicateStmt(Stmt):
     name: Token
-    subject: Token
-    type: Type
-    condition: Expr
+    params: list[ParamSpec]
+    body: Expr
 
     def accept(self, visitor: Stmt.Visitor[T]) -> T:
         return visitor.visit_predicate_stmt(self)
@@ -116,6 +123,9 @@ class Expr(ABC):
         @abstractmethod
         def visit_unary_expr(self, expr: UnaryExpr) -> T: ...
 
+        @abstractmethod
+        def visit_call_expr(self, expr: CallExpr) -> T: ...
+
         @abstractmethod
         def visit_get_expr(self, expr: GetExpr) -> T: ...
 
@@ -161,6 +171,16 @@ class UnaryExpr(Expr):
         return visitor.visit_unary_expr(self)
 
 
+@dataclass(frozen=True)
+class CallExpr(Expr):
+    callee: Expr
+    arguments: list[Expr]
+    keywords: dict[str, Expr]
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_call_expr(self)
+
+
 @dataclass(frozen=True)
 class GetExpr(Expr):
     expr: Expr
@@ -279,9 +299,7 @@ class ExtensionType(Type):
 
 @dataclass(frozen=True)
 class FunctionType(Type):
-    pos_args: list[Argument]
-    args: list[Argument]
-    kw_args: list[Argument]
+    params: ParamSpec
     returns: Type
 
     @dataclass(frozen=True, kw_only=True)

midas/ast/printer.py

@@ -150,13 +150,17 @@ class MidasAstPrinter(
         self._write_line("PredicateStmt")
         with self._child_level():
             self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line(f'subject: "{stmt.subject.lexeme}"')
-            self._write_line("type")
+            self._write_line("params")
+            with self._child_level():
+                for i, spec in enumerate(stmt.params):
+                    self._idx = i
+                    if i == len(stmt.params) - 1:
+                        self._mark_last()
+                    self._visit_param_spec(spec)
+
+            self._write_line("body", last=True)
             with self._child_level(single=True):
-                stmt.type.accept(self)
-            self._write_line("condition", last=True)
-            with self._child_level(single=True):
-                stmt.condition.accept(self)
+                stmt.body.accept(self)
 
     # Expressions
 
@@ -195,6 +199,29 @@ class MidasAstPrinter(
             with self._child_level(single=True):
                 expr.right.accept(self)
 
+    def visit_call_expr(self, expr: m.CallExpr) -> None:
+        self._write_line("CallExpr")
+        with self._child_level():
+            self._write_line("callee")
+            with self._child_level(single=True):
+                expr.callee.accept(self)
+            self._write_line("arguments")
+            with self._child_level():
+                for i, arg in enumerate(expr.arguments):
+                    self._idx = i
+                    if i == len(expr.arguments) - 1:
+                        self._mark_last()
+                    arg.accept(self)
+            self._write_line("keywords", last=True)
+            with self._child_level():
+                for i, (name, arg) in enumerate(expr.keywords.items()):
+                    self._idx = i
+                    if i == len(expr.keywords) - 1:
+                        self._mark_last()
+                    self._write_line(name)
+                    with self._child_level(single=True):
+                        arg.accept(self)
+
     def visit_get_expr(self, expr: m.GetExpr):
         self._write_line("GetExpr")
         with self._child_level():
@@ -276,34 +303,41 @@ class MidasAstPrinter(
     def visit_function_type(self, type: m.FunctionType) -> None:
         self._write_line("FunctionType")
         with self._child_level():
-            self._write_line("pos_args")
-            with self._child_level():
-                for i, arg in enumerate(type.pos_args):
-                    self._idx = i
-                    if i == len(type.pos_args) - 1:
-                        self._mark_last()
-                    self._print_function_arg(arg)
-
-            self._write_line("args")
-            with self._child_level():
-                for i, arg in enumerate(type.args):
-                    self._idx = i
-                    if i == len(type.args) - 1:
-                        self._mark_last()
-                    self._print_function_arg(arg)
-
-            self._write_line("kw_args")
-            with self._child_level():
-                for i, arg in enumerate(type.kw_args):
-                    self._idx = i
-                    if i == len(type.kw_args) - 1:
-                        self._mark_last()
-                    self._print_function_arg(arg)
+            self._write_line("params")
+            with self._child_level(single=True):
+                self._visit_param_spec(type.params)
 
             self._write_line("returns", last=True)
             with self._child_level(single=True):
                 type.returns.accept(self)
 
+    def _visit_param_spec(self, spec: m.ParamSpec) -> None:
+        self._write_line("ParamSpec")
+        with self._child_level():
+            self._write_line("pos")
+            with self._child_level():
+                for i, arg in enumerate(spec.pos):
+                    self._idx = i
+                    if i == len(spec.pos) - 1:
+                        self._mark_last()
+                    self._print_function_arg(arg)
+
+            self._write_line("mixed")
+            with self._child_level():
+                for i, arg in enumerate(spec.mixed):
+                    self._idx = i
+                    if i == len(spec.mixed) - 1:
+                        self._mark_last()
+                    self._print_function_arg(arg)
+
+            self._write_line("kw", last=True)
+            with self._child_level():
+                for i, arg in enumerate(spec.kw):
+                    self._idx = i
+                    if i == len(spec.kw) - 1:
+                        self._mark_last()
+                    self._print_function_arg(arg)
+
     def _print_function_arg(self, arg: m.FunctionType.Argument) -> None:
         self._write_line("Argument")
         with self._child_level():
@@ -367,10 +401,9 @@ class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str]
 
     def visit_predicate_stmt(self, stmt: m.PredicateStmt):
         name: str = stmt.name.lexeme
-        subject: str = stmt.subject.lexeme
-        type: str = stmt.type.accept(self)
-        condition: str = stmt.condition.accept(self)
-        return self.indented(f"predicate {name}({subject}: {type}) = {condition}")
+        sig: str = "".join(self._visit_param_spec(spec) for spec in stmt.params)
+        body: str = stmt.body.accept(self)
+        return self.indented(f"predicate {name}{sig} = {body}")
 
     def visit_logical_expr(self, expr: m.LogicalExpr):
         left: str = expr.left.accept(self)
@@ -389,6 +422,12 @@ class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str]
         right: str = expr.right.accept(self)
         return f"{operator}{right}"
 
+    def visit_call_expr(self, expr: m.CallExpr) -> str:
+        args: list[str] = [arg.accept(self) for arg in expr.arguments] + [
+            f"{name}={arg.accept(self)}" for name, arg in expr.keywords.items()
+        ]
+        return f"{expr.callee.accept(self)}({', '.join(args)})"
+
     def visit_get_expr(self, expr: m.GetExpr):
         expr_: str = expr.expr.accept(self)
         name: str = expr.name.lexeme
@@ -436,9 +475,13 @@ class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str]
         return f"{type.base.accept(self)} & {type.extension.accept(self)}"
 
     def visit_function_type(self, type: m.FunctionType) -> str:
-        pos_args: list[str] = [self._print_arg(arg) for arg in type.pos_args]
-        mixed_args: list[str] = [self._print_arg(arg) for arg in type.args]
-        kw_args: list[str] = [self._print_arg(arg) for arg in type.kw_args]
+        spec: str = self._visit_param_spec(type.params)
+        return f"fn {spec} -> {type.returns.accept(self)}"
+
+    def _visit_param_spec(self, spec: m.ParamSpec) -> str:
+        pos_args: list[str] = [self._print_arg(arg) for arg in spec.pos]
+        mixed_args: list[str] = [self._print_arg(arg) for arg in spec.mixed]
+        kw_args: list[str] = [self._print_arg(arg) for arg in spec.kw]
         args: list[str] = pos_args
 
         if len(pos_args) != 0:
@@ -447,8 +490,7 @@ class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str]
         if len(kw_args) != 0:
             args.append("*")
         args += kw_args
-
-        return f"fn ({', '.join(args)}) -> {type.returns.accept(self)}"
+        return f"({', '.join(args)})"
 
     def _print_arg(self, arg: m.FunctionType.Argument) -> str:
         res: str = ""
@@ -745,6 +787,27 @@ class PythonAstPrinter(
                         self._mark_last()
                     item.accept(self)
 
+    def visit_dict_expr(self, expr: p.DictExpr) -> None:
+        self._write_line("DictExpr")
+        with self._child_level():
+            self._write_line("keys")
+            with self._child_level():
+                for i, key in enumerate(expr.keys):
+                    self._idx = i
+                    if i == len(expr.keys) - 1:
+                        self._mark_last()
+                    if key is None:
+                        self._write_line("None")
+                    else:
+                        key.accept(self)
+            self._write_line("values", last=True)
+            with self._child_level():
+                for i, value in enumerate(expr.values):
+                    self._idx = i
+                    if i == len(expr.values) - 1:
+                        self._mark_last()
+                    value.accept(self)
+
     def visit_subscript_expr(self, expr: p.SubscriptExpr) -> None:
         self._write_line("SubscriptExpr")
         with self._child_level():

midas/ast/python.py

@@ -259,6 +259,9 @@ class Expr(ABC):
         @abstractmethod
         def visit_list_expr(self, expr: ListExpr) -> T: ...
 
+        @abstractmethod
+        def visit_dict_expr(self, expr: DictExpr) -> T: ...
+
         @abstractmethod
         def visit_subscript_expr(self, expr: SubscriptExpr) -> T: ...
 
@@ -370,6 +373,15 @@ class ListExpr(Expr):
         return visitor.visit_list_expr(self)
 
 
+@dataclass(frozen=True)
+class DictExpr(Expr):
+    keys: list[Optional[Expr]]
+    values: list[Expr]
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_dict_expr(self)
+
+
 @dataclass(frozen=True)
 class SubscriptExpr(Expr):
     object: Expr

midas/checker/builtins.midas

@@ -150,3 +150,32 @@ extend list[T] {
 
     prop __doc__: str
 }
+
+extend dict[K, V] {
+    def copy: fn() -> dict[K, V]
+    def keys: fn() -> list[K]  // TODO: use builtin types
+    def values: fn() -> list[V]  // TODO: use builtin types
+    // def items: fn() -> list[tuple[K, V]]  // TODO: use builtin types
+    
+    // def get: fn(key: K, default: None = None, /) -> V | None
+    def get: fn(key: K, default: V, /) -> V
+    // def get: fn[T](key: K, default: T, /) -> V | T
+    def pop: fn(key: K, /) -> V
+    def pop: fn(key: K, default: V, /) -> V
+    // def pop: fn[T](key: K, default: T, /) -> V | T
+    def __len__: fn() -> int
+    def __getitem__: fn(key: K, /) -> V
+    def __setitem__: fn(key: K, value: V, /) -> None
+    def __delitem__: fn(key: K, /) -> None
+    // def __iter__: fn() -> Iterator[K]
+    def __eq__: fn(value: object, /) -> bool
+    // def __reversed__: fn() -> Iterator[K]
+    
+    def __or__: fn(value: dict[K, V], /) -> dict[K, V]
+    // def __or__: fn[K2, V2](value: dict[K2, V2], /) -> dict[K | K2, V | V2]
+    def __ror__: fn(value: dict[K, V], /) -> dict[K, V]
+    // def __ror__: fn[K2, V2](value: dict[K2, V2], /) -> dict[K | K2, V | V2]
+    // def __ior__: fn(value: SupportsKeysAndGetItem[K, V], /) -> dict[K, V]
+    // def __ior__: fn(value: Iterable[tuple[K, V]], /) -> dict[K, V]
+
+}

midas/checker/builtins.py

@@ -39,3 +39,14 @@ def define_builtins(reg: TypesRegistry):
             body=BaseType(name="list"),
         ),
     )
+    dict = reg.define_type(
+        "dict",
+        GenericType(
+            name="dict",
+            params=[
+                TypeVar(name="K", bound=None),
+                TypeVar(name="V", bound=None),
+            ],
+            body=BaseType(name="dict"),
+        ),
+    )

midas/checker/midas.py

@@ -1,4 +1,5 @@
 import logging
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
@@ -9,9 +10,11 @@ from midas.checker.reporter import FileReporter, Reporter
 from midas.checker.types import (
     AliasType,
     ComplexType,
+    ConstraintType,
     ExtensionType,
     Function,
     GenericType,
+    Predicate,
     Type,
     TypeVar,
     UnknownType,
@@ -21,6 +24,13 @@ from midas.lexer.token import Token
 from midas.parser.midas import MidasParser
 
 
+@dataclass(frozen=True, kw_only=True)
+class TypedParamSpec:
+    pos: list[Function.Argument]
+    mixed: list[Function.Argument]
+    kw: list[Function.Argument]
+
+
 class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[Type]):
     """A resolver which evaluates Midas type definitions and build a registry"""
 
@@ -37,6 +47,8 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[Type
         builtins_path: Path = (Path(__file__).parent / "builtins.midas").resolve()
         self.process(builtins_path.read_text(), str(builtins_path))
 
+        self._bool: Type = self.get_type("bool")
+
     def process(self, source: str, path: Optional[str]):
         self.reporter = self.reporter.for_file(path)
         lexer: MidasLexer = MidasLexer(source)
@@ -106,7 +118,24 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[Type
             )
 
     def visit_predicate_stmt(self, stmt: m.PredicateStmt) -> None:
-        self.reporter.warning(stmt.location, "PredicateStmt not yet supported")
+        params: list[TypedParamSpec] = [
+            self._visit_param_spec(spec) for spec in stmt.params
+        ]
+        type: Type = self._bool
+        for spec in reversed(params):
+            type = Function(
+                pos_args=spec.pos,
+                args=spec.mixed,
+                kw_args=spec.kw,
+                returns=type,
+            )
+        self.types.define_predicate(
+            stmt.name.lexeme,
+            Predicate(
+                type=type,
+                body=stmt.body,
+            ),
+        )
 
     def visit_logical_expr(self, expr: m.LogicalExpr) -> None:
         self.reporter.warning(expr.location, "LogicalExpr not yet supported")
@@ -117,6 +146,9 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[Type
     def visit_unary_expr(self, expr: m.UnaryExpr) -> None:
         self.reporter.warning(expr.location, "UnaryExpr not yet supported")
 
+    def visit_call_expr(self, expr: m.CallExpr) -> None:
+        self.reporter.warning(expr.location, "CallExpr not yet supported")
+
     def visit_get_expr(self, expr: m.GetExpr) -> None:
         self.reporter.warning(expr.location, "GetExpr not yet supported")
 
@@ -153,10 +185,10 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[Type
             return UnknownType()
 
     def visit_constraint_type(self, type: m.ConstraintType) -> Type:
-        type_: Type = type.type.accept(self)
-        type.constraint.accept(self)
-        # TODO
-        return UnknownType()
+        return ConstraintType(
+            type=type.type.accept(self),
+            constraint=type.constraint,
+        )
 
     def visit_complex_type(self, type: m.ComplexType) -> ComplexType:
         return ComplexType(
@@ -172,8 +204,17 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[Type
         )
 
     def visit_function_type(self, type: m.FunctionType) -> Type:
-        n_pos_args: int = len(type.pos_args)
-        n_args: int = len(type.args)
+        params: TypedParamSpec = self._visit_param_spec(type.params)
+        return Function(
+            pos_args=params.pos,
+            args=params.mixed,
+            kw_args=params.kw,
+            returns=type.returns.accept(self),
+        )
+
+    def _visit_param_spec(self, spec: m.ParamSpec) -> TypedParamSpec:
+        n_pos: int = len(spec.pos)
+        n_mixed: int = len(spec.mixed)
 
         def process_arg(arg: m.FunctionType.Argument, i: int) -> Function.Argument:
             return Function.Argument(
@@ -183,14 +224,10 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[None], m.Type.Visitor[Type
                 required=arg.required,
             )
 
-        return Function(
-            pos_args=[process_arg(arg, i) for i, arg in enumerate(type.pos_args)],
-            args=[process_arg(arg, i + n_pos_args) for i, arg in enumerate(type.args)],
-            kw_args=[
-                process_arg(arg, i + n_pos_args + n_args)
-                for i, arg in enumerate(type.kw_args)
-            ],
-            returns=type.returns.accept(self),
+        return TypedParamSpec(
+            pos=[process_arg(arg, i) for i, arg in enumerate(spec.pos)],
+            mixed=[process_arg(arg, i + n_pos) for i, arg in enumerate(spec.mixed)],
+            kw=[process_arg(arg, i + n_pos + n_mixed) for i, arg in enumerate(spec.kw)],
         )
 
     def _resolve_type_params(self, params: list[m.TypeParam]):

midas/checker/preamble.py

@@ -61,7 +61,7 @@ class Preamble(Environment):
         # TODO: more specific arg types
         self._def_function(
             name=name,
-            pos=[Param("object", TopType())],
+            pos=[Param("object", TopType(), required=False)],
             returns=self._types.get_type(name),
         )
 

midas/checker/python.py

@@ -552,6 +552,46 @@ class PythonTyper(
         )
         return self.types.apply_generic(list_type, [UnknownType()])
 
+    def visit_dict_expr(self, expr: p.DictExpr) -> Type:
+        dict_type: Type = self.types.get_type("dict")
+
+        key_types: list[Type] = []
+        value_types: list[Type] = []
+        for key, value in zip(expr.keys, expr.values):
+            if key is None:
+                self.reporter.warning(
+                    value.location, "Dictionary unpacking not supported"
+                )
+                continue
+            key_types.append(self.type_of(key))
+            value_types.append(self.type_of(value))
+
+        key_types = self.types.reduce_types(key_types)
+        value_types = self.types.reduce_types(value_types)
+
+        if len(key_types) == 0 or len(value_types) == 0:
+            return dict_type
+
+        key_type: Type = UnknownType()
+        value_type: Type = UnknownType()
+
+        if len(key_types) == 1:
+            key_type = key_types[0]
+        else:
+            self.reporter.error(
+                expr.location,
+                f"Heterogeneous dict keys: {key_types}",
+            )
+
+        if len(value_types) == 1:
+            value_type = value_types[0]
+        else:
+            self.reporter.error(
+                expr.location,
+                f"Heterogeneous dict values: {value_types}",
+            )
+        return self.types.apply_generic(dict_type, [key_type, value_type])
+
     def visit_subscript_expr(self, expr: p.SubscriptExpr) -> Type:
         object: Type = self.type_of(expr.object)
         operation: Optional[Type] = self.types.lookup_member(object, "__getitem__")

midas/checker/registry.py

@@ -11,6 +11,7 @@ from midas.checker.types import (
     Function,
     GenericType,
     OverloadedFunction,
+    Predicate,
     TopType,
     Type,
     TypeVar,
@@ -24,6 +25,7 @@ class TypesRegistry:
         self.logger: logging.Logger = logging.getLogger("TypesRegistry")
         self._types: dict[str, Type] = {}
         self._members: dict[str, dict[str, Type]] = {}
+        self._predicates: dict[str, Predicate] = {}
 
     def get_type(self, name: str) -> Type:
         """Get a type from its name
@@ -81,6 +83,11 @@ class TypesRegistry:
         else:
             members[member_name] = member_type
 
+    def define_predicate(self, name: str, predicate: Predicate):
+        if name in self._predicates:
+            raise ValueError(f"Predicate {name} already defined")
+        self._predicates[name] = predicate
+
     def is_subtype(self, type1: Type, type2: Type) -> bool:
         """Check whether `type1` is a subtype of `type2`
 

midas/checker/resolver.py

@@ -213,6 +213,13 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
         for item in expr.items:
             self.resolve(item)
 
+    def visit_dict_expr(self, expr: p.DictExpr) -> None:
+        for key in expr.keys:
+            if key is not None:
+                self.resolve(key)
+        for value in expr.values:
+            self.resolve(value)
+
     def visit_subscript_expr(self, expr: p.SubscriptExpr) -> None:
         self.resolve(expr.object)
         self.resolve(expr.index)

midas/checker/types.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Optional, assert_never
+
+import midas.ast.midas as m
+from midas.ast.printer import MidasPrinter
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -130,6 +133,16 @@ class AppliedType:
         return f"{self.name}[{', '.join(map(str, self.args))}]"
 
 
+@dataclass(frozen=True, kw_only=True)
+class ConstraintType:
+    type: Type
+    constraint: m.Expr
+
+    def __str__(self) -> str:
+        printer = MidasPrinter()
+        return f"{self.type} where {printer.print(self.constraint)}"
+
+
 def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
     def sub_argument(arg: Function.Argument):
         return Function.Argument(
@@ -195,6 +208,12 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
                 body=substitute_typevars(body, substitutions),
             )
 
+        case ConstraintType():
+            return ConstraintType(
+                type=substitute_typevars(type.type, substitutions),
+                constraint=type.constraint,
+            )
+
         case TypeVar(name=name):
             if name in substitutions:
                 return substitutions[name]
@@ -203,9 +222,13 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
         case UnknownType() | UnitType():
             return type
 
-        case _:
+        case TopType() | GenericType():
             raise NotImplementedError(f"Unsupported type {type}")
 
+        # Ensure exhaustiveness
+        case _:
+            assert_never(type)
+
 
 def unfold_type(type: Type) -> Type:
     match type:
@@ -215,6 +238,12 @@ def unfold_type(type: Type) -> Type:
             return type
 
 
+@dataclass(frozen=True, kw_only=True)
+class Predicate:
+    type: Type
+    body: m.Expr
+
+
 Type = (
     TopType
     | BaseType
@@ -228,4 +257,5 @@ Type = (
     | TypeVar
     | GenericType
     | AppliedType
+    | ConstraintType
 )

midas/generator/constraints.py

@@ -0,0 +1,92 @@
+import ast
+
+import midas.ast.midas as m
+from midas.lexer.token import TokenType
+
+LOGICAL_OPERATORS: dict[TokenType, type[ast.boolop]] = {
+    TokenType.AND: ast.And,
+    # TokenType.OR: ast.Or,
+}
+
+BINARY_OPERATORS: dict[TokenType, type[ast.operator]] = {
+    # 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.MINUS: ast.USub,
+}
+
+COMPARISON_OPERATORS: dict[TokenType, type[ast.cmpop]] = {
+    TokenType.GREATER: ast.Gt,
+    TokenType.GREATER_EQUAL: ast.GtE,
+    TokenType.LESS: ast.Lt,
+    TokenType.LESS_EQUAL: ast.LtE,
+    TokenType.EQUAL_EQUAL: ast.Eq,
+    TokenType.BANG_EQUAL: ast.NotEq,
+}
+
+
+class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
+    def visit_logical_expr(self, expr: m.LogicalExpr) -> ast.expr:
+        return ast.BoolOp(
+            op=LOGICAL_OPERATORS[expr.operator.type](),
+            values=[
+                expr.left.accept(self),
+                expr.right.accept(self),
+            ],
+        )
+
+    def visit_binary_expr(self, expr: m.BinaryExpr) -> ast.expr:
+        op: TokenType = expr.operator.type
+        if op in BINARY_OPERATORS:
+            return ast.BinOp(
+                left=expr.left.accept(self),
+                op=BINARY_OPERATORS[op](),
+                right=expr.right.accept(self),
+            )
+        if op in COMPARISON_OPERATORS:
+            return ast.Compare(
+                left=expr.left.accept(self),
+                ops=[COMPARISON_OPERATORS[op]()],
+                comparators=[expr.right.accept(self)],
+            )
+        raise ValueError(f"Unexpected binary operator {op}")
+
+    def visit_unary_expr(self, expr: m.UnaryExpr) -> ast.expr:
+        return ast.UnaryOp(
+            op=UNARY_OPERATORS[expr.operator.type](),
+            operand=expr.right.accept(self),
+        )
+
+    def visit_call_expr(self, expr: m.CallExpr) -> ast.expr:
+        return ast.Call(
+            func=expr.callee.accept(self),
+            args=[arg.accept(self) for arg in expr.arguments],
+            keywords=[
+                ast.keyword(arg=name, value=arg.accept(self))
+                for name, arg in expr.keywords.items()
+            ],
+        )
+
+    def visit_get_expr(self, expr: m.GetExpr) -> ast.expr:
+        return ast.Attribute(
+            value=expr.expr.accept(self),
+            attr=expr.name.lexeme,
+        )
+
+    def visit_variable_expr(self, expr: m.VariableExpr) -> ast.expr:
+        # TODO: lookup predicate
+        return ast.Name(id=expr.name.lexeme)
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr) -> ast.expr:
+        return expr.accept(self)
+
+    def visit_literal_expr(self, expr: m.LiteralExpr) -> ast.expr:
+        return ast.Constant(value=expr.value)
+
+    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> ast.expr:
+        return ast.Name(id="_")

midas/generator/generator.py

@@ -2,15 +2,18 @@ import ast
 import shutil
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional
+from typing import Optional, assert_never
 
-from midas.ast.location import Location
+import midas.ast.midas as m
 import midas.ast.python as p
+from midas.ast.location import Location
+from midas.ast.printer import MidasPrinter
 from midas.checker.types import (
     AliasType,
     AppliedType,
     BaseType,
     ComplexType,
+    ConstraintType,
     ExtensionType,
     Function,
     GenericType,
@@ -19,7 +22,9 @@ from midas.checker.types import (
     Type,
     TypeVar,
     UnitType,
+    UnknownType,
 )
+from midas.generator.constraints import ConstraintGenerator
 from midas.utils import TypedAST
 
 
@@ -45,6 +50,8 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         self._alias_count: int = 0
         self._scopes: list[Scope] = []
 
+        self._constraint_generator: ConstraintGenerator = ConstraintGenerator()
+
     def generate_ast(self, typed_ast: TypedAST, src_path: Path) -> ast.AST:
         self.rel_src_path = src_path.relative_to(self.workdir)
         self._typed_ast = typed_ast
@@ -139,6 +146,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
             elts=[item.accept(self) for item in expr.items],
         )
 
+    def visit_dict_expr(self, expr: p.DictExpr) -> ast.expr:
+        return ast.Dict(
+            keys=[key.accept(self) if key is not None else None for key in expr.keys],
+            values=[value.accept(self) for value in expr.values],
+        )
+
     def visit_subscript_expr(self, expr: p.SubscriptExpr) -> ast.expr:
         return ast.Subscript(
             value=expr.object.accept(self),
@@ -270,6 +283,9 @@ 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):
         match type:
+            case UnknownType():
+                pass
+
             case BaseType(name=name):
                 self._add_assert(
                     ast.Call(
@@ -295,8 +311,15 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                     self._make_cast_assert_message(src_location, expr, type),
                 )
 
-            case AppliedType():
-                self._make_cast_asserts(src_location, expr, type.body)
+            case AppliedType(body=body):
+                self._make_cast_asserts(src_location, expr, body)
+
+            case ConstraintType(type=base, constraint=constraint):
+                self._make_cast_asserts(src_location, expr, base)
+                self._make_constraint_assert(src_location, expr, constraint)
+
+            case TypeVar():
+                raise RuntimeError("Unexpected TypeVar")
 
             case (
                 TopType()
@@ -308,8 +331,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
             ):
                 raise NotImplementedError(f"Can't make assertion for type {type}")
 
-            case TypeVar():
-                raise RuntimeError("Unexpected TypeVar")
+            # Ensure exhaustiveness
+            case _:
+                assert_never(type)
 
     def _make_cast_assert_message(
         self, location: Location, expr: ast.expr, type: Type
@@ -333,3 +357,23 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 ast.Constant(f" to {type}"),
             ]
         )
+
+    def _make_constraint_assert(
+        self, src_location: Location, expr: ast.expr, constraint: m.Expr
+    ):
+        test: ast.expr = constraint.accept(self._constraint_generator)
+        self._add_assert(
+            test,
+            self._make_constraint_assert_message(src_location, expr, constraint),
+        )
+
+    def _make_constraint_assert_message(
+        self, location: Location, expr: ast.expr, constraint: m.Expr
+    ) -> ast.expr:
+        printer = MidasPrinter()
+        constraint_str: str = printer.print(constraint)
+        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}'"
+        )

midas/parser/midas.py

@@ -3,6 +3,7 @@ from typing import Optional
 from midas.ast.location import Location
 from midas.ast.midas import (
     BinaryExpr,
+    CallExpr,
     ComplexType,
     ConstraintType,
     Expr,
@@ -17,6 +18,7 @@ from midas.ast.midas import (
     MemberKind,
     MemberStmt,
     NamedType,
+    ParamSpec,
     PredicateStmt,
     Stmt,
     Type,
@@ -265,6 +267,9 @@ class MidasParser(Parser):
         Returns:
             Expr: the parsed constraint expression
         """
+        return self.expression()
+
+    def expression(self) -> Expr:
         return self.and_()
 
     def and_(self) -> Expr:
@@ -331,7 +336,55 @@ class MidasParser(Parser):
             right: Expr = self.unary()
             location: Location = Location.span(operator.get_location(), right.location)
             return UnaryExpr(location=location, operator=operator, right=right)
-        return self.reference()
+        return self.call()
+
+    def call(self) -> Expr:
+        expr: Expr = self.reference()
+        if self.match(TokenType.LEFT_PAREN):
+            expr = self.finish_call(expr)
+        return expr
+
+    def finish_call(self, callee: Expr) -> Expr:
+        l_paren: Token = self.previous()
+        pos_args: list[Expr] = []
+        kw_args: dict[str, Expr] = {}
+        keywords: bool = False
+        while not self.match(TokenType.RIGHT_PAREN):
+            if self.check_identifier() and self.check_next(TokenType.EQUAL):
+                keywords = True
+                keyword: Token = self.advance()
+                value: Expr = self.expression()
+                name: str = keyword.lexeme
+                if name in kw_args:
+                    self.error(
+                        self.peek(),
+                        f"Multiple values passed for '{name}', only the last occurrence will be used",
+                    )
+                kw_args[name] = value
+            else:
+                value = self.expression()
+                if self.check(TokenType.EQUAL):
+                    if keywords:
+                        raise self.error(self.peek(), "Invalid keyword argument name")
+                    else:
+                        raise self.error(
+                            self.peek(),
+                            "Cannot pass positional arguments after a keyword argument",
+                        )
+                pos_args.append(value)
+
+            if not self.match(TokenType.COMMA):
+                break
+
+        r_paren: Token = self.consume(
+            TokenType.RIGHT_PAREN, "Expected ')' after arguments."
+        )
+        return CallExpr(
+            location=l_paren.location_to(r_paren),
+            callee=callee,
+            arguments=pos_args,
+            keywords=kw_args,
+        )
 
     def reference(self) -> Expr:
         """Parse an attribute access expression or a simpler expression
@@ -453,23 +506,35 @@ class MidasParser(Parser):
             PredicateStmt: the parsed predicate declaration statement
         """
         keyword: Token = self.previous()
+
         name: Token = self.consume_identifier("Expected predicate name")
-        self.consume(TokenType.LEFT_PAREN, "Expected '(' before predicate subject")
-        subject: Token = self.consume_identifier("Expected subject name")
-        self.consume(TokenType.COLON, "Expected ':' after subject name")
-        type: Type = self.type_expr()
-        self.consume(TokenType.RIGHT_PAREN, "Expected ')' after predicate subject")
+
+        params: list[ParamSpec] = []
+        while self.check(TokenType.LEFT_PAREN):
+            params.append(self.function_args())
+
         self.consume(TokenType.EQUAL, "Expected '=' after predicate subject")
-        condition: Expr = self.constraint()
+        body: Expr = self.constraint()
         return PredicateStmt(
             location=keyword.location_to(self.previous()),
             name=name,
-            subject=subject,
-            type=type,
-            condition=condition,
+            params=params,
+            body=body,
         )
 
     def function(self) -> FunctionType:
+        params: ParamSpec = self.function_args()
+
+        self.consume(TokenType.ARROW, "Expected '->' before result type")
+        result: Type = self.type_expr()
+
+        return FunctionType(
+            location=params.l_paren.location_to(self.previous()),
+            params=params,
+            returns=result,
+        )
+
+    def function_args(self) -> ParamSpec:
         l_paren: Token = self.consume(
             TokenType.LEFT_PAREN, "Expected '(' before function parameters"
         )
@@ -526,14 +591,4 @@ class MidasParser(Parser):
                 self.error(token, "Unnamed mixed argument")
 
         self.consume(TokenType.RIGHT_PAREN, "Expected ')' after function parameters")
-
-        self.consume(TokenType.ARROW, "Expected '->' before result type")
-        result: Type = self.type_expr()
-
-        return FunctionType(
-            location=l_paren.location_to(self.previous()),
-            pos_args=pos_args,
-            args=args,
-            kw_args=kw_args,
-            returns=result,
-        )
+        return ParamSpec(l_paren=l_paren, pos=pos_args, mixed=args, kw=kw_args)

midas/parser/python.py

@@ -10,6 +10,7 @@ from midas.ast.python import (
     CastExpr,
     CompareExpr,
     ConstraintType,
+    DictExpr,
     Expr,
     ExpressionStmt,
     ForStmt,
@@ -447,6 +448,16 @@ class PythonParser:
                     items=[self.parse_expr(item) for item in items],
                 )
 
+            case ast.Dict(keys=keys, values=values):
+                return DictExpr(
+                    location=location,
+                    keys=[
+                        self.parse_expr(key) if key is not None else None
+                        for key in keys
+                    ],
+                    values=[self.parse_expr(value) for value in values],
+                )
+
             case ast.Subscript(value=value, slice=index):
                 return SubscriptExpr(
                     location=location,

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

@@ -2582,18 +2582,21 @@
           "name": "__sub__",
           "type": {
             "_type": "FunctionType",
-            "pos_args": [
-              {
-                "name": null,
-                "type": {
-                  "_type": "NamedType",
-                  "name": "GeoLocation"
-                },
-                "required": true
-              }
-            ],
-            "args": [],
-            "kw_args": [],
+            "params": {
+              "_type": "ParamSpec",
+              "pos": [
+                {
+                  "name": null,
+                  "type": {
+                    "_type": "NamedType",
+                    "name": "GeoLocation"
+                  },
+                  "required": true
+                }
+              ],
+              "mixed": [],
+              "kw": []
+            },
             "returns": {
               "_type": "GenericType",
               "type": {
@@ -2673,18 +2676,21 @@
           "name": "__sub__",
           "type": {
             "_type": "FunctionType",
-            "pos_args": [
-              {
-                "name": null,
-                "type": {
-                  "_type": "NamedType",
-                  "name": "Latitude"
-                },
-                "required": true
-              }
-            ],
-            "args": [],
-            "kw_args": [],
+            "params": {
+              "_type": "ParamSpec",
+              "pos": [
+                {
+                  "name": null,
+                  "type": {
+                    "_type": "NamedType",
+                    "name": "Latitude"
+                  },
+                  "required": true
+                }
+              ],
+              "mixed": [],
+              "kw": []
+            },
             "returns": {
               "_type": "GenericType",
               "type": {
@@ -2713,18 +2719,21 @@
           "name": "__sub__",
           "type": {
             "_type": "FunctionType",
-            "pos_args": [
-              {
-                "name": null,
-                "type": {
-                  "_type": "NamedType",
-                  "name": "Longitude"
-                },
-                "required": true
-              }
-            ],
-            "args": [],
-            "kw_args": [],
+            "params": {
+              "_type": "ParamSpec",
+              "pos": [
+                {
+                  "name": null,
+                  "type": {
+                    "_type": "NamedType",
+                    "name": "Longitude"
+                  },
+                  "required": true
+                }
+              ],
+              "mixed": [],
+              "kw": []
+            },
             "returns": {
               "_type": "GenericType",
               "type": {
@@ -2745,12 +2754,24 @@
     {
       "_type": "PredicateStmt",
       "name": "Positive",
-      "subject": "v",
-      "type": {
-        "_type": "NamedType",
-        "name": "float"
-      },
-      "condition": {
+      "params": [
+        {
+          "_type": "ParamSpec",
+          "pos": [],
+          "mixed": [
+            {
+              "name": "v",
+              "type": {
+                "_type": "NamedType",
+                "name": "float"
+              },
+              "required": true
+            }
+          ],
+          "kw": []
+        }
+      ],
+      "body": {
         "_type": "BinaryExpr",
         "left": {
           "_type": "VariableExpr",
@@ -2766,12 +2787,24 @@
     {
       "_type": "PredicateStmt",
       "name": "StrictlyPositive",
-      "subject": "v",
-      "type": {
-        "_type": "NamedType",
-        "name": "float"
-      },
-      "condition": {
+      "params": [
+        {
+          "_type": "ParamSpec",
+          "pos": [],
+          "mixed": [
+            {
+              "name": "v",
+              "type": {
+                "_type": "NamedType",
+                "name": "float"
+              },
+              "required": true
+            }
+          ],
+          "kw": []
+        }
+      ],
+      "body": {
         "_type": "BinaryExpr",
         "left": {
           "_type": "VariableExpr",
@@ -2787,12 +2820,24 @@
     {
       "_type": "PredicateStmt",
       "name": "Equatorial",
-      "subject": "loc",
-      "type": {
-        "_type": "NamedType",
-        "name": "GeoLocation"
-      },
-      "condition": {
+      "params": [
+        {
+          "_type": "ParamSpec",
+          "pos": [],
+          "mixed": [
+            {
+              "name": "loc",
+              "type": {
+                "_type": "NamedType",
+                "name": "GeoLocation"
+              },
+              "required": true
+            }
+          ],
+          "kw": []
+        }
+      ],
+      "body": {
         "_type": "GroupingExpr",
         "expr": {
           "_type": "BinaryExpr",
@@ -2827,12 +2872,24 @@
     {
       "_type": "PredicateStmt",
       "name": "Arctic",
-      "subject": "loc",
-      "type": {
-        "_type": "NamedType",
-        "name": "GeoLocation"
-      },
-      "condition": {
+      "params": [
+        {
+          "_type": "ParamSpec",
+          "pos": [],
+          "mixed": [
+            {
+              "name": "loc",
+              "type": {
+                "_type": "NamedType",
+                "name": "GeoLocation"
+              },
+              "required": true
+            }
+          ],
+          "kw": []
+        }
+      ],
+      "body": {
         "_type": "GroupingExpr",
         "expr": {
           "_type": "BinaryExpr",

tests/serializer/midas.py

@@ -2,6 +2,7 @@ from typing import Optional, Sequence
 
 from midas.ast.midas import (
     BinaryExpr,
+    CallExpr,
     ComplexType,
     ConstraintType,
     Expr,
@@ -15,6 +16,7 @@ from midas.ast.midas import (
     LogicalExpr,
     MemberStmt,
     NamedType,
+    ParamSpec,
     PredicateStmt,
     Stmt,
     Type,
@@ -78,9 +80,8 @@ class MidasAstJsonSerializer(
         return {
             "_type": "PredicateStmt",
             "name": stmt.name.lexeme,
-            "subject": stmt.subject.lexeme,
-            "type": stmt.type.accept(self),
-            "condition": stmt.condition.accept(self),
+            "params": [self._serialize_param_spec(spec) for spec in stmt.params],
+            "body": stmt.body.accept(self),
         }
 
     def visit_logical_expr(self, expr: LogicalExpr) -> dict:
@@ -106,6 +107,14 @@ class MidasAstJsonSerializer(
             "right": expr.right.accept(self),
         }
 
+    def visit_call_expr(self, expr: CallExpr) -> dict:
+        return {
+            "_type": "CallExpr",
+            "callee": expr.callee.accept(self),
+            "arguments": self._serialize_list(expr.arguments),
+            "keywords": {name: arg.accept(self) for name, arg in expr.keywords.items()},
+        }
+
     def visit_get_expr(self, expr: GetExpr) -> dict:
         return {
             "_type": "GetExpr",
@@ -163,15 +172,21 @@ class MidasAstJsonSerializer(
     def visit_function_type(self, type: FunctionType) -> dict:
         return {
             "_type": "FunctionType",
-            "pos_args": [self._serialize_func_arg(arg) for arg in type.pos_args],
-            "args": [self._serialize_func_arg(arg) for arg in type.args],
-            "kw_args": [self._serialize_func_arg(arg) for arg in type.kw_args],
+            "params": self._serialize_param_spec(type.params),
             "returns": type.returns.accept(self),
         }
 
+    def _serialize_param_spec(self, spec: ParamSpec) -> dict:
+        return {
+            "_type": "ParamSpec",
+            "pos": [self._serialize_func_arg(arg) for arg in spec.pos],
+            "mixed": [self._serialize_func_arg(arg) for arg in spec.mixed],
+            "kw": [self._serialize_func_arg(arg) for arg in spec.kw],
+        }
+
     def _serialize_func_arg(self, arg: FunctionType.Argument) -> dict:
         return {
-            "name": arg.name,
+            "name": arg.name.lexeme if arg.name is not None else None,
             "type": arg.type.accept(self),
             "required": arg.required,
         }

tests/serializer/python.py

@@ -9,6 +9,7 @@ from midas.ast.python import (
     CastExpr,
     CompareExpr,
     ConstraintType,
+    DictExpr,
     Expr,
     ExpressionStmt,
     ForStmt,
@@ -278,6 +279,13 @@ class PythonAstJsonSerializer(
             "items": [item.accept(self) for item in expr.items],
         }
 
+    def visit_dict_expr(self, expr: DictExpr) -> dict:
+        return {
+            "_type": "DictExpr",
+            "keys": [self._serialize_optional(key) for key in expr.keys],
+            "values": self._serialize_list(expr.values),
+        }
+
     def visit_subscript_expr(self, expr: SubscriptExpr) -> dict:
         return {
             "_type": "SubscriptExpr",