refactor: rebrand function parameters and unify spec

rename function arguments to parameters where it was wrong, and add ParamSpec for Python AST, like for Midas

gen/gen.py

@@ -1,3 +1,9 @@
+"""
+Helper script to generate AST nodes for Midas and Python.
+
+Takes in simple templates and generates full dataclasses and a visitor interface
+"""
+
 import re
 from pathlib import Path
 

gen/midas.py

@@ -29,9 +29,9 @@ class MemberKind(Enum):
 @dataclass(frozen=True, kw_only=True)
 class ParamSpec:
     l_paren: Token
-    pos: list[FunctionType.Argument]
+    pos: list[FunctionType.Parameter]
-    mixed: list[FunctionType.Argument]
+    mixed: list[FunctionType.Parameter]
-    kw: list[FunctionType.Argument]
+    kw: list[FunctionType.Parameter]
 
 
 ###<
@@ -150,7 +150,7 @@ class FunctionType:
     returns: Type
 
     @dataclass(frozen=True, kw_only=True)
-    class Argument:
+    class Parameter:
         location: Optional[Location] = None
         name: Optional[Token]
         type: Type

gen/python.py

@@ -12,6 +12,21 @@ from midas.ast.location import Location
 ###<
 
 
+###> Preamble
+@dataclass(frozen=True, kw_only=True)
+class ParamSpec:
+    pos: list[Function.Parameter]
+    mixed: list[Function.Parameter]
+    kw: list[Function.Parameter]
+
+    @property
+    def all(self) -> list[Function.Parameter]:
+        return self.pos + self.mixed + self.kw
+
+
+###<
+
+
 ###> MidasType | Type annotations | node
 class BaseType:
     base: str
@@ -42,25 +57,17 @@ class ExpressionStmt:
 
 class Function:
     name: str
-    posonlyargs: list[Argument]
+    params: ParamSpec
-    args: list[Argument]
-    sink: Optional[Argument]
-    kwonlyargs: list[Argument]
-    kw_sink: Optional[Argument]
     returns: Optional[MidasType]
     body: list[Stmt]
 
     @dataclass(frozen=True, kw_only=True)
-    class Argument:
+    class Parameter:
         location: Optional[Location] = None
         name: str
         type: Optional[MidasType]
         default: Optional[Expr]
 
-    @property
-    def all_args(self) -> list[Argument]:
-        return self.posonlyargs + self.args + self.kwonlyargs
-
 
 class TypeAssign:
     name: str

midas/ast/location.py

@@ -13,6 +13,8 @@ class HasLocation(Protocol):
 
 @dataclass(frozen=True, kw_only=True)
 class Location:
+    """Information about the location of an AST node"""
+
     lineno: int
     col_offset: int
     end_lineno: Optional[int]
@@ -29,6 +31,16 @@ class Location:
 
     @staticmethod
     def span(start: Location, end: Location) -> Location:
+        """Create a new location spanning from one location to another
+
+        Args:
+            start (Location): the starting location
+            end (Location): the end location
+
+        Returns:
+            Location: a new location spanning from the start of `start`
+                to the end of `end`
+        """
         return Location(
             lineno=start.lineno,
             col_offset=start.col_offset,

midas/ast/midas.py

@@ -30,9 +30,9 @@ class MemberKind(Enum):
 @dataclass(frozen=True, kw_only=True)
 class ParamSpec:
     l_paren: Token
-    pos: list[FunctionType.Argument]
+    pos: list[FunctionType.Parameter]
-    mixed: list[FunctionType.Argument]
+    mixed: list[FunctionType.Parameter]
-    kw: list[FunctionType.Argument]
+    kw: list[FunctionType.Parameter]
 
 
 ##############
@@ -318,7 +318,7 @@ class FunctionType(Type):
     returns: Type
 
     @dataclass(frozen=True, kw_only=True)
-    class Argument:
+    class Parameter:
         location: Optional[Location] = None
         name: Optional[Token]
         type: Type

midas/ast/printer.py

@@ -1,896 +0,0 @@
-from __future__ import annotations
-
-import ast
-import io
-from contextlib import contextmanager
-from enum import Enum, auto
-from typing import Generator, Generic, Optional, Protocol, TypeVar
-
-import midas.ast.midas as m
-import midas.ast.python as p
-
-
-class _Level(Enum):
-    EMPTY = auto()
-    ACTIVE = auto()
-    LAST = auto()
-
-
-class Expr(Protocol):
-    def accept(self, printer: AstPrinter) -> None: ...
-
-
-T = TypeVar("T", bound=Expr)
-
-
-class AstPrinter(Generic[T]):
-    LAST_CHILD = "└── "
-    CHILD = "├── "
-    VERTICAL = "│   "
-    EMPTY = "    "
-
-    def __init__(self):
-        self._levels: list[_Level] = []
-        self._idx: Optional[int] = None
-        self._buf: io.StringIO = io.StringIO()
-
-    def print(self, expr: T):
-        self._buf = io.StringIO()
-        expr.accept(self)
-        return self._buf.getvalue()
-
-    @contextmanager
-    def _child_level(self, single: bool = False) -> Generator[None, None, None]:
-        self._levels.append(_Level.LAST if single else _Level.ACTIVE)
-        try:
-            yield
-        finally:
-            self._levels.pop()
-
-    def _mark_last(self):
-        if self._levels:
-            self._levels[-1] = _Level.LAST
-
-    def _write_line(self, text: str, *, last: bool = False):
-        if last:
-            self._mark_last()
-        indent: str = self._build_indent()
-        if self._idx is not None:
-            text = f"[{self._idx}] {text}"
-            self._idx = None
-        self._buf.write(indent + text + "\n")
-
-    def _build_indent(self) -> str:
-        parts: list[str] = []
-        for level in self._levels[:-1]:
-            parts.append(self.EMPTY if level == _Level.EMPTY else self.VERTICAL)
-        if self._levels:
-            if self._levels[-1] == _Level.LAST:
-                parts.append(self.LAST_CHILD)
-                self._levels[-1] = _Level.EMPTY
-            else:
-                parts.append(self.CHILD)
-        return "".join(parts)
-
-    def _write_optional_child(
-        self, label: str, child: Optional[T], *, last: bool = False
-    ):
-        if last:
-            self._mark_last()
-        if child is None:
-            self._write_line(f"{label}: None")
-        else:
-            self._write_line(label)
-            with self._child_level(single=True):
-                child.accept(self)
-
-
-class MidasAstPrinter(
-    AstPrinter, m.Expr.Visitor[None], m.Stmt.Visitor[None], m.Type.Visitor[None]
-):
-    # Statements
-
-    def visit_type_stmt(self, stmt: m.TypeStmt) -> None:
-        self._write_line("TypeStmt")
-        with self._child_level():
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line("params")
-            with self._child_level():
-                for i, param in enumerate(stmt.params):
-                    self._idx = i
-                    if i == len(stmt.params) - 1:
-                        self._mark_last()
-                    self._print_type_param(param)
-            self._write_line("type", last=True)
-            with self._child_level(single=True):
-                stmt.type.accept(self)
-
-    def visit_alias_stmt(self, stmt: m.AliasStmt) -> None:
-        self._write_line("AliasStmt")
-        with self._child_level():
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line("type", last=True)
-            with self._child_level(single=True):
-                stmt.type.accept(self)
-
-    def _print_type_param(self, param: m.TypeParam) -> None:
-        self._write_line("Param")
-        with self._child_level():
-            self._write_line(f'name: "{param.name.lexeme}"')
-            self._write_optional_child("bound", param.bound, last=True)
-
-    def visit_member_stmt(self, stmt: m.MemberStmt):
-        self._write_line("MemberStmt")
-        with self._child_level():
-            self._write_line(f"kind: {stmt.kind.name}")
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line("type", last=True)
-            with self._child_level(single=True):
-                stmt.type.accept(self)
-
-    def visit_extend_stmt(self, stmt: m.ExtendStmt) -> None:
-        self._write_line("ExtendStmt")
-        with self._child_level():
-            self._write_line("params")
-            with self._child_level():
-                for i, param in enumerate(stmt.params):
-                    self._idx = i
-                    if i == len(stmt.params) - 1:
-                        self._mark_last()
-                    self._print_type_param(param)
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            self._write_line("params")
-            with self._child_level():
-                for i, param in enumerate(stmt.params):
-                    self._idx = i
-                    if i == len(stmt.params) - 1:
-                        self._mark_last()
-                    self._print_type_param(param)
-            self._write_line("members", last=True)
-            with self._child_level():
-                for i, member in enumerate(stmt.members):
-                    self._idx = i
-                    if i == len(stmt.members) - 1:
-                        self._mark_last()
-                    member.accept(self)
-
-    def visit_predicate_stmt(self, stmt: m.PredicateStmt):
-        self._write_line("PredicateStmt")
-        with self._child_level():
-            self._write_line(f'name: "{stmt.name.lexeme}"')
-            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.body.accept(self)
-
-    # Expressions
-
-    def visit_logical_expr(self, expr: m.LogicalExpr):
-        self._write_line("LogicalExpr")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level(single=True):
-                expr.left.accept(self)
-
-            self._write_line(f"operator: {expr.operator.lexeme}")
-
-            self._write_line("right", last=True)
-            with self._child_level(single=True):
-                expr.right.accept(self)
-
-    def visit_binary_expr(self, expr: m.BinaryExpr):
-        self._write_line("BinaryExpr")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level(single=True):
-                expr.left.accept(self)
-
-            self._write_line(f"operator: {expr.operator.lexeme}")
-
-            self._write_line("right", last=True)
-            with self._child_level(single=True):
-                expr.right.accept(self)
-
-    def visit_unary_expr(self, expr: m.UnaryExpr):
-        self._write_line("UnaryExpr")
-        with self._child_level():
-            self._write_line(f"operator: {expr.operator.lexeme}")
-
-            self._write_line("right", last=True)
-            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():
-            self._write_line("expr")
-            with self._child_level(single=True):
-                expr.expr.accept(self)
-            self._write_line(f'name: "{expr.name.lexeme}"', last=True)
-
-    def visit_variable_expr(self, expr: m.VariableExpr):
-        self._write_line("VariableExpr")
-        with self._child_level():
-            self._write_line(f'name: "{expr.name.lexeme}"', last=True)
-
-    def visit_grouping_expr(self, expr: m.GroupingExpr):
-        self._write_line("GroupingExpr")
-        with self._child_level():
-            self._write_line("expr", last=True)
-            with self._child_level(single=True):
-                expr.expr.accept(self)
-
-    def visit_literal_expr(self, expr: m.LiteralExpr) -> None:
-        self._write_line("LiteralExpr")
-        with self._child_level():
-            self._write_line(f"value: {expr.value}", last=True)
-
-    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> None:
-        self._write_line("WildcardExpr")
-
-    def visit_named_type(self, type: m.NamedType) -> None:
-        self._write_line("NamedType")
-        with self._child_level():
-            self._write_line(f'name: "{type.name.lexeme}"', last=True)
-
-    def visit_generic_type(self, type: m.GenericType) -> None:
-        self._write_line("GenericType")
-        with self._child_level():
-            self._write_line("type")
-            with self._child_level():
-                type.type.accept(self)
-            self._write_line("args", last=True)
-            with self._child_level():
-                for i, param in enumerate(type.args):
-                    self._idx = i
-                    if i == len(type.args) - 1:
-                        self._mark_last()
-                    param.accept(self)
-
-    def visit_constraint_type(self, type: m.ConstraintType) -> None:
-        self._write_line("ConstraintType")
-        with self._child_level():
-            self._write_line("type")
-            with self._child_level(single=True):
-                type.type.accept(self)
-            self._write_line("constraint", last=True)
-            with self._child_level(single=True):
-                type.constraint.accept(self)
-
-    def visit_complex_type(self, type: m.ComplexType) -> None:
-        self._write_line("ComplexType")
-        with self._child_level():
-            self._write_line("members", last=True)
-            with self._child_level():
-                for i, member in enumerate(type.members):
-                    self._idx = i
-                    if i == len(type.members) - 1:
-                        self._mark_last()
-                    member.accept(self)
-
-    def visit_extension_type(self, type: m.ExtensionType) -> None:
-        self._write_line("ExtensionType")
-        with self._child_level():
-            self._write_line("base")
-            with self._child_level(single=True):
-                type.base.accept(self)
-            self._write_line("extension", last=True)
-            with self._child_level(single=True):
-                type.extension.accept(self)
-
-    def visit_function_type(self, type: m.FunctionType) -> None:
-        self._write_line("FunctionType")
-        with self._child_level():
-            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():
-            name: str = "None"
-            if arg.name is not None:
-                name = f'"{arg.name.lexeme}"'
-            self._write_line(f"name: {name}")
-            self._write_line("type")
-            with self._child_level(single=True):
-                arg.type.accept(self)
-            self._write_line(f"required: {arg.required}", last=True)
-
-    def visit_frame_type(self, type: m.FrameType) -> None:
-        self._write_line("FrameType")
-        with self._child_level(single=True):
-            self._write_line("columns")
-            with self._child_level():
-                for i, column in enumerate(type.columns):
-                    self._idx = i
-                    if i == len(type.columns) - 1:
-                        self._mark_last()
-                    self._print_frame_column(column)
-
-    def _print_frame_column(self, column: m.FrameType.Column) -> None:
-        self._write_line("Column")
-        with self._child_level():
-            self._write_line(f'name: "{column.name.lexeme}"')
-            self._write_line("type")
-            with self._child_level(single=True):
-                column.type.accept(self)
-
-
-class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str]):
-    def __init__(self, indent: int = 4):
-        self.indent: int = indent
-        self.level: int = 0
-
-    def indented(self, text: str) -> str:
-        return " " * (self.level * self.indent) + text
-
-    def print(self, expr: m.Expr | m.Stmt | m.Type) -> str:
-        self.level = 0
-        return expr.accept(self)
-
-    def visit_type_stmt(self, stmt: m.TypeStmt) -> str:
-        template: str = ""
-        if len(stmt.params) != 0:
-            params: list[str] = [self._print_type_param(param) for param in stmt.params]
-            template = f"[{', '.join(params)}]"
-        res: str = f"type {stmt.name.lexeme}{template} = {stmt.type.accept(self)}"
-        return self.indented(res)
-
-    def visit_alias_stmt(self, stmt: m.AliasStmt) -> str:
-        return self.indented(f"alias {stmt.name.lexeme} = {stmt.type.accept(self)}")
-
-    def _print_type_param(self, param: m.TypeParam) -> str:
-        res: str = param.name.lexeme
-        if param.bound is not None:
-            res += "<:" + param.bound.accept(self)
-        return res
-
-    def visit_member_stmt(self, stmt: m.MemberStmt):
-        keyword: str = {
-            m.MemberKind.PROPERTY: "prop",
-            m.MemberKind.METHOD: "def",
-        }.get(stmt.kind, "")
-        res: str = f"{keyword} {stmt.name.lexeme}: {stmt.type.accept(self)}"
-        return self.indented(res)
-
-    def visit_extend_stmt(self, stmt: m.ExtendStmt):
-        template: str = ""
-        if len(stmt.params) != 0:
-            params: list[str] = [self._print_type_param(param) for param in stmt.params]
-            template = f"[{', '.join(params)}]"
-        res: str = self.indented(f"extend {stmt.name.lexeme}{template}")
-        res += " {\n"
-        self.level += 1
-        for member in stmt.members:
-            res += member.accept(self) + "\n"
-        self.level -= 1
-        res += self.indented("}")
-        return res
-
-    def visit_predicate_stmt(self, stmt: m.PredicateStmt):
-        name: str = stmt.name.lexeme
-        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)
-        operator: str = expr.operator.lexeme
-        right: str = expr.right.accept(self)
-        return f"{left} {operator} {right}"
-
-    def visit_binary_expr(self, expr: m.BinaryExpr):
-        left: str = expr.left.accept(self)
-        operator: str = expr.operator.lexeme
-        right: str = expr.right.accept(self)
-        return f"{left} {operator} {right}"
-
-    def visit_unary_expr(self, expr: m.UnaryExpr):
-        operator: str = expr.operator.lexeme
-        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
-        return f"{expr_}.{name}"
-
-    def visit_variable_expr(self, expr: m.VariableExpr):
-        return expr.name.lexeme
-
-    def visit_grouping_expr(self, expr: m.GroupingExpr):
-        expr_: str = expr.expr.accept(self)
-        return f"({expr_})"
-
-    def visit_literal_expr(self, expr: m.LiteralExpr):
-        return str(expr.value)
-
-    def visit_wildcard_expr(self, expr: m.WildcardExpr):
-        return "_"
-
-    def visit_named_type(self, type: m.NamedType) -> str:
-        return type.name.lexeme
-
-    def visit_generic_type(self, type: m.GenericType) -> str:
-        res: str = type.type.accept(self)
-        if len(type.args) != 0:
-            args: list[str] = [param.accept(self) for param in type.args]
-            res += f"[{', '.join(args)}]"
-        return res
-
-    def visit_constraint_type(self, type: m.ConstraintType) -> str:
-        res: str = type.type.accept(self)
-        res += " where " + type.constraint.accept(self)
-        return res
-
-    def visit_complex_type(self, type: m.ComplexType) -> str:
-        res: str = "{\n"
-        self.level += 1
-        for member in type.members:
-            res += member.accept(self)
-            res += "\n"
-        self.level -= 1
-        res += self.indented("}")
-        return res
-
-    def visit_extension_type(self, type: m.ExtensionType) -> str:
-        return f"{type.base.accept(self)} & {type.extension.accept(self)}"
-
-    def visit_function_type(self, type: m.FunctionType) -> str:
-        spec: str = self._visit_param_spec(type.params)
-        return f"fn {spec} -> {type.returns.accept(self)}"
-
-    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:
-            args.append("/")
-        args += mixed_args
-        if len(kw_args) != 0:
-            args.append("*")
-        args += kw_args
-        return f"({', '.join(args)})"
-
-    def _print_arg(self, arg: m.FunctionType.Argument) -> str:
-        res: str = ""
-        if arg.name is not None:
-            res += arg.name.lexeme
-            res += ": "
-        res += arg.type.accept(self)
-        if not arg.required:
-            res += "?"
-        return res
-
-    def visit_frame_type(self, type: m.FrameType) -> str:
-        res: str = self.indented("Frame[")
-        if len(type.columns) != 0:
-            res += "\n"
-            self.level += 1
-            columns: list[str] = []
-            for column in type.columns:
-                columns.append(self.indented(self._print_frame_column(column)))
-            res += ",\n".join(columns)
-            self.level -= 1
-            res += "\n"
-        res += "]"
-        return res
-
-    def _print_frame_column(self, column: m.FrameType.Column) -> str:
-        return f"{column.name.lexeme}: {column.type.accept(self)}"
-
-
-class PythonAstPrinter(
-    AstPrinter,
-    p.MidasType.Visitor[None],
-    p.Stmt.Visitor[None],
-    p.Expr.Visitor[None],
-):
-    def visit_base_type(self, node: p.BaseType) -> None:
-        self._write_line("BaseType")
-        with self._child_level():
-            self._write_line(f"base: {node.base}")
-            self._write_line("args:", last=True)
-            with self._child_level():
-                for i, arg in enumerate(node.args):
-                    self._idx = i
-                    if i == len(node.args) - 1:
-                        self._mark_last()
-                    arg.accept(self)
-
-    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():
-            self._write_line(f"name: {node.name}")
-            self._write_optional_child("type", node.type, last=True)
-
-    def visit_frame_type(self, node: p.FrameType) -> None:
-        self._write_line("FrameType")
-        with self._child_level():
-            self._write_line("columns", last=True)
-            with self._child_level():
-                for i, col in enumerate(node.columns):
-                    self._idx = i
-                    if i == len(node.columns) - 1:
-                        self._mark_last()
-                    col.accept(self)
-
-    def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
-        stmt.expr.accept(self)
-
-    def visit_function(self, stmt: p.Function) -> None:
-        self._write_line("Function")
-        with self._child_level():
-            self._write_line(f"name: {stmt.name}")
-
-            self._write_line("posonlyargs")
-            with self._child_level():
-                for i, arg in enumerate(stmt.posonlyargs):
-                    self._idx = i
-                    if i == len(stmt.posonlyargs) - 1:
-                        self._mark_last()
-                    self._print_argument(arg)
-
-            self._write_line("args")
-            with self._child_level():
-                for i, arg in enumerate(stmt.args):
-                    self._idx = i
-                    if i == len(stmt.args) - 1:
-                        self._mark_last()
-                    self._print_argument(arg)
-
-            self._write_line("kwonlyargs")
-            with self._child_level():
-                for i, arg in enumerate(stmt.kwonlyargs):
-                    self._idx = i
-                    if i == len(stmt.kwonlyargs) - 1:
-                        self._mark_last()
-                    self._print_argument(arg)
-
-            self._write_optional_child("returns", stmt.returns)
-            self._write_line("body", last=True)
-            with self._child_level():
-                for i, body_stmt in enumerate(stmt.body):
-                    self._idx = i
-                    if i == len(stmt.body) - 1:
-                        self._mark_last()
-                    body_stmt.accept(self)
-
-    def _print_argument(self, arg: p.Function.Argument) -> None:
-        self._write_line("FunctionArgument")
-        with self._child_level():
-            self._write_line(f"name: {arg.name}")
-            self._write_optional_child("type", arg.type, last=True)
-
-    def visit_type_assign(self, stmt: p.TypeAssign) -> None:
-        self._write_line("TypeAssign")
-        with self._child_level():
-            self._write_line(f"name: {stmt.name}")
-            self._write_line("type", last=True)
-            with self._child_level(single=True):
-                stmt.type.accept(self)
-
-    def visit_assign_stmt(self, stmt: p.AssignStmt) -> None:
-        self._write_line("AssignStmt")
-        with self._child_level():
-            self._write_line("targets")
-            with self._child_level():
-                for i, target in enumerate(stmt.targets):
-                    self._idx = i
-                    if i == len(stmt.targets) - 1:
-                        self._mark_last()
-                    target.accept(self)
-            self._write_line("value", last=True)
-            with self._child_level(single=True):
-                stmt.value.accept(self)
-
-    def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
-        self._write_line("ReturnStmt")
-        with self._child_level():
-            self._write_optional_child("value", stmt.value, last=True)
-
-    def visit_if_stmt(self, stmt: p.IfStmt) -> None:
-        self._write_line("IfStmt")
-        with self._child_level():
-            self._write_line("test")
-            with self._child_level(single=True):
-                stmt.test.accept(self)
-            self._write_line("body")
-            with self._child_level():
-                for i, body_stmt in enumerate(stmt.body):
-                    self._idx = i
-                    if i == len(stmt.body) - 1:
-                        self._mark_last()
-                    body_stmt.accept(self)
-            self._write_line("orelse", last=True)
-            with self._child_level():
-                for i, else_stmt in enumerate(stmt.orelse):
-                    self._idx = i
-                    if i == len(stmt.orelse) - 1:
-                        self._mark_last()
-                    else_stmt.accept(self)
-
-    def visit_pass(self, stmt: p.Pass) -> None:
-        self._write_line("Pass")
-
-    def visit_for_stmt(self, stmt: p.ForStmt) -> None:
-        self._write_line("ForStmt")
-        with self._child_level():
-            self._write_line("target")
-            with self._child_level(single=True):
-                stmt.target.accept(self)
-            self._write_line("iterator")
-            with self._child_level(single=True):
-                stmt.iterator.accept(self)
-            self._write_line("body", last=True)
-            with self._child_level():
-                for i, body_stmt in enumerate(stmt.body):
-                    self._idx = i
-                    if i == len(stmt.body) - 1:
-                        self._mark_last()
-                    body_stmt.accept(self)
-
-    def visit_raw_stmt(self, stmt: p.RawStmt) -> None:
-        self._write_line("RawStmt")
-        with self._child_level(single=True):
-            self._write_line(f"stmt: {ast.unparse(stmt.stmt)}")
-
-    def visit_binary_expr(self, expr: p.BinaryExpr) -> None:
-        self._write_line("BinaryExpr")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level(single=True):
-                expr.left.accept(self)
-
-            self._write_line(f"operator: {expr.operator.__class__.__name__}")
-
-            self._write_line("right", last=True)
-            with self._child_level(single=True):
-                expr.right.accept(self)
-
-    def visit_compare_expr(self, expr: p.CompareExpr) -> None:
-        self._write_line("CompareExpr")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level(single=True):
-                expr.left.accept(self)
-
-            self._write_line(f"operator: {expr.operator.__class__.__name__}")
-
-            self._write_line("right", last=True)
-            with self._child_level(single=True):
-                expr.right.accept(self)
-
-    def visit_unary_expr(self, expr: p.UnaryExpr) -> None:
-        self._write_line("UnaryExpr")
-        with self._child_level():
-            self._write_line(f"operator: {expr.operator.__class__.__name__}")
-
-            self._write_line("right", last=True)
-            with self._child_level(single=True):
-                expr.right.accept(self)
-
-    def visit_call_expr(self, expr: p.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: p.GetExpr) -> None:
-        self._write_line("GetExpr")
-        with self._child_level():
-            self._write_line("object")
-            with self._child_level(single=True):
-                expr.object.accept(self)
-            self._write_line(f"name: {expr.name}", last=True)
-
-    def visit_literal_expr(self, expr: p.LiteralExpr) -> None:
-        self._write_line("LiteralExpr")
-        with self._child_level(single=True):
-            self._write_line(f"value: {expr.value!r}")
-
-    def visit_variable_expr(self, expr: p.VariableExpr) -> None:
-        self._write_line("VariableExpr")
-        with self._child_level(single=True):
-            self._write_line(f"name: {expr.name}")
-
-    def visit_logical_expr(self, expr: p.LogicalExpr) -> None:
-        self._write_line("LogicalExpr")
-        with self._child_level():
-            self._write_line("left")
-            with self._child_level(single=True):
-                expr.left.accept(self)
-
-            self._write_line(f"operator: {expr.operator.__class__.__name__}")
-
-            self._write_line("right", last=True)
-            with self._child_level(single=True):
-                expr.right.accept(self)
-
-    def visit_cast_expr(self, expr: p.CastExpr) -> None:
-        self._write_line("CastExpr")
-        with self._child_level():
-            self._write_line("type")
-            with self._child_level(single=True):
-                expr.type.accept(self)
-            self._write_line("expr")
-            with self._child_level(single=True):
-                expr.expr.accept(self)
-            self._write_line(f"unsafe: {expr.unsafe}", last=True)
-
-    def visit_ternary_expr(self, expr: p.TernaryExpr) -> None:
-        self._write_line("TernaryExpr")
-        with self._child_level():
-            self._write_line("test")
-            with self._child_level(single=True):
-                expr.test.accept(self)
-
-            self._write_line("if_true")
-            with self._child_level(single=True):
-                expr.if_true.accept(self)
-
-            self._write_line("if_false", last=True)
-            with self._child_level(single=True):
-                expr.if_false.accept(self)
-
-    def visit_list_expr(self, expr: p.ListExpr) -> None:
-        self._write_line("ListExpr")
-        with self._child_level():
-            self._write_line("items", last=True)
-            with self._child_level():
-                for i, item in enumerate(expr.items):
-                    self._idx = i
-                    if i == len(expr.items) - 1:
-                        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():
-            self._write_line("object")
-            with self._child_level(single=True):
-                expr.object.accept(self)
-            self._write_line("index", last=True)
-            with self._child_level(single=True):
-                expr.index.accept(self)
-
-    def visit_slice_expr(self, expr: p.SliceExpr) -> None:
-        self._write_line("SliceExpr")
-        with self._child_level():
-            self._write_optional_child("lower", expr.lower)
-            self._write_optional_child("upper", expr.upper)
-            self._write_optional_child("step", expr.step, last=True)
-
-    def visit_tuple_expr(self, expr: p.TupleExpr) -> None:
-        self._write_line("TupleExpr")
-        with self._child_level():
-            self._write_line("items", last=True)
-            with self._child_level():
-                for i, item in enumerate(expr.items):
-                    self._idx = i
-                    if i == len(expr.items) - 1:
-                        self._mark_last()
-                    item.accept(self)
-
-    def visit_raw_expr(self, expr: p.RawExpr) -> None:
-        self._write_line("RawExpr")
-        with self._child_level(single=True):
-            self._write_line(f"expr: {ast.unparse(expr.expr)}")

midas/ast/printer/__init__.py

@@ -0,0 +1,3 @@
+from .midas import MidasPrinter as MidasPrinter
+from .midas_ast import MidasAstPrinter as MidasAstPrinter
+from .python_ast import PythonAstPrinter as PythonAstPrinter

midas/ast/printer/base.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import io
+from contextlib import contextmanager
+from enum import Enum, auto
+from typing import Callable, Generator, Generic, Optional, Protocol, Sequence, TypeVar
+
+
+class _Level(Enum):
+    EMPTY = auto()
+    ACTIVE = auto()
+    LAST = auto()
+
+
+class Expr(Protocol):
+    def accept(self, printer: AstPrinter) -> None: ...
+
+
+T = TypeVar("T", bound=Expr)
+
+
+class AstPrinter(Generic[T]):
+    LAST_CHILD = "└── "
+    CHILD = "├── "
+    VERTICAL = "│   "
+    EMPTY = "    "
+
+    def __init__(self):
+        self._levels: list[_Level] = []
+        self._idx: Optional[int] = None
+        self._buf: io.StringIO = io.StringIO()
+
+    def print(self, expr: T):
+        self._buf = io.StringIO()
+        expr.accept(self)
+        return self._buf.getvalue()
+
+    @contextmanager
+    def _child_level(self, single: bool = False) -> Generator[None, None, None]:
+        self._levels.append(_Level.LAST if single else _Level.ACTIVE)
+        try:
+            yield
+        finally:
+            self._levels.pop()
+
+    def _mark_last(self):
+        if self._levels:
+            self._levels[-1] = _Level.LAST
+
+    def _write_line(self, text: str, *, last: bool = False):
+        if last:
+            self._mark_last()
+        indent: str = self._build_indent()
+        if self._idx is not None:
+            text = f"[{self._idx}] {text}"
+            self._idx = None
+        self._buf.write(indent + text + "\n")
+
+    def _build_indent(self) -> str:
+        parts: list[str] = []
+        for level in self._levels[:-1]:
+            parts.append(self.EMPTY if level == _Level.EMPTY else self.VERTICAL)
+        if self._levels:
+            if self._levels[-1] == _Level.LAST:
+                parts.append(self.LAST_CHILD)
+                self._levels[-1] = _Level.EMPTY
+            else:
+                parts.append(self.CHILD)
+        return "".join(parts)
+
+    def _write_optional_child(
+        self, label: str, child: Optional[T], *, last: bool = False
+    ):
+        if last:
+            self._mark_last()
+        if child is None:
+            self._write_line(f"{label}: None")
+        else:
+            self._write_line(label)
+            with self._child_level(single=True):
+                child.accept(self)
+
+    def _write_sequence(
+        self,
+        label: str,
+        list_: Sequence[T],
+        *,
+        last: bool = False,
+        print_func: Optional[Callable[[T], None]] = None,
+    ):
+        if last:
+            self._mark_last()
+
+        self._write_line(label)
+        with self._child_level():
+            for i, item in enumerate(list_):
+                self._idx = i
+                if i == len(list_) - 1:
+                    self._mark_last()
+                if print_func is not None:
+                    print_func(item)
+                else:
+                    item.accept(self)

midas/ast/printer/midas.py

@@ -0,0 +1,183 @@
+import midas.ast.midas as m
+
+
+class MidasPrinter(
+    m.Expr.Visitor[str],
+    m.Stmt.Visitor[str],
+    m.Type.Visitor[str],
+):
+    def __init__(self, indent: int = 4):
+        self.indent: int = indent
+        self.level: int = 0
+
+    def indented(self, text: str) -> str:
+        return " " * (self.level * self.indent) + text
+
+    def print(self, expr: m.Expr | m.Stmt | m.Type) -> str:
+        self.level = 0
+        return expr.accept(self)
+
+    # Statements
+
+    def visit_type_stmt(self, stmt: m.TypeStmt) -> str:
+        template: str = ""
+        if len(stmt.params) != 0:
+            params: list[str] = [self._print_type_param(param) for param in stmt.params]
+            template = f"[{', '.join(params)}]"
+        res: str = f"type {stmt.name.lexeme}{template} = {stmt.type.accept(self)}"
+        return self.indented(res)
+
+    def visit_alias_stmt(self, stmt: m.AliasStmt) -> str:
+        return self.indented(f"alias {stmt.name.lexeme} = {stmt.type.accept(self)}")
+
+    def _print_type_param(self, param: m.TypeParam) -> str:
+        res: str = param.name.lexeme
+        if param.bound is not None:
+            res += "<:" + param.bound.accept(self)
+        return res
+
+    def visit_member_stmt(self, stmt: m.MemberStmt):
+        keyword: str = {
+            m.MemberKind.PROPERTY: "prop",
+            m.MemberKind.METHOD: "def",
+        }.get(stmt.kind, "")
+        res: str = f"{keyword} {stmt.name.lexeme}: {stmt.type.accept(self)}"
+        return self.indented(res)
+
+    def visit_extend_stmt(self, stmt: m.ExtendStmt):
+        template: str = ""
+        if len(stmt.params) != 0:
+            params: list[str] = [self._print_type_param(param) for param in stmt.params]
+            template = f"[{', '.join(params)}]"
+        res: str = self.indented(f"extend {stmt.name.lexeme}{template}")
+        res += " {\n"
+        self.level += 1
+        for member in stmt.members:
+            res += member.accept(self) + "\n"
+        self.level -= 1
+        res += self.indented("}")
+        return res
+
+    def visit_predicate_stmt(self, stmt: m.PredicateStmt):
+        name: str = stmt.name.lexeme
+        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}")
+
+    # Expressions
+
+    def visit_logical_expr(self, expr: m.LogicalExpr):
+        left: str = expr.left.accept(self)
+        operator: str = expr.operator.lexeme
+        right: str = expr.right.accept(self)
+        return f"{left} {operator} {right}"
+
+    def visit_binary_expr(self, expr: m.BinaryExpr):
+        left: str = expr.left.accept(self)
+        operator: str = expr.operator.lexeme
+        right: str = expr.right.accept(self)
+        return f"{left} {operator} {right}"
+
+    def visit_unary_expr(self, expr: m.UnaryExpr):
+        operator: str = expr.operator.lexeme
+        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
+        return f"{expr_}.{name}"
+
+    def visit_variable_expr(self, expr: m.VariableExpr):
+        return expr.name.lexeme
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr):
+        expr_: str = expr.expr.accept(self)
+        return f"({expr_})"
+
+    def visit_literal_expr(self, expr: m.LiteralExpr):
+        return str(expr.value)
+
+    def visit_wildcard_expr(self, expr: m.WildcardExpr):
+        return "_"
+
+    # Types
+
+    def visit_named_type(self, type: m.NamedType) -> str:
+        return type.name.lexeme
+
+    def visit_generic_type(self, type: m.GenericType) -> str:
+        res: str = type.type.accept(self)
+        if len(type.args) != 0:
+            args: list[str] = [param.accept(self) for param in type.args]
+            res += f"[{', '.join(args)}]"
+        return res
+
+    def visit_constraint_type(self, type: m.ConstraintType) -> str:
+        res: str = type.type.accept(self)
+        res += " where " + type.constraint.accept(self)
+        return res
+
+    def visit_complex_type(self, type: m.ComplexType) -> str:
+        res: str = "{\n"
+        self.level += 1
+        for member in type.members:
+            res += member.accept(self)
+            res += "\n"
+        self.level -= 1
+        res += self.indented("}")
+        return res
+
+    def visit_extension_type(self, type: m.ExtensionType) -> str:
+        return f"{type.base.accept(self)} & {type.extension.accept(self)}"
+
+    def visit_function_type(self, type: m.FunctionType) -> str:
+        spec: str = self._visit_param_spec(type.params)
+        return f"fn {spec} -> {type.returns.accept(self)}"
+
+    def _visit_param_spec(self, spec: m.ParamSpec) -> str:
+        pos: list[str] = [self._print_param(param) for param in spec.pos]
+        mixed: list[str] = [self._print_param(param) for param in spec.mixed]
+        kw: list[str] = [self._print_param(param) for param in spec.kw]
+        params: list[str] = pos
+
+        if len(pos) != 0:
+            params.append("/")
+        params += mixed
+        if len(kw) != 0:
+            params.append("*")
+        params += kw
+        return f"({', '.join(params)})"
+
+    def _print_param(self, param: m.FunctionType.Parameter) -> str:
+        res: str = ""
+        if param.name is not None:
+            res += param.name.lexeme
+            res += ": "
+        res += param.type.accept(self)
+        if not param.required:
+            res += "?"
+        return res
+
+    def visit_frame_type(self, type: m.FrameType) -> str:
+        res: str = self.indented("Frame[")
+        if len(type.columns) != 0:
+            res += "\n"
+            self.level += 1
+            columns: list[str] = []
+            for column in type.columns:
+                columns.append(self.indented(self._print_frame_column(column)))
+            res += ",\n".join(columns)
+            self.level -= 1
+            res += "\n"
+        res += "]"
+        return res
+
+    def _print_frame_column(self, column: m.FrameType.Column) -> str:
+        return f"{column.name.lexeme}: {column.type.accept(self)}"

midas/ast/printer/midas_ast.py

@@ -0,0 +1,253 @@
+import midas.ast.midas as m
+from midas.ast.printer.base import AstPrinter
+
+
+class MidasAstPrinter(
+    AstPrinter,
+    m.Expr.Visitor[None],
+    m.Stmt.Visitor[None],
+    m.Type.Visitor[None],
+):
+    # Statements
+
+    def visit_type_stmt(self, stmt: m.TypeStmt) -> None:
+        self._write_line("TypeStmt")
+        with self._child_level():
+            self._write_line(f'name: "{stmt.name.lexeme}"')
+            self._write_sequence(
+                "params",
+                stmt.params,
+                print_func=self._print_type_param,
+            )
+            self._write_line("type", last=True)
+            with self._child_level(single=True):
+                stmt.type.accept(self)
+
+    def visit_alias_stmt(self, stmt: m.AliasStmt) -> None:
+        self._write_line("AliasStmt")
+        with self._child_level():
+            self._write_line(f'name: "{stmt.name.lexeme}"')
+            self._write_line("type", last=True)
+            with self._child_level(single=True):
+                stmt.type.accept(self)
+
+    def _print_type_param(self, param: m.TypeParam) -> None:
+        self._write_line("Param")
+        with self._child_level():
+            self._write_line(f'name: "{param.name.lexeme}"')
+            self._write_optional_child("bound", param.bound, last=True)
+
+    def visit_member_stmt(self, stmt: m.MemberStmt):
+        self._write_line("MemberStmt")
+        with self._child_level():
+            self._write_line(f"kind: {stmt.kind.name}")
+            self._write_line(f'name: "{stmt.name.lexeme}"')
+            self._write_line("type", last=True)
+            with self._child_level(single=True):
+                stmt.type.accept(self)
+
+    def visit_extend_stmt(self, stmt: m.ExtendStmt) -> None:
+        self._write_line("ExtendStmt")
+        with self._child_level():
+            self._write_line(f'name: "{stmt.name.lexeme}"')
+            self._write_sequence(
+                "params",
+                stmt.params,
+                print_func=self._print_type_param,
+            )
+            self._write_sequence("members", stmt.members, last=True)
+
+    def visit_predicate_stmt(self, stmt: m.PredicateStmt):
+        self._write_line("PredicateStmt")
+        with self._child_level():
+            self._write_line(f'name: "{stmt.name.lexeme}"')
+            self._write_sequence(
+                "params",
+                stmt.params,
+                print_func=self._visit_param_spec,
+            )
+            self._write_line("body", last=True)
+            with self._child_level(single=True):
+                stmt.body.accept(self)
+
+    # Expressions
+
+    def visit_logical_expr(self, expr: m.LogicalExpr):
+        self._write_line("LogicalExpr")
+        with self._child_level():
+            self._write_line("left")
+            with self._child_level(single=True):
+                expr.left.accept(self)
+
+            self._write_line(f"operator: {expr.operator.lexeme}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_binary_expr(self, expr: m.BinaryExpr):
+        self._write_line("BinaryExpr")
+        with self._child_level():
+            self._write_line("left")
+            with self._child_level(single=True):
+                expr.left.accept(self)
+
+            self._write_line(f"operator: {expr.operator.lexeme}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_unary_expr(self, expr: m.UnaryExpr):
+        self._write_line("UnaryExpr")
+        with self._child_level():
+            self._write_line(f"operator: {expr.operator.lexeme}")
+
+            self._write_line("right", last=True)
+            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_sequence("arguments", expr.arguments)
+            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():
+            self._write_line("expr")
+            with self._child_level(single=True):
+                expr.expr.accept(self)
+            self._write_line(f'name: "{expr.name.lexeme}"', last=True)
+
+    def visit_variable_expr(self, expr: m.VariableExpr):
+        self._write_line("VariableExpr")
+        with self._child_level():
+            self._write_line(f'name: "{expr.name.lexeme}"', last=True)
+
+    def visit_grouping_expr(self, expr: m.GroupingExpr):
+        self._write_line("GroupingExpr")
+        with self._child_level():
+            self._write_line("expr", last=True)
+            with self._child_level(single=True):
+                expr.expr.accept(self)
+
+    def visit_literal_expr(self, expr: m.LiteralExpr) -> None:
+        self._write_line("LiteralExpr")
+        with self._child_level():
+            self._write_line(f"value: {expr.value}", last=True)
+
+    def visit_wildcard_expr(self, expr: m.WildcardExpr) -> None:
+        self._write_line("WildcardExpr")
+
+    # Types
+
+    def visit_named_type(self, type: m.NamedType) -> None:
+        self._write_line("NamedType")
+        with self._child_level():
+            self._write_line(f'name: "{type.name.lexeme}"', last=True)
+
+    def visit_generic_type(self, type: m.GenericType) -> None:
+        self._write_line("GenericType")
+        with self._child_level():
+            self._write_line("type")
+            with self._child_level():
+                type.type.accept(self)
+            self._write_sequence("args", type.args, last=True)
+
+    def visit_constraint_type(self, type: m.ConstraintType) -> None:
+        self._write_line("ConstraintType")
+        with self._child_level():
+            self._write_line("type")
+            with self._child_level(single=True):
+                type.type.accept(self)
+            self._write_line("constraint", last=True)
+            with self._child_level(single=True):
+                type.constraint.accept(self)
+
+    def visit_complex_type(self, type: m.ComplexType) -> None:
+        self._write_line("ComplexType")
+        with self._child_level():
+            self._write_sequence("members", type.members, last=True)
+
+    def visit_extension_type(self, type: m.ExtensionType) -> None:
+        self._write_line("ExtensionType")
+        with self._child_level():
+            self._write_line("base")
+            with self._child_level(single=True):
+                type.base.accept(self)
+            self._write_line("extension", last=True)
+            with self._child_level(single=True):
+                type.extension.accept(self)
+
+    def visit_function_type(self, type: m.FunctionType) -> None:
+        self._write_line("FunctionType")
+        with self._child_level():
+            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_sequence(
+                "pos",
+                spec.pos,
+                print_func=self._print_param,
+            )
+            self._write_sequence(
+                "mixed",
+                spec.mixed,
+                print_func=self._print_param,
+            )
+            self._write_sequence(
+                "kw",
+                spec.kw,
+                print_func=self._print_param,
+                last=True,
+            )
+
+    def _print_param(self, param: m.FunctionType.Parameter) -> None:
+        self._write_line("Parameter")
+        with self._child_level():
+            name: str = "None"
+            if param.name is not None:
+                name = f'"{param.name.lexeme}"'
+            self._write_line(f"name: {name}")
+            self._write_line("type")
+            with self._child_level(single=True):
+                param.type.accept(self)
+            self._write_line(f"required: {param.required}", last=True)
+
+    def visit_frame_type(self, type: m.FrameType) -> None:
+        self._write_line("FrameType")
+        with self._child_level(single=True):
+            self._write_sequence(
+                "columns",
+                type.columns,
+                print_func=self._print_frame_column,
+            )
+
+    def _print_frame_column(self, column: m.FrameType.Column) -> None:
+        self._write_line("Column")
+        with self._child_level():
+            self._write_line(f'name: "{column.name.lexeme}"')
+            self._write_line("type")
+            with self._child_level(single=True):
+                column.type.accept(self)

midas/ast/printer/python_ast.py

@@ -0,0 +1,285 @@
+import ast
+
+import midas.ast.python as p
+from midas.ast.printer.base import AstPrinter
+
+
+class PythonAstPrinter(
+    AstPrinter,
+    p.MidasType.Visitor[None],
+    p.Stmt.Visitor[None],
+    p.Expr.Visitor[None],
+):
+    # Types
+
+    def visit_base_type(self, node: p.BaseType) -> None:
+        self._write_line("BaseType")
+        with self._child_level():
+            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():
+            self._write_line(f"name: {node.name}")
+            self._write_optional_child("type", node.type, last=True)
+
+    def visit_frame_type(self, node: p.FrameType) -> None:
+        self._write_line("FrameType")
+        with self._child_level(single=True):
+            self._write_sequence("columns", node.columns)
+
+    # Statements
+
+    def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
+        stmt.expr.accept(self)
+
+    def visit_function(self, stmt: p.Function) -> None:
+        self._write_line("Function")
+        with self._child_level():
+            self._write_line(f"name: {stmt.name}")
+            self._write_line("params")
+            with self._child_level():
+                self._print_param_spec(stmt.params)
+
+            self._write_optional_child("returns", stmt.returns)
+            self._write_sequence("body", stmt.body, last=True)
+
+    def _print_param_spec(self, spec: p.ParamSpec) -> None:
+        self._write_line("ParamSpec")
+        with self._child_level():
+            self._write_sequence(
+                "pos",
+                spec.pos,
+                print_func=self._print_param,
+            )
+            self._write_sequence(
+                "mixed",
+                spec.mixed,
+                print_func=self._print_param,
+            )
+            self._write_sequence(
+                "kw",
+                spec.kw,
+                print_func=self._print_param,
+                last=True,
+            )
+
+    def _print_param(self, param: p.Function.Parameter) -> None:
+        self._write_line("Parameter")
+        with self._child_level():
+            self._write_line(f"name: {param.name}")
+            self._write_optional_child("type", param.type, last=True)
+
+    def visit_type_assign(self, stmt: p.TypeAssign) -> None:
+        self._write_line("TypeAssign")
+        with self._child_level():
+            self._write_line(f"name: {stmt.name}")
+            self._write_line("type", last=True)
+            with self._child_level(single=True):
+                stmt.type.accept(self)
+
+    def visit_assign_stmt(self, stmt: p.AssignStmt) -> None:
+        self._write_line("AssignStmt")
+        with self._child_level():
+            self._write_sequence("targets", stmt.targets)
+            self._write_line("value", last=True)
+            with self._child_level(single=True):
+                stmt.value.accept(self)
+
+    def visit_return_stmt(self, stmt: p.ReturnStmt) -> None:
+        self._write_line("ReturnStmt")
+        with self._child_level():
+            self._write_optional_child("value", stmt.value, last=True)
+
+    def visit_if_stmt(self, stmt: p.IfStmt) -> None:
+        self._write_line("IfStmt")
+        with self._child_level():
+            self._write_line("test")
+            with self._child_level(single=True):
+                stmt.test.accept(self)
+            self._write_sequence("body", stmt.body)
+            self._write_sequence("orelse", stmt.orelse, last=True)
+
+    def visit_pass(self, stmt: p.Pass) -> None:
+        self._write_line("Pass")
+
+    def visit_for_stmt(self, stmt: p.ForStmt) -> None:
+        self._write_line("ForStmt")
+        with self._child_level():
+            self._write_line("target")
+            with self._child_level(single=True):
+                stmt.target.accept(self)
+            self._write_line("iterator")
+            with self._child_level(single=True):
+                stmt.iterator.accept(self)
+            self._write_sequence("body", stmt.body, last=True)
+
+    def visit_raw_stmt(self, stmt: p.RawStmt) -> None:
+        self._write_line("RawStmt")
+        with self._child_level(single=True):
+            self._write_line(f"stmt: {ast.unparse(stmt.stmt)}")
+
+    # Expressions
+
+    def visit_binary_expr(self, expr: p.BinaryExpr) -> None:
+        self._write_line("BinaryExpr")
+        with self._child_level():
+            self._write_line("left")
+            with self._child_level(single=True):
+                expr.left.accept(self)
+
+            self._write_line(f"operator: {expr.operator.__class__.__name__}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_compare_expr(self, expr: p.CompareExpr) -> None:
+        self._write_line("CompareExpr")
+        with self._child_level():
+            self._write_line("left")
+            with self._child_level(single=True):
+                expr.left.accept(self)
+
+            self._write_line(f"operator: {expr.operator.__class__.__name__}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_unary_expr(self, expr: p.UnaryExpr) -> None:
+        self._write_line("UnaryExpr")
+        with self._child_level():
+            self._write_line(f"operator: {expr.operator.__class__.__name__}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_call_expr(self, expr: p.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_sequence("arguments", expr.arguments)
+            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: p.GetExpr) -> None:
+        self._write_line("GetExpr")
+        with self._child_level():
+            self._write_line("object")
+            with self._child_level(single=True):
+                expr.object.accept(self)
+            self._write_line(f"name: {expr.name}", last=True)
+
+    def visit_literal_expr(self, expr: p.LiteralExpr) -> None:
+        self._write_line("LiteralExpr")
+        with self._child_level(single=True):
+            self._write_line(f"value: {expr.value!r}")
+
+    def visit_variable_expr(self, expr: p.VariableExpr) -> None:
+        self._write_line("VariableExpr")
+        with self._child_level(single=True):
+            self._write_line(f"name: {expr.name}")
+
+    def visit_logical_expr(self, expr: p.LogicalExpr) -> None:
+        self._write_line("LogicalExpr")
+        with self._child_level():
+            self._write_line("left")
+            with self._child_level(single=True):
+                expr.left.accept(self)
+
+            self._write_line(f"operator: {expr.operator.__class__.__name__}")
+
+            self._write_line("right", last=True)
+            with self._child_level(single=True):
+                expr.right.accept(self)
+
+    def visit_cast_expr(self, expr: p.CastExpr) -> None:
+        self._write_line("CastExpr")
+        with self._child_level():
+            self._write_line("type")
+            with self._child_level(single=True):
+                expr.type.accept(self)
+            self._write_line("expr")
+            with self._child_level(single=True):
+                expr.expr.accept(self)
+            self._write_line(f"unsafe: {expr.unsafe}", last=True)
+
+    def visit_ternary_expr(self, expr: p.TernaryExpr) -> None:
+        self._write_line("TernaryExpr")
+        with self._child_level():
+            self._write_line("test")
+            with self._child_level(single=True):
+                expr.test.accept(self)
+
+            self._write_line("if_true")
+            with self._child_level(single=True):
+                expr.if_true.accept(self)
+
+            self._write_line("if_false", last=True)
+            with self._child_level(single=True):
+                expr.if_false.accept(self)
+
+    def visit_list_expr(self, expr: p.ListExpr) -> None:
+        self._write_line("ListExpr")
+        with self._child_level():
+            self._write_sequence("items", expr.items, last=True)
+
+    def visit_dict_expr(self, expr: p.DictExpr) -> None:
+        self._write_line("DictExpr")
+        with self._child_level():
+            self._write_sequence(
+                "keys",
+                expr.keys,
+                print_func=lambda k: (
+                    self._write_line("None") if k is None else k.accept(self)
+                ),
+            )
+            self._write_sequence("values", expr.values, last=True)
+
+    def visit_subscript_expr(self, expr: p.SubscriptExpr) -> None:
+        self._write_line("SubscriptExpr")
+        with self._child_level():
+            self._write_line("object")
+            with self._child_level(single=True):
+                expr.object.accept(self)
+            self._write_line("index", last=True)
+            with self._child_level(single=True):
+                expr.index.accept(self)
+
+    def visit_slice_expr(self, expr: p.SliceExpr) -> None:
+        self._write_line("SliceExpr")
+        with self._child_level():
+            self._write_optional_child("lower", expr.lower)
+            self._write_optional_child("upper", expr.upper)
+            self._write_optional_child("step", expr.step, last=True)
+
+    def visit_tuple_expr(self, expr: p.TupleExpr) -> None:
+        self._write_line("TupleExpr")
+        with self._child_level():
+            self._write_sequence("items", expr.items, last=True)
+
+    def visit_raw_expr(self, expr: p.RawExpr) -> None:
+        self._write_line("RawExpr")
+        with self._child_level(single=True):
+            self._write_line(f"expr: {ast.unparse(expr.expr)}")

midas/ast/python.py

@@ -14,6 +14,16 @@ from midas.ast.location import Location
 
 T = TypeVar("T")
 
+@dataclass(frozen=True, kw_only=True)
+class ParamSpec:
+    pos: list[Function.Parameter]
+    mixed: list[Function.Parameter]
+    kw: list[Function.Parameter]
+
+    @property
+    def all(self) -> list[Function.Parameter]:
+        return self.pos + self.mixed + self.kw
+
 
 ####################
 # Type annotations #
@@ -128,25 +138,17 @@ class ExpressionStmt(Stmt):
 @dataclass(frozen=True)
 class Function(Stmt):
     name: str
-    posonlyargs: list[Argument]
+    params: ParamSpec
-    args: list[Argument]
-    sink: Optional[Argument]
-    kwonlyargs: list[Argument]
-    kw_sink: Optional[Argument]
     returns: Optional[MidasType]
     body: list[Stmt]
 
     @dataclass(frozen=True, kw_only=True)
-    class Argument:
+    class Parameter:
         location: Optional[Location] = None
         name: str
         type: Optional[MidasType]
         default: Optional[Expr]
 
-    @property
-    def all_args(self) -> list[Argument]:
-        return self.posonlyargs + self.args + self.kwonlyargs
-
     def accept(self, visitor: Stmt.Visitor[T]) -> T:
         return visitor.visit_function(self)
 

midas/checker/builtins.py

@@ -14,12 +14,15 @@ if TYPE_CHECKING:
     from midas.checker.registry import TypesRegistry
 
 
-# Hard-coded subtype relationships between builtin types
-# Circular dependencies and diamond inheritance MUST be avoided
 BUILTIN_SUBTYPES: dict[str, set[str]] = {
     "object": {"float", "list", "dict", "str", "bytes", "tuple"},
     "float": {"int"},
 }
+"""
+Hard-coded subtype relationships between builtin types
+
+Circular dependencies and diamond inheritance MUST be avoided
+"""
 
 
 def define_builtins(reg: TypesRegistry):

midas/checker/checker.py

@@ -10,6 +10,11 @@ from midas.utils import TypedAST
 
 
 class TypeChecker:
+    """Type checking dispatcher
+
+    Contains a typer for Midas and one for Python, as well as the types registry
+    """
+
     def __init__(self):
         self.types: TypesRegistry = TypesRegistry()
         self.reporter: Reporter = Reporter()

midas/checker/diagnostic.py

@@ -14,6 +14,15 @@ class DiagnosticType(StrEnum):
 
 @dataclass(frozen=True)
 class Diagnostic:
+    """Information about a diagnostic (warning, errors, etc.)
+
+    Holds a location, a diagnostic type and a message.
+    Optionally bound to a file path
+
+    Returns:
+        _type_: _description_
+    """
+
     file_path: Optional[str]
     location: Location
     type: DiagnosticType
@@ -21,6 +30,18 @@ class Diagnostic:
 
     @property
     def location_str(self) -> str:
+        """The diagnostic type and location as a human readable string
+
+        The location is formatted as "<Type> in <file> from L<start_line>:<start_col> to <end_line>:<end_col>",
+        for example: "Error in /home/user/Desktop/script.py from L12:5 to L12:8"
+
+        If the file is `None`, the "in ..." section is excluded from the result.<br>
+        If the location's end is not specified, the formulation "at L<start_line>:<start_col>" is used.
+
+        Returns:
+            str: _description_
+        """
+
         start_loc: str = f"L{self.location.lineno}:{self.location.col_offset+1}"
         end_loc: Optional[str] = ""
         if (

midas/checker/dispatcher.py

@@ -30,9 +30,9 @@ TypedExpr = tuple[E, Type]
 
 @dataclass(frozen=True, kw_only=True)
 class MappedArgument(Generic[E]):
-    expr: E
+    arg_expr: E
-    type: Type
+    arg_type: Type
-    argument: Function.Argument
+    parameter: Function.Parameter
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -219,11 +219,11 @@ class CallDispatcher(Generic[E]):
         """
         valid: bool = True
         for arg in arguments:
-            if not self.types.is_subtype(arg.type, arg.argument.type):
+            if not self.types.is_subtype(arg.arg_type, arg.parameter.type):
                 if report_errors:
                     self.reporter.error(
-                        arg.expr.location,
+                        arg.arg_expr.location,
-                        f"Wrong type for argument '{arg.argument.name}', expected {arg.argument.type}, got {arg.type}",
+                        f"Wrong type for argument '{arg.parameter.name}', expected {arg.parameter.type}, got {arg.arg_type}",
                     )
                 valid = False
         return valid
@@ -347,28 +347,30 @@ class CallDispatcher(Generic[E]):
             tuple[bool, list[MappedArgument]]: a boolean reporting whether
             the call is valid and the list of mapped arguments
         """
-        set_args: set[str] = set()
+        set_params: set[str] = set()
 
         required_positional: list[str] = [
-            arg.name for arg in function.pos_args + function.args if arg.required
+            param.name
+            for param in function.params.pos + function.params.mixed
+            if param.required
         ]
         required_keyword: list[str] = [
-            arg.name for arg in function.kw_args if arg.required
+            param.name for param in function.params.kw if param.required
         ]
 
         mapped: list[MappedArgument[E]] = []
 
-        pos_params: list[Function.Argument] = list(function.pos_args)
+        pos_params: list[Function.Parameter] = list(function.params.pos)
-        mixed_params: list[Function.Argument] = list(function.args)
+        mixed_params: list[Function.Parameter] = list(function.params.mixed)
-        kw_params: dict[str, Function.Argument] = {
+        kw_params: dict[str, Function.Parameter] = {
-            arg.name: arg for arg in function.kw_args
+            param.name: param for param in function.params.kw
         }
 
         valid_call: bool = True
 
         # TODO: handle *args and **kwargs sinks
         for arg in positional:
-            param: Function.Argument
+            param: Function.Parameter
             if len(pos_params) != 0:
                 param = pos_params.pop(0)
             elif len(mixed_params) != 0:
@@ -385,27 +387,27 @@ class CallDispatcher(Generic[E]):
                 required_positional.remove(name)
             if name in required_keyword:
                 required_keyword.remove(name)
-            set_args.add(name)
+            set_params.add(name)
             mapped.append(
                 MappedArgument(
-                    expr=arg[0],
+                    arg_expr=arg[0],
-                    type=arg[1],
+                    arg_type=arg[1],
-                    argument=param,
+                    parameter=param,
                 )
             )
 
-        kw_params.update({arg.name: arg for arg in mixed_params})
+        kw_params.update({param.name: param for param in mixed_params})
         for name, arg in keywords.items():
-            param: Function.Argument
+            param: Function.Parameter
             if name not in kw_params:
                 if report_errors:
-                    if name in set_args:
+                    if name in set_params:
                         self.reporter.error(
-                            arg[0].location, f"Multiple values for argument '{name}'"
+                            arg[0].location, f"Multiple values for parameter '{name}'"
                         )
                     else:
                         self.reporter.error(
-                            arg[0].location, f"Unknown keyword argument '{name}'"
+                            arg[0].location, f"Unknown keyword parameter '{name}'"
                         )
                 valid_call = False
                 continue
@@ -414,40 +416,40 @@ class CallDispatcher(Generic[E]):
                 required_positional.remove(name)
             if name in required_keyword:
                 required_keyword.remove(name)
-            set_args.add(name)
+            set_params.add(name)
             mapped.append(
                 MappedArgument(
-                    expr=arg[0],
+                    arg_expr=arg[0],
-                    type=arg[1],
+                    arg_type=arg[1],
-                    argument=param,
+                    parameter=param,
                 )
             )
 
-        def join_args(args: list[str]) -> str:
+        def join_params(params: list[str]) -> str:
-            args = list(map(lambda a: f"'{a}'", args))
+            params = list(map(lambda p: f"'{p}'", params))
-            if len(args) == 0:
+            if len(params) == 0:
                 return ""
-            if len(args) == 1:
+            if len(params) == 1:
-                return args[0]
+                return params[0]
-            return ", ".join(args[:-1]) + " and " + args[-1]
+            return ", ".join(params[:-1]) + " and " + params[-1]
 
         if len(required_positional) != 0:
             plural: str = "" if len(required_positional) == 1 else "s"
-            args: str = join_args(required_positional)
+            params: str = join_params(required_positional)
             if report_errors:
                 self.reporter.error(
                     location,
-                    f"Missing required positional argument{plural}: {args}",
+                    f"Missing required positional argument{plural}: {params}",
                 )
             valid_call = False
 
         if len(required_keyword) != 0:
             plural: str = "" if len(required_keyword) == 1 else "s"
-            args: str = join_args(required_keyword)
+            params: str = join_params(required_keyword)
             if report_errors:
                 self.reporter.error(
                     location,
-                    f"Missing required keyword argument{plural}: {args}",
+                    f"Missing required keyword argument{plural}: {params}",
                 )
             valid_call = False
 
@@ -474,11 +476,11 @@ class CallDispatcher(Generic[E]):
         """
         by_expr: dict[E, Type] = {}
         for arg in mapped1:
-            by_expr[arg.expr] = arg.argument.type
+            by_expr[arg.arg_expr] = arg.parameter.type
 
         for arg in mapped2:
-            type2: Type = arg.argument.type
+            type2: Type = arg.parameter.type
-            type1: Type = by_expr[arg.expr]
+            type1: Type = by_expr[arg.arg_expr]
             if not self.types.is_subtype(type1, type2):
                 return False
         return True

midas/checker/evaluator.py

@@ -158,15 +158,17 @@ class Evaluator(m.Expr.Visitor[Any]):
         return res
 
     def _map_args(self, function: Function, args: list[Any], kwargs: dict[str, Any]):
-        positional: list[Function.Argument] = function.pos_args + function.args
+        positional: list[Function.Parameter] = (
-        keywords: dict[str, Function.Argument] = {
+            function.params.pos + function.params.mixed
-            arg.name: arg for arg in function.args + function.kw_args
+        )
+        keywords: dict[str, Function.Parameter] = {
+            param.name: param for param in function.params.mixed + function.params.kw
         }
 
         for i, arg in enumerate(args):
-            param: Function.Argument = positional[i]
+            param: Function.Parameter = positional[i]
             self.set_value(param.name, arg)
 
         for name, arg in kwargs.items():
-            param: Function.Argument = keywords[name]
+            param: Function.Parameter = keywords[name]
             self.set_value(param.name, arg)

midas/checker/frames/column_groupby_methods.py

@@ -7,7 +7,14 @@ import midas.ast.python as p
 from midas.ast.location import Location
 from midas.checker.dispatcher import CallResult
 from midas.checker.frames.utils import MethodRegistry, method
-from midas.checker.types import ColumnGroupBy, ColumnType, Function, TopType, Type
+from midas.checker.types import (
+    ColumnGroupBy,
+    ColumnType,
+    Function,
+    ParamSpec,
+    TopType,
+    Type,
+)
 
 if TYPE_CHECKING:
     from midas.checker.python import TypedExpr
@@ -38,31 +45,31 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def _aggregate(
         self,
         call: Call,
-        args: list[str | tuple[str, str, bool]] = [],
+        params: list[str | tuple[str, str, bool]] = [],
         *,
         preserve_inner_type: bool = False,
     ) -> Type:
-        real_args: list[Function.Argument] = []
+        real_params: list[Function.Parameter] = []
-        for i, arg in enumerate(args):
+        for i, param in enumerate(params):
-            match arg:
+            match param:
                 case str() as name:
-                    arg = Function.Argument(
+                    param = Function.Parameter(
                         pos=i,
                         name=name,
                         type=self.types.get_type(self.NAMED_ARGS[name]),
                         required=False,
                     )
                 case (name, type, required):
-                    arg = Function.Argument(
+                    param = Function.Parameter(
                         pos=i,
                         name=name,
                         type=self.types.get_type(type),
                         required=required,
                     )
-            real_args.append(arg)
+            real_params.append(param)
 
         signature = Function(
-            args=real_args,
+            params=ParamSpec(mixed=real_params),
             returns=(
                 call.groupby.column
                 if preserve_inner_type

midas/checker/frames/column_methods.py

@@ -13,6 +13,7 @@ from midas.checker.types import (
     ColumnType,
     Function,
     GenericType,
+    ParamSpec,
     TopType,
     Type,
     TypeVar,
@@ -82,14 +83,16 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
             name="add",
             params=[param_type],
             body=Function(
-                args=[
+                params=ParamSpec(
-                    Function.Argument(
+                    mixed=[
-                        pos=0,
+                        Function.Parameter(
-                        name="other",
+                            pos=0,
-                        type=ColumnType(type=param_type),
+                            name="other",
-                        required=True,
+                            type=ColumnType(type=param_type),
-                    ),
+                            required=True,
-                ],
+                        ),
+                    ],
+                ),
                 returns=self._element_binary_op(call, method),
             ),
         )
@@ -163,20 +166,22 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
     def _aggregate(
         self,
         call: Call,
-        kwargs: list[Function.Argument] = [],
+        kwargs: list[Function.Parameter] = [],
         *,
         preserve_inner_type: bool = False,
     ) -> Type:
         signature = Function(
-            kw_args=[
+            params=ParamSpec(
-                Function.Argument(
+                kw=[
-                    pos=0,
+                    Function.Parameter(
-                    name="axis",
+                        pos=0,
-                    type=TopType(),
+                        name="axis",
-                    required=False,
+                        type=TopType(),
-                ),
+                        required=False,
-                *kwargs,
+                    ),
-            ],
+                    *kwargs,
+                ],
+            ),
             returns=call.column if preserve_inner_type else ColumnType(type=TopType()),
         )
 
@@ -221,7 +226,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         return self._aggregate(
             call,
             [
-                Function.Argument(
+                Function.Parameter(
                     pos=1,
                     name="ddof",
                     type=self.types.get_type("int"),
@@ -239,7 +244,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         return self._aggregate(
             call,
             [
-                Function.Argument(
+                Function.Parameter(
                     pos=1,
                     name="var",
                     type=self.types.get_type("int"),
@@ -251,14 +256,16 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
     @method()
     def head(self, call: Call) -> Type:
         signature = Function(
-            args=[
+            params=ParamSpec(
-                Function.Argument(
+                mixed=[
-                    pos=0,
+                    Function.Parameter(
-                    name="n",
+                        pos=0,
-                    type=self.types.get_type("int"),
+                        name="n",
-                    required=False,
+                        type=self.types.get_type("int"),
-                ),
+                        required=False,
-            ],
+                    ),
+                ],
+            ),
             returns=call.column,
         )
 
@@ -273,14 +280,16 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
     @method()
     def tail(self, call: Call) -> Type:
         signature = Function(
-            args=[
+            params=ParamSpec(
-                Function.Argument(
+                mixed=[
-                    pos=0,
+                    Function.Parameter(
-                    name="n",
+                        pos=0,
-                    type=self.types.get_type("int"),
+                        name="n",
-                    required=False,
+                        type=self.types.get_type("int"),
-                ),
+                        required=False,
-            ],
+                    ),
+                ],
+            ),
             returns=call.column,
         )
 
@@ -296,52 +305,33 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
     def groupby(self, call: Call) -> Type:
         bool_: Type = self.types.get_type("bool")
         function: Function = Function(
-            args=[
+            params=ParamSpec(
-                Function.Argument(
+                mixed=[
-                    pos=0,
+                    Function.Parameter(
-                    name="by",
+                        pos=0,
-                    type=TopType(),
+                        name="by",
-                    required=False,
+                        type=TopType(),
-                ),
+                        required=False,
-                Function.Argument(
+                    ),
-                    pos=1,
+                    Function.Parameter(
-                    name="level",
+                        pos=1,
-                    type=TopType(),
+                        name="level",
-                    required=False,
+                        type=TopType(),
-                ),
+                        required=False,
-            ],
+                    ),
-            kw_args=[
+                ],
-                Function.Argument(
+                kw=[
-                    pos=2,
+                    Function.Parameter(
-                    name="as_index",
+                        pos=i + 2,
-                    type=bool_,
+                        name=name,
-                    required=False,
+                        type=bool_,
-                ),
+                        required=False,
-                Function.Argument(
+                    )
-                    pos=3,
+                    for i, name in enumerate(
-                    name="sort",
+                        ["as_index", "sort", "group_keys", "observed", "dropna"]
-                    type=bool_,
+                    )
-                    required=False,
+                ],
-                ),
+            ),
-                Function.Argument(
-                    pos=4,
-                    name="group_keys",
-                    type=bool_,
-                    required=False,
-                ),
-                Function.Argument(
-                    pos=5,
-                    name="observed",
-                    type=bool_,
-                    required=False,
-                ),
-                Function.Argument(
-                    pos=6,
-                    name="dropna",
-                    type=bool_,
-                    required=False,
-                ),
-            ],
             returns=ColumnGroupBy(column=call.column),
         )
 

midas/checker/frames/frame_methods.py

@@ -14,6 +14,7 @@ from midas.checker.types import (
     FrameGroupBy,
     Function,
     OverloadedFunction,
+    ParamSpec,
     TopType,
     Type,
     UnknownType,
@@ -150,14 +151,16 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         # TODO: support scalar, sequence, Series, dict operand
         # Build signature with new schema and generic operand
         signature = Function(
-            args=[
+            params=ParamSpec(
-                Function.Argument(
+                mixed=[
-                    pos=0,
+                    Function.Parameter(
-                    name="other",
+                        pos=0,
-                    type=DataFrameType(columns=[]),
+                        name="other",
-                    required=True,
+                        type=DataFrameType(columns=[]),
-                ),
+                        required=True,
-            ],
+                    ),
+                ],
+            ),
             returns=self._element_binary_op(call, method),
         )
 
@@ -227,29 +230,33 @@ class FrameMethodRegistry(MethodRegistry[Call]):
     def eq(self, call: Call) -> Type:
         return self._element_wise(call, "__eq__")
 
-    def _aggregate(self, call: Call, kwargs: list[Function.Argument] = []) -> Type:
+    def _aggregate(self, call: Call, kwargs: list[Function.Parameter] = []) -> Type:
         with_axis = Function(
-            kw_args=[
+            params=ParamSpec(
-                Function.Argument(
+                kw=[
-                    pos=0,
+                    Function.Parameter(
-                    name="axis",
+                        pos=0,
-                    type=self.types.get_type("int"),
+                        name="axis",
-                    required=False,
+                        type=self.types.get_type("int"),
-                ),
+                        required=False,
-                *kwargs,
+                    ),
-            ],
+                    *kwargs,
+                ],
+            ),
             returns=ColumnType(type=TopType()),
         )
         without_axis = Function(
-            kw_args=[
+            params=ParamSpec(
-                Function.Argument(
+                kw=[
-                    pos=0,
+                    Function.Parameter(
-                    name="axis",
+                        pos=0,
-                    type=self.types.get_type("None"),
+                        name="axis",
-                    required=True,
+                        type=self.types.get_type("None"),
-                ),
+                        required=True,
-                *kwargs,
+                    ),
-            ],
+                    *kwargs,
+                ],
+            ),
             returns=TopType(),
         )
         overload = OverloadedFunction(
@@ -300,7 +307,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         return self._aggregate(
             call,
             [
-                Function.Argument(
+                Function.Parameter(
                     pos=1,
                     name="ddof",
                     type=self.types.get_type("int"),
@@ -318,7 +325,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         return self._aggregate(
             call,
             [
-                Function.Argument(
+                Function.Parameter(
                     pos=1,
                     name="var",
                     type=self.types.get_type("int"),
@@ -330,14 +337,16 @@ class FrameMethodRegistry(MethodRegistry[Call]):
     @method()
     def head(self, call: Call) -> Type:
         signature = Function(
-            args=[
+            params=ParamSpec(
-                Function.Argument(
+                mixed=[
-                    pos=0,
+                    Function.Parameter(
-                    name="n",
+                        pos=0,
-                    type=self.types.get_type("int"),
+                        name="n",
-                    required=False,
+                        type=self.types.get_type("int"),
-                ),
+                        required=False,
-            ],
+                    ),
+                ],
+            ),
             returns=call.frame,
         )
 
@@ -352,14 +361,16 @@ class FrameMethodRegistry(MethodRegistry[Call]):
     @method()
     def tail(self, call: Call) -> Type:
         signature = Function(
-            args=[
+            params=ParamSpec(
-                Function.Argument(
+                mixed=[
-                    pos=0,
+                    Function.Parameter(
-                    name="n",
+                        pos=0,
-                    type=self.types.get_type("int"),
+                        name="n",
-                    required=False,
+                        type=self.types.get_type("int"),
-                ),
+                        required=False,
-            ],
+                    ),
+                ],
+            ),
             returns=call.frame,
         )
 
@@ -375,52 +386,33 @@ class FrameMethodRegistry(MethodRegistry[Call]):
     def groupby(self, call: Call) -> Type:
         bool_: Type = self.types.get_type("bool")
         function: Function = Function(
-            args=[
+            params=ParamSpec(
-                Function.Argument(
+                mixed=[
-                    pos=0,
+                    Function.Parameter(
-                    name="by",
+                        pos=0,
-                    type=TopType(),
+                        name="by",
-                    required=False,
+                        type=TopType(),
-                ),
+                        required=False,
-                Function.Argument(
+                    ),
-                    pos=1,
+                    Function.Parameter(
-                    name="level",
+                        pos=1,
-                    type=TopType(),
+                        name="level",
-                    required=False,
+                        type=TopType(),
-                ),
+                        required=False,
-            ],
+                    ),
-            kw_args=[
+                ],
-                Function.Argument(
+                kw=[
-                    pos=2,
+                    Function.Parameter(
-                    name="as_index",
+                        pos=i + 2,
-                    type=bool_,
+                        name=name,
-                    required=False,
+                        type=bool_,
-                ),
+                        required=False,
-                Function.Argument(
+                    )
-                    pos=3,
+                    for i, name in enumerate(
-                    name="sort",
+                        ["as_index", "sort", "group_keys", "observed", "dropna"]
-                    type=bool_,
+                    )
-                    required=False,
+                ],
-                ),
+            ),
-                Function.Argument(
-                    pos=4,
-                    name="group_keys",
-                    type=bool_,
-                    required=False,
-                ),
-                Function.Argument(
-                    pos=5,
-                    name="observed",
-                    type=bool_,
-                    required=False,
-                ),
-                Function.Argument(
-                    pos=6,
-                    name="dropna",
-                    type=bool_,
-                    required=False,
-                ),
-            ],
             returns=FrameGroupBy(frame=call.frame),
         )
 

midas/checker/midas.py

@@ -21,6 +21,7 @@ from midas.checker.types import (
     ExtensionType,
     Function,
     GenericType,
+    ParamSpec,
     Predicate,
     Type,
     TypeVar,
@@ -32,13 +33,6 @@ 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 ReturnException(Exception):
     pass
 
@@ -47,7 +41,7 @@ class ReturnException(Exception):
 class MappedArgument:
     expr: m.Expr
     type: Type
-    argument: Function.Argument
+    argument: Function.Parameter
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -196,9 +190,7 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
                 self._predicate_params[param.name.lexeme] = param.type.accept(self)
 
         type: Type = self.type_of(stmt.body)
-        params: list[TypedParamSpec] = [
+        params: list[ParamSpec] = [self._visit_param_spec(spec) for spec in stmt.params]
-            self._visit_param_spec(spec) for spec in stmt.params
-        ]
 
         if not self._is_valid_predicate(type):
             self.reporter.error(
@@ -209,9 +201,7 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
             type = self._bool
             for spec in reversed(params):
                 type = Function(
-                    pos_args=spec.pos,
+                    params=spec,
-                    args=spec.mixed,
-                    kw_args=spec.kw,
                     returns=type,
                 )
         self._predicate_params = {}
@@ -386,30 +376,34 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
         )
 
     def visit_function_type(self, type: m.FunctionType) -> Type:
-        params: TypedParamSpec = self._visit_param_spec(type.params)
         return Function(
-            pos_args=params.pos,
+            params=self._visit_param_spec(type.params),
-            args=params.mixed,
-            kw_args=params.kw,
             returns=type.returns.accept(self),
         )
 
-    def _visit_param_spec(self, spec: m.ParamSpec) -> TypedParamSpec:
+    def _visit_param_spec(self, spec: m.ParamSpec) -> ParamSpec:
         n_pos: int = len(spec.pos)
         n_mixed: int = len(spec.mixed)
 
-        def process_arg(arg: m.FunctionType.Argument, i: int) -> Function.Argument:
+        def process_param(
-            return Function.Argument(
+            param: m.FunctionType.Parameter, i: int
+        ) -> Function.Parameter:
+            return Function.Parameter(
                 pos=i,
-                name=arg.name.lexeme if arg.name is not None else str(i),
+                name=param.name.lexeme if param.name is not None else str(i),
-                type=arg.type.accept(self),
+                type=param.type.accept(self),
-                required=arg.required,
+                required=param.required,
             )
 
-        return TypedParamSpec(
+        return ParamSpec(
-            pos=[process_arg(arg, i) for i, arg in enumerate(spec.pos)],
+            pos=[process_param(param, i) for i, param in enumerate(spec.pos)],
-            mixed=[process_arg(arg, i + n_pos) for i, arg in enumerate(spec.mixed)],
+            mixed=[
-            kw=[process_arg(arg, i + n_pos + n_mixed) for i, arg in enumerate(spec.kw)],
+                process_param(param, i + n_pos) for i, param in enumerate(spec.mixed)
+            ],
+            kw=[
+                process_param(param, i + n_pos + n_mixed)
+                for i, param in enumerate(spec.kw)
+            ],
         )
 
     def visit_frame_type(self, type: m.FrameType) -> Type:

midas/checker/preamble.py

@@ -7,6 +7,7 @@ from midas.checker.types import (
     Function,
     GenericType,
     OverloadedFunction,
+    ParamSpec,
     TopType,
     Type,
     TypeVar,
@@ -132,9 +133,9 @@ class Preamble(Environment):
         returns: Type = UnitType(),
         type_vars: list[TypeVar] = [],
     ) -> Type:
-        def map_args(params: list[Param], offset: int) -> list[Function.Argument]:
+        def map_params(params: list[Param], offset: int) -> list[Function.Parameter]:
             return [
-                Function.Argument(
+                Function.Parameter(
                     pos=i + offset,
                     name=param.name,
                     type=param.type,
@@ -144,9 +145,11 @@ class Preamble(Environment):
             ]
 
         function = Function(
-            pos_args=map_args(pos, 0),
+            params=ParamSpec(
-            args=map_args(mixed, len(pos)),
+                pos=map_params(pos, 0),
-            kw_args=map_args(kw, len(pos) + len(mixed)),
+                mixed=map_params(mixed, len(pos)),
+                kw=map_params(kw, len(pos) + len(mixed)),
+            ),
             returns=returns,
         )
         if len(type_vars) != 0:

midas/checker/python.py

@@ -31,6 +31,7 @@ from midas.checker.types import (
     FrameGroupBy,
     Function,
     GenericType,
+    ParamSpec,
     TopType,
     TupleType,
     Type,
@@ -59,7 +60,7 @@ class UndefinedMethodException(Exception):
 class MappedArgument:
     expr: p.Expr
     type: Type
-    argument: Function.Argument
+    argument: Function.Parameter
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -289,61 +290,64 @@ class PythonTyper(
 
     def visit_function(self, stmt: p.Function) -> None:
         env: Environment = Environment(self.env)
-        pos_args: list[Function.Argument] = []
+        pos: list[Function.Parameter] = []
-        args: list[Function.Argument] = []
+        mixed: list[Function.Parameter] = []
-        kw_args: list[Function.Argument] = []
+        kw: list[Function.Parameter] = []
 
-        def eval_arg_type(arg: p.Function.Argument) -> Type:
+        def eval_param_type(param: p.Function.Parameter) -> Type:
-            if arg.type is not None:
+            if param.type is not None:
-                return self.resolve_type_expr(arg.type)
+                return self.resolve_type_expr(param.type)
-            if arg.default is not None:
+            if param.default is not None:
-                return self.type_of(arg.default)
+                return self.type_of(param.default)
             return UnknownType()
 
-        pos: int = 0
+        position: int = 0
-        for arg in stmt.posonlyargs:
+        for param in stmt.params.pos:
-            pos_args.append(
+            pos.append(
-                Function.Argument(
+                Function.Parameter(
-                    pos=pos,
+                    pos=position,
-                    name=arg.name,
+                    name=param.name,
-                    type=eval_arg_type(arg),
+                    type=eval_param_type(param),
-                    required=arg.default is None,
+                    required=param.default is None,
                 )
             )
-            pos += 1
+            position += 1
-        for arg in stmt.args:
+        for param in stmt.params.mixed:
-            args.append(
+            mixed.append(
-                Function.Argument(
+                Function.Parameter(
-                    pos=pos,
+                    pos=position,
-                    name=arg.name,
+                    name=param.name,
-                    type=eval_arg_type(arg),
+                    type=eval_param_type(param),
-                    required=arg.default is None,
+                    required=param.default is None,
                 )
             )
-            pos += 1
+            position += 1
-        for arg in stmt.kwonlyargs:
+        for param in stmt.params.kw:
-            kw_args.append(
+            kw.append(
-                Function.Argument(
+                Function.Parameter(
-                    pos=pos,  # not relevant
+                    pos=position,  # not relevant
-                    name=arg.name,
+                    name=param.name,
-                    type=eval_arg_type(arg),
+                    type=eval_param_type(param),
-                    required=arg.default is None,
+                    required=param.default is None,
                 )
             )
-            pos += 1
+            position += 1
 
-        all_args: list[Function.Argument] = pos_args + args + kw_args
+        param_spec: ParamSpec = ParamSpec(
-        for arg in all_args:
+            pos=pos,
-            env.define(arg.name, arg.type)
+            mixed=mixed,
+            kw=kw,
+        )
+        all_params: list[Function.Parameter] = pos + mixed + kw
+        for param in all_params:
+            env.define(param.name, param.type)
 
         returns_hint: Optional[Type] = None
         if stmt.returns is not None:
             returns_hint = self.resolve_type_expr(stmt.returns)
             # Early define to handle simple fully-typed recursion
             inside_function: Function = Function(
-                pos_args=pos_args,
+                params=param_spec,
-                args=args,
-                kw_args=kw_args,
                 returns=returns_hint,
             )
             self.env.define(stmt.name, inside_function)
@@ -375,13 +379,11 @@ class PythonTyper(
 
         # TODO: handle *args and **kwargs sinks
         function: Type = Function(
-            pos_args=pos_args,
+            params=param_spec,
-            args=args,
-            kw_args=kw_args,
             returns=returns,
         )
         generic_params: list[TypeVar] = []
-        all_types: list[Type] = [arg.type for arg in all_args] + [returns]
+        all_types: list[Type] = [param.type for param in all_params] + [returns]
         for type in all_types:
             if isinstance(type, TypeVar):
                 if type not in generic_params:

midas/checker/registry.py

@@ -234,92 +234,100 @@ class TypesRegistry:
         if not self.is_subtype(func1.returns, func2.returns):
             return False
 
-        pos1: list[Function.Argument] = func1.pos_args
+        pos1: list[Function.Parameter] = func1.params.pos
-        mixed1: list[Function.Argument] = func1.args
+        mixed1: list[Function.Parameter] = func1.params.mixed
-        kw1: dict[str, Function.Argument] = {a.name: a for a in func1.kw_args}
+        kw1: dict[str, Function.Parameter] = {
-        pos2: list[Function.Argument] = func2.pos_args
+            param.name: param for param in func1.params.kw
-        mixed2: list[Function.Argument] = func2.args
+        }
-        kw2: dict[str, Function.Argument] = {a.name: a for a in func2.kw_args}
+        pos2: list[Function.Parameter] = func2.params.pos
+        mixed2: list[Function.Parameter] = func2.params.mixed
+        kw2: dict[str, Function.Parameter] = {
+            param.name: param for param in func2.params.kw
+        }
 
-        mixed_by_pos: dict[int, Function.Argument] = {arg.pos: arg for arg in mixed2}
+        mixed_by_pos: dict[int, Function.Parameter] = {
-        mixed_by_name: dict[str, Function.Argument] = {arg.name: arg for arg in mixed2}
+            param.pos: param for param in mixed2
+        }
+        mixed_by_name: dict[str, Function.Parameter] = {
+            param.name: param for param in mixed2
+        }
 
-        def is_arg_subtype(sub: Function.Argument, sup: Function.Argument) -> bool:
+        def is_arg_subtype(sub: Function.Parameter, sup: Function.Parameter) -> bool:
             if not self.is_subtype(sub.type, sup.type):
                 return False
             if not sup.required and sub.required:
                 return False
             return True
 
-        for arg1 in pos1:
+        for param1 in pos1:
-            arg2: Function.Argument
+            param2: Function.Parameter
-            if arg1.pos < len(pos2):
+            if param1.pos < len(pos2):
-                arg2 = pos2[arg1.pos]
+                param2 = pos2[param1.pos]
-            elif arg1.pos in mixed_by_pos:
+            elif param1.pos in mixed_by_pos:
-                arg2 = mixed_by_pos[arg1.pos]
+                param2 = mixed_by_pos[param1.pos]
-            elif not arg1.required:
+            elif not param1.required:
                 continue
             else:
                 return False
-            if not is_arg_subtype(arg2, arg1):
+            if not is_arg_subtype(param2, param1):
                 return False
 
-        for name, arg1 in kw1.items():
+        for name, param1 in kw1.items():
-            arg2: Function.Argument
+            param2: Function.Parameter
             if name in kw2:
-                arg2 = kw2[name]
+                param2 = kw2[name]
             elif name in mixed_by_name:
-                arg2 = mixed_by_name[name]
+                param2 = mixed_by_name[name]
-            elif not arg1.required:
+            elif not param1.required:
                 continue
             else:
                 return False
-            if not is_arg_subtype(arg2, arg1):
+            if not is_arg_subtype(param2, param1):
                 return False
 
-        for arg1 in mixed1:
+        for param1 in mixed1:
-            pos_arg2: Optional[Function.Argument] = None
+            pos_param2: Optional[Function.Parameter] = None
-            kw_arg2: Optional[Function.Argument] = None
+            kw_param2: Optional[Function.Parameter] = None
-            if arg1.name in kw2:
+            if param1.name in kw2:
-                kw_arg2 = kw2[arg1.name]
+                kw_param2 = kw2[param1.name]
-            elif arg1.name in mixed_by_name:
+            elif param1.name in mixed_by_name:
-                kw_arg2 = mixed_by_name[arg1.name]
+                kw_param2 = mixed_by_name[param1.name]
-            if arg1.pos < len(pos2):
+            if param1.pos < len(pos2):
-                pos_arg2 = pos2[arg1.pos]
+                pos_param2 = pos2[param1.pos]
-            elif arg1.pos in mixed_by_pos:
+            elif param1.pos in mixed_by_pos:
-                pos_arg2 = mixed_by_pos[arg1.pos]
+                pos_param2 = mixed_by_pos[param1.pos]
 
             # No match in func2 and arg is required
-            if pos_arg2 is None and kw_arg2 is None and arg1.required:
+            if pos_param2 is None and kw_param2 is None and param1.required:
                 return False
 
             # Matching keyword argument
-            if kw_arg2 is not None and not is_arg_subtype(kw_arg2, arg1):
+            if kw_param2 is not None and not is_arg_subtype(kw_param2, param1):
                 return False
 
             # Matching positional argument
-            if pos_arg2 is not None and not is_arg_subtype(pos_arg2, arg1):
+            if pos_param2 is not None and not is_arg_subtype(pos_param2, param1):
                 return False
 
-        mixed_positions: set[int] = {a.pos for a in mixed1}
+        mixed_positions: set[int] = {param.pos for param in mixed1}
-        mixed_names: set[str] = {a.name for a in mixed1}
+        mixed_names: set[str] = {param.name for param in mixed1}
-        for arg2 in pos2:
+        for param2 in pos2:
-            if not arg2.required:
+            if not param2.required:
                 continue
-            if arg2.pos >= len(pos1) and arg2.pos not in mixed_positions:
+            if param2.pos >= len(pos1) and param2.pos not in mixed_positions:
                 return False
 
-        for name, arg2 in kw2.items():
+        for name, param2 in kw2.items():
-            if not arg2.required:
+            if not param2.required:
                 continue
             if name not in kw1 and name not in mixed_names:
                 return False
 
-        for arg2 in mixed2:
+        for param2 in mixed2:
-            if arg2.required:
+            if param2.required:
                 continue
-            pos_match: bool = arg2.pos < len(pos1) or arg2.pos in mixed_positions
+            pos_match: bool = param2.pos < len(pos1) or param2.pos in mixed_positions
-            kw_match: bool = arg2.name in kw1 or arg2.name in mixed_names
+            kw_match: bool = param2.name in kw1 or param2.name in mixed_names
             if not pos_match or not kw_match:
                 return False
 

midas/checker/resolver.py

@@ -93,7 +93,7 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
             function (p.Function): the function to resolve
         """
         self.begin_scope()
-        for param in function.all_args:
+        for param in function.params.all:
             self.declare(param.name)
             self.define(param.name)
         self.resolve(*function.body)

midas/checker/types.py

@@ -45,28 +45,14 @@ class UnitType:
 
 @dataclass(frozen=True, kw_only=True)
 class Function:
-    pos_args: list[Argument] = field(default_factory=list)
+    params: ParamSpec
-    args: list[Argument] = field(default_factory=list)
-    kw_args: list[Argument] = field(default_factory=list)
     returns: Type
 
     def __str__(self) -> str:
-        args: list[str] = []
+        return f"{self.params} -> {self.returns}"
-        if len(self.pos_args) != 0:
-            args += list(map(str, self.pos_args))
-            args.append("/")
-
-        if len(self.args) != 0:
-            args += list(map(str, self.args))
-
-        if len(self.kw_args) != 0:
-            args.append("*")
-            args += list(map(str, self.kw_args))
-
-        return f"({', '.join(args)}) -> {self.returns}"
 
     @dataclass(frozen=True, kw_only=True)
-    class Argument:
+    class Parameter:
         pos: int
         name: str
         type: Type
@@ -77,6 +63,28 @@ class Function:
             return f"{self.name}: {self.type}{opt}"
 
 
+@dataclass(frozen=True, kw_only=True)
+class ParamSpec:
+    pos: list[Function.Parameter] = field(default_factory=list)
+    mixed: list[Function.Parameter] = field(default_factory=list)
+    kw: list[Function.Parameter] = field(default_factory=list)
+
+    def __str__(self) -> str:
+        params: list[str] = []
+        if len(self.pos) != 0:
+            params += list(map(str, self.pos))
+            params.append("/")
+
+        if len(self.mixed) != 0:
+            params += list(map(str, self.mixed))
+
+        if len(self.kw) != 0:
+            params.append("*")
+            params += list(map(str, self.kw))
+
+        return f"({', '.join(params)})"
+
+
 @dataclass(frozen=True, kw_only=True)
 class OverloadedFunction:
     overloads: list[Type]
@@ -204,12 +212,19 @@ class ColumnGroupBy:
 
 
 def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
-    def sub_argument(arg: Function.Argument):
+    def sub_parameter(param: Function.Parameter):
-        return Function.Argument(
+        return Function.Parameter(
-            pos=arg.pos,
+            pos=param.pos,
-            name=arg.name,
+            name=param.name,
-            type=substitute_typevars(arg.type, substitutions),
+            type=substitute_typevars(param.type, substitutions),
-            required=arg.required,
+            required=param.required,
+        )
+
+    def sub_param_spec(spec: ParamSpec):
+        return ParamSpec(
+            pos=list(map(sub_parameter, spec.pos)),
+            mixed=list(map(sub_parameter, spec.mixed)),
+            kw=list(map(sub_parameter, spec.kw)),
         )
 
     def sub_column(col: DataFrameType.Column):
@@ -235,15 +250,11 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
             )
 
         case Function(
-            pos_args=pos_args,
+            params=params,
-            args=args,
-            kw_args=kw_args,
             returns=returns,
         ):
             return Function(
-                pos_args=list(map(sub_argument, pos_args)),
+                params=sub_param_spec(params),
-                args=list(map(sub_argument, args)),
-                kw_args=list(map(sub_argument, kw_args)),
                 returns=substitute_typevars(returns, substitutions),
             )
 
@@ -351,14 +362,14 @@ def unfold_type(type: Type) -> Type:
 
 
 def to_annotation(type: Type) -> str:
-    def _args_annotation(func: Function) -> str:
+    def _params_annotation(spec: ParamSpec) -> str:
-        if len(func.kw_args) != 0:
+        if len(spec.kw) != 0:
             return "..."
 
-        args: str = ", ".join(
+        params: str = ", ".join(
-            to_annotation(arg.type) for arg in func.pos_args + func.args
+            to_annotation(param.type) for param in spec.pos + spec.mixed
         )
-        return f"[{args}]"
+        return f"[{params}]"
 
     match type:
         case TopType():
@@ -376,8 +387,8 @@ def to_annotation(type: Type) -> str:
         case UnitType():
             return "None"
 
-        case Function(returns=returns):
+        case Function(params=params, returns=returns):
-            params_annot: str = _args_annotation(type)
+            params_annot: str = _params_annotation(params)
             return f"Callable[{params_annot}, {to_annotation(returns)}]"
 
         case OverloadedFunction():

midas/checker/unifier.py

@@ -8,6 +8,7 @@ from midas.checker.types import (
     DataFrameType,
     Function,
     GenericType,
+    ParamSpec,
     TopType,
     Type,
     TypeVar,
@@ -29,25 +30,26 @@ class Unifier:
         keywords: dict[str, Type],
     ) -> Optional[Type]:
         concrete_func: Function = Function(
-            pos_args=[
+            params=ParamSpec(
-                Function.Argument(
+                pos=[
-                    pos=i,
+                    Function.Parameter(
-                    name=str(i),
+                        pos=i,
-                    type=arg,
+                        name=str(i),
-                    required=True,
+                        type=arg,
-                )
+                        required=True,
-                for i, arg in enumerate(positional)
+                    )
-            ],
+                    for i, arg in enumerate(positional)
-            args=[],
+                ],
-            kw_args=[
+                kw=[
-                Function.Argument(
+                    Function.Parameter(
-                    pos=len(positional) + i,
+                        pos=len(positional) + i,
-                    name=name,
+                        name=name,
-                    type=arg,
+                        type=arg,
-                    required=True,
+                        required=True,
-                )
+                    )
-                for i, (name, arg) in enumerate(keywords.items())
+                    for i, (name, arg) in enumerate(keywords.items())
-            ],
+                ],
+            ),
             returns=TopType(),  # TODO: use expected type
         )
         return self.unify_generic(type, concrete_func, match_return=False)
@@ -125,7 +127,7 @@ class Unifier:
                 return self.match(template_column, concrete_column)
 
             case (Function(), Function()):
-                mapped: list[tuple[Function.Argument, Function.Argument]] = (
+                mapped: list[tuple[Function.Parameter, Function.Parameter]] = (
                     self.map_params(template, concrete)
                 )
                 substitutions: dict[str, Type] = {}
@@ -161,19 +163,23 @@ class Unifier:
 
     def map_params(
         self, func1: Function, func2: Function
-    ) -> list[tuple[Function.Argument, Function.Argument]]:
+    ) -> list[tuple[Function.Parameter, Function.Parameter]]:
-        pos1: list[Function.Argument] = func1.pos_args
+        pos1: list[Function.Parameter] = func1.params.pos
-        mixed1: list[Function.Argument] = func1.args
+        mixed1: list[Function.Parameter] = func1.params.mixed
-        kw1: list[Function.Argument] = func1.kw_args
+        kw1: list[Function.Parameter] = func1.params.kw
 
-        pos2: list[Function.Argument] = func2.pos_args
+        pos2: list[Function.Parameter] = func2.params.pos
-        mixed2: list[Function.Argument] = func2.args
+        mixed2: list[Function.Parameter] = func2.params.mixed
-        kw2: list[Function.Argument] = func2.kw_args
+        kw2: list[Function.Parameter] = func2.params.kw
 
-        mapped: list[tuple[Function.Argument, Function.Argument]] = []
+        mapped: list[tuple[Function.Parameter, Function.Parameter]] = []
 
-        by_pos2: dict[int, Function.Argument] = {arg.pos: arg for arg in pos2 + mixed2}
+        by_pos2: dict[int, Function.Parameter] = {
-        by_name2: dict[str, Function.Argument] = {arg.name: arg for arg in mixed2 + kw2}
+            param.pos: param for param in pos2 + mixed2
+        }
+        by_name2: dict[str, Function.Parameter] = {
+            param.name: param for param in mixed2 + kw2
+        }
 
         for arg1 in pos1:
             if (arg2 := by_pos2.get(arg1.pos)) is not None:

midas/checker/variance.py

@@ -77,14 +77,14 @@ class VarianceInferrer:
         match type:
             # Arguments are negative positions -> flip polarity
             # Return is positive position -> keep polarity
-            case Function(pos_args=pos_args, args=mixed_args, kw_args=kw_args):
+            case Function(params=spec):
-                all_args: list[Function.Argument] = pos_args + mixed_args + kw_args
+                all_params: list[Function.Parameter] = spec.pos + spec.mixed + spec.kw
-                for arg in all_args:
+                for param in all_params:
                     self.walk(
-                        arg.type,
+                        param.type,
                         -polarity,
                         base_name,
-                        path + [f"arg:'{arg.name}'"],
+                        path + [f"param:'{param.name}'"],
                     )
 
                 self.walk(type.returns, polarity, base_name, path + ["return"])
@@ -109,10 +109,10 @@ class VarianceInferrer:
                     Variance.COVARIANT: 1,
                     Variance.CONTRAVARIANT: -1,
                 }
-                for arg, param in zip(args, params):
+                for param, param in zip(args, params):
                     param_polarity: Polarity = polarities[param.variance]
                     self.walk(
-                        arg,
+                        param,
                         cast(Polarity, polarity * param_polarity),
                         base_name,
                         path + [f"applied:'{name}'"],

midas/cli/highlighter.py

@@ -157,15 +157,18 @@ class PythonHighlighter(
 
     def visit_function(self, stmt: p.Function) -> None:
         self.wrap(stmt, "function")
-        for arg in stmt.posonlyargs + stmt.args + stmt.kwonlyargs:
+        self._highlight_param_spec(stmt.params)
-            self._highlight_function_argument(arg)
         for body_stmt in stmt.body:
             body_stmt.accept(self)
 
-    def _highlight_function_argument(self, arg: p.Function.Argument) -> None:
+    def _highlight_param_spec(self, spec: p.ParamSpec) -> None:
-        self.wrap(arg, "argument")
+        for param in spec.all:
-        if arg.type is not None:
+            self._highlight_function_param(param)
-            arg.type.accept(self)
+
+    def _highlight_function_param(self, param: p.Function.Parameter) -> None:
+        self.wrap(param, "parameter")
+        if param.type is not None:
+            param.type.accept(self)
 
     def visit_type_assign(self, stmt: p.TypeAssign) -> None:
         stmt.type.accept(self)

midas/cli/hl_python.css

@@ -23,7 +23,7 @@ span {
         --col: 215, 103, 224;
     }
 
-    &.argument {
+    &.parameter {
         --col: 103, 192, 224;
     }
 }

midas/generator/constraints.py

@@ -5,6 +5,7 @@ import midas.ast.midas as m
 from midas.checker.registry import TypesRegistry
 from midas.checker.types import (
     Function,
+    ParamSpec,
     Predicate,
     Type,
     to_annotation,
@@ -54,16 +55,16 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
                 return expr.accept(self)
             case _:
                 func = Function(
-                    pos_args=[],
+                    params=ParamSpec(
-                    args=[
+                        mixed=[
-                        Function.Argument(
+                            Function.Parameter(
-                            pos=0,
+                                pos=0,
-                            name="_",
+                                name="_",
-                            type=self.types.get_type("Any"),
+                                type=self.types.get_type("Any"),
-                            required=True,
+                                required=True,
-                        )
+                            )
-                    ],
+                        ],
-                    kw_args=[],
+                    ),
                     returns=self.types.get_type("bool"),
                 )
                 alias: str = self.make_alias(None)
@@ -94,28 +95,28 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
             )
         return self.make_func(name, [ast.Return(value=body)], predicate.type)
 
-    def make_args(self, func: Function) -> ast.arguments:
+    def make_args(self, params: ParamSpec) -> ast.arguments:
         return ast.arguments(
             posonlyargs=[
                 ast.arg(
-                    arg=arg.name,
+                    arg=param.name,
-                    annotation=ast.Constant(value=to_annotation(arg.type)),
+                    annotation=ast.Constant(value=to_annotation(param.type)),
                 )
-                for arg in func.pos_args
+                for param in params.pos
             ],
             args=[
                 ast.arg(
-                    arg=arg.name,
+                    arg=param.name,
-                    annotation=ast.Constant(value=to_annotation(arg.type)),
+                    annotation=ast.Constant(value=to_annotation(param.type)),
                 )
-                for arg in func.args
+                for param in params.mixed
             ],
             kwonlyargs=[
                 ast.arg(
-                    arg=arg.name,
+                    arg=param.name,
-                    annotation=ast.Constant(value=to_annotation(arg.type)),
+                    annotation=ast.Constant(value=to_annotation(param.type)),
                 )
-                for arg in func.kw_args
+                for param in params.kw
             ],
             defaults=[],
             kw_defaults=[],
@@ -125,11 +126,11 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
         self, name: str, inner_body: list[ast.stmt], type: Type, level: int = 0
     ) -> ast.stmt:
         match type:
-            case Function(returns=Function()):
+            case Function(params=params, returns=Function()):
                 inner_name: str = f"inner{level}"
                 return ast.FunctionDef(
                     name=name,
-                    args=self.make_args(type),
+                    args=self.make_args(params),
                     body=[
                         self.make_func(inner_name, inner_body, type.returns, level + 1),
                         ast.Return(value=ast.Name(id=inner_name)),
@@ -138,10 +139,10 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
                     decorator_list=[],
                 )
 
-            case Function():
+            case Function(params=params):
                 return ast.FunctionDef(
                     name=name,
-                    args=self.make_args(type),
+                    args=self.make_args(params),
                     body=inner_body,
                     returns=ast.Constant(value=to_annotation(type.returns)),
                     decorator_list=[],

midas/generator/generator.py

@@ -250,25 +250,26 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
             value=self.convert(stmt.expr),
         )
 
+    def make_args(self, params: p.ParamSpec) -> ast.arguments:
+        return ast.arguments(
+            posonlyargs=[ast.arg(arg=param.name) for param in params.pos],
+            args=[ast.arg(arg=param.name) for param in params.mixed],
+            kwonlyargs=[ast.arg(arg=param.name) for param in params.kw],
+            defaults=[
+                self.convert(param.default)
+                for param in params.pos + params.mixed
+                if param.default is not None
+            ],
+            kw_defaults=[
+                self.convert(param.default) if param.default is not None else None
+                for param in params.kw
+            ],
+        )
+
     def visit_function(self, stmt: p.Function) -> ast.stmt:
         return ast.FunctionDef(
             name=stmt.name,
-            args=ast.arguments(
+            args=self.make_args(stmt.params),
-                posonlyargs=[ast.arg(arg=arg.name) for arg in stmt.posonlyargs],
-                vararg=None,
-                args=[ast.arg(arg=arg.name) for arg in stmt.args],
-                kwonlyargs=[ast.arg(arg=arg.name) for arg in stmt.kwonlyargs],
-                kwarg=None,
-                defaults=[
-                    self.convert(arg.default)
-                    for arg in stmt.posonlyargs + stmt.args
-                    if arg.default is not None
-                ],
-                kw_defaults=[
-                    self.convert(arg.default) if arg.default is not None else None
-                    for arg in stmt.kwonlyargs
-                ],
-            ),
             body=self._visit_body(stmt.body),
             decorator_list=[],
         )

midas/generator/stubs.py

@@ -17,6 +17,7 @@ from midas.checker.types import (
     Function,
     GenericType,
     OverloadedFunction,
+    ParamSpec,
     TopType,
     TupleType,
     Type,
@@ -328,7 +329,7 @@ class StubsGenerator:
                 return [
                     ast.FunctionDef(
                         name=name,
-                        args=self.dump_args(method, with_self=True),
+                        args=self.dump_params(method.params, with_self=True),
                         returns=self.dump_type(method.returns),
                         body=[ast.Expr(value=Empty)],
                         decorator_list=[ast.Name(id="overload")] if overloaded else [],
@@ -348,24 +349,33 @@ class StubsGenerator:
                     )
                 ]
 
-    def dump_args(self, func: Function, with_self: bool = False) -> ast.arguments:
+    def dump_params(self, params: ParamSpec, with_self: bool = False) -> ast.arguments:
         pos: list[ast.arg] = [
-            ast.arg(arg=f"_{arg.pos}", annotation=self.dump_type(arg.type))
+            ast.arg(
-            for arg in func.pos_args
+                arg=f"_{param.pos}",
+                annotation=self.dump_type(param.type),
+            )
+            for param in params.pos
         ]
         mixed: list[ast.arg] = [
-            ast.arg(arg=arg.name, annotation=self.dump_type(arg.type))
+            ast.arg(
-            for arg in func.args
+                arg=param.name,
+                annotation=self.dump_type(param.type),
+            )
+            for param in params.mixed
         ]
         kw: list[ast.arg] = [
-            ast.arg(arg=arg.name, annotation=self.dump_type(arg.type))
+            ast.arg(
-            for arg in func.kw_args
+                arg=param.name,
+                annotation=self.dump_type(param.type),
+            )
+            for param in params.kw
         ]
         defaults: list[ast.expr] = [
-            Empty for arg in func.pos_args + func.args if not arg.required
+            Empty for param in params.pos + params.mixed if not param.required
         ]
         kw_defaults: list[Optional[ast.expr]] = [
-            None if arg.required else Empty for arg in func.kw_args
+            None if param.required else Empty for param in params.kw
         ]
         if with_self:
             arg = ast.arg(arg="self", annotation=None)
@@ -391,7 +401,7 @@ class StubsGenerator:
             body=[
                 ast.FunctionDef(
                     name="__call__",
-                    args=self.dump_args(func, with_self=True),
+                    args=self.dump_params(func.params, with_self=True),
                     returns=self.dump_type(func.returns),
                     body=[ast.Expr(value=Empty)],
                     decorator_list=[],

midas/parser/midas.py

@@ -587,7 +587,7 @@ class MidasParser(Parser):
 
         params: list[ParamSpec] = []
         while self.check(TokenType.LEFT_PAREN):
-            params.append(self.function_args())
+            params.append(self.function_params())
 
         self.consume(TokenType.EQUAL, "Expected '=' after predicate subject")
         body: Expr = self.constraint()
@@ -599,7 +599,7 @@ class MidasParser(Parser):
         )
 
     def function(self) -> FunctionType:
-        params: ParamSpec = self.function_args()
+        params: ParamSpec = self.function_params()
 
         self.consume(TokenType.ARROW, "Expected '->' before result type")
         result: Type = self.type_expr()
@@ -610,36 +610,38 @@ class MidasParser(Parser):
             returns=result,
         )
 
-    def function_args(self) -> ParamSpec:
+    def function_params(self) -> ParamSpec:
         l_paren: Token = self.consume(
             TokenType.LEFT_PAREN, "Expected '(' before function parameters"
         )
-        pos_args: list[FunctionType.Argument] = []
+        pos: list[FunctionType.Parameter] = []
-        args: list[FunctionType.Argument] = []
+        mixed: list[FunctionType.Parameter] = []
-        kw_args: list[FunctionType.Argument] = []
+        kw: list[FunctionType.Parameter] = []
 
-        args_first_tokens: list[Token] = []
+        mixed_first_tokens: list[Token] = []
 
         section: int = 0
         while not self.is_at_end() and not self.check(TokenType.RIGHT_PAREN):
             match section:
                 case 0 if self.match(TokenType.SLASH):
-                    pos_args = args
+                    pos = mixed
-                    args = []
+                    mixed = []
-                    args_first_tokens = []
+                    mixed_first_tokens = []
                     section = 1
                 case 0 | 1 if self.match(TokenType.STAR):
                     section = 2
                 case _:
-                    # Record first token of mixed argument for errors if unnamed
+                    # Record first token of mixed parameters for errors if unnamed
                     if section != 2:
-                        args_first_tokens.append(self.peek())
+                        mixed_first_tokens.append(self.peek())
 
                     name: Optional[Token] = None
                     if section == 2:
-                        name = self.consume_identifier("Expected keyword argument name")
+                        name = self.consume_identifier(
+                            "Expected keyword parameter name"
+                        )
                         self.consume(
-                            TokenType.COLON, "Expected ':' after argument name"
+                            TokenType.COLON, "Expected ':' after parameter name"
                         )
                     elif self.check_identifier() and self.check_next(TokenType.COLON):
                         name = self.advance()
@@ -647,24 +649,24 @@ class MidasParser(Parser):
 
                     type: Type = self.type_expr()
                     optional: bool = self.match(TokenType.QMARK)
-                    arg = FunctionType.Argument(
+                    param = FunctionType.Parameter(
                         location=None,
                         name=name,
                         type=type,
                         required=not optional,
                     )
                     if section == 2:
-                        kw_args.append(arg)
+                        kw.append(param)
                     else:
-                        args.append(arg)
+                        mixed.append(param)
 
             if not self.match(TokenType.COMMA):
                 break
 
-        for arg, token in zip(args, args_first_tokens):
+        for param, token in zip(mixed, mixed_first_tokens):
-            if arg.name is None:
+            if param.name is None:
                 # Not raised because we can keep parsing
-                self.error(token, "Unnamed mixed argument")
+                self.error(token, "Unnamed mixed parameter")
 
         self.consume(TokenType.RIGHT_PAREN, "Expected ')' after function parameters")
-        return ParamSpec(l_paren=l_paren, pos=pos_args, mixed=args, kw=kw_args)
+        return ParamSpec(l_paren=l_paren, pos=pos, mixed=mixed, kw=kw)

midas/parser/python.py

@@ -23,6 +23,7 @@ from midas.ast.python import (
     LiteralExpr,
     LogicalExpr,
     MidasType,
+    ParamSpec,
     RawExpr,
     RawStmt,
     ReturnStmt,
@@ -212,27 +213,10 @@ class PythonParser:
         match node:
             case ast.FunctionDef(
                 name=name,
-                args=ast.arguments(
+                args=args,
-                    posonlyargs=posonlyargs,
-                    args=args,
-                    vararg=sink,
-                    kwonlyargs=kwonlyargs,
-                    kwarg=kw_sink,
-                    defaults=defaults,
-                    kw_defaults=kw_defaults,
-                ),
                 returns=returns,
                 body=raw_body,
             ):
-
-                def parse_args(
-                    args_list: list[ast.arg], defaults: list[Optional[Expr]]
-                ) -> list[Function.Argument]:
-                    return [
-                        self._parse_function_argument(arg, default)
-                        for arg, default in zip(args_list, defaults)
-                    ]
-
                 body: list[Stmt] = []
                 for stmt in raw_body:
                     stmts = self.parse_stmt(stmt)
@@ -241,54 +225,58 @@ class PythonParser:
                     elif stmts is not None:
                         body.extend(stmts)
 
-                parsed_defaults: list[Optional[Expr]] = [
-                    self.parse_expr(default) for default in defaults
-                ]
-                n_posargs: int = len(posonlyargs)
-                n_args: int = len(args)
-                n_all_posargs = n_posargs + n_args
-                parsed_defaults = [
-                    None,
-                ] * (n_all_posargs - len(defaults)) + parsed_defaults
-
-                posargs_defaults: list[Optional[Expr]] = parsed_defaults[:n_posargs]
-                args_defaults: list[Optional[Expr]] = parsed_defaults[n_posargs:]
-                kwargs_defaults: list[Optional[Expr]] = [
-                    self.parse_expr(default) if default is not None else None
-                    for default in kw_defaults
-                ]
-
                 return Function(
                     location=loc,
                     name=name,
-                    posonlyargs=parse_args(posonlyargs, posargs_defaults),
+                    params=self._parse_param_spec(args),
-                    args=parse_args(args, args_defaults),
-                    sink=(
-                        self._parse_function_argument(sink, None)
-                        if sink is not None
-                        else None
-                    ),
-                    kwonlyargs=parse_args(kwonlyargs, kwargs_defaults),
-                    kw_sink=(
-                        self._parse_function_argument(kw_sink, None)
-                        if kw_sink is not None
-                        else None
-                    ),
                     returns=self._parse_type(returns) if returns is not None else None,
                     body=body,
                 )
             case _:
                 print(f"Unsupported function definition: {ast.unparse(node)}")
 
-    def _parse_function_argument(
+    def _parse_param_spec(self, args: ast.arguments) -> ParamSpec:
+        def parse_params(
+            args_list: list[ast.arg], defaults: list[Optional[Expr]]
+        ) -> list[Function.Parameter]:
+            return [
+                self._parse_function_parameter(arg, default)
+                for arg, default in zip(args_list, defaults)
+            ]
+
+        defaults: list[ast.expr] = args.defaults
+        parsed_defaults: list[Optional[Expr]] = [
+            self.parse_expr(default) for default in defaults
+        ]
+        n_pos: int = len(args.posonlyargs)
+        n_mixed: int = len(args.args)
+        n_all_pos = n_pos + n_mixed
+        parsed_defaults = [
+            None,
+        ] * (n_all_pos - len(defaults)) + parsed_defaults
+
+        pos_defaults: list[Optional[Expr]] = parsed_defaults[:n_pos]
+        mixed_defaults: list[Optional[Expr]] = parsed_defaults[n_pos:]
+        kw_defaults: list[Optional[Expr]] = [
+            self.parse_expr(default) if default is not None else None
+            for default in args.kw_defaults
+        ]
+
+        return ParamSpec(
+            pos=parse_params(args.posonlyargs, pos_defaults),
+            mixed=parse_params(args.args, mixed_defaults),
+            kw=parse_params(args.kwonlyargs, kw_defaults),
+        )
+
+    def _parse_function_parameter(
         self, arg: ast.arg, default: Optional[Expr]
-    ) -> Function.Argument:
+    ) -> Function.Parameter:
         loc: Location = Location.from_ast(arg)
         name: str = arg.arg
         type: Optional[MidasType] = None
         if arg.annotation is not None:
             type = self._parse_type(arg.annotation)
-        return Function.Argument(
+        return Function.Parameter(
             location=loc,
             name=name,
             type=type,

tests/serializer/midas.py

@@ -188,16 +188,16 @@ class MidasAstJsonSerializer(
     def _serialize_param_spec(self, spec: ParamSpec) -> dict:
         return {
             "_type": "ParamSpec",
-            "pos": [self._serialize_func_arg(arg) for arg in spec.pos],
+            "pos": [self._serialize_func_param(arg) for arg in spec.pos],
-            "mixed": [self._serialize_func_arg(arg) for arg in spec.mixed],
+            "mixed": [self._serialize_func_param(arg) for arg in spec.mixed],
-            "kw": [self._serialize_func_arg(arg) for arg in spec.kw],
+            "kw": [self._serialize_func_param(arg) for arg in spec.kw],
         }
 
-    def _serialize_func_arg(self, arg: FunctionType.Argument) -> dict:
+    def _serialize_func_param(self, param: FunctionType.Parameter) -> dict:
         return {
-            "name": arg.name.lexeme if arg.name is not None else None,
+            "name": param.name.lexeme if param.name is not None else None,
-            "type": arg.type.accept(self),
+            "type": param.type.accept(self),
-            "required": arg.required,
+            "required": param.required,
         }
 
     def visit_extension_type(self, type: ExtensionType) -> dict:

tests/serializer/python.py

@@ -22,6 +22,7 @@ from midas.ast.python import (
     LiteralExpr,
     LogicalExpr,
     MidasType,
+    ParamSpec,
     Pass,
     RawExpr,
     RawStmt,
@@ -128,32 +129,30 @@ class PythonAstJsonSerializer(
             "expr": stmt.expr.accept(self),
         }
 
-    def _serialize_argument(self, arg: Function.Argument) -> dict:
-        return {
-            "name": arg.name,
-            "type": self._serialize_optional(arg.type),
-            "default": self._serialize_optional(arg.default),
-        }
-
     def visit_function(self, stmt: Function) -> dict:
         return {
             "_type": "Function",
             "name": stmt.name,
-            "posonlyargs": [self._serialize_argument(arg) for arg in stmt.posonlyargs],
+            "params": self._serialize_param_spec(stmt.params),
-            "args": [self._serialize_argument(arg) for arg in stmt.args],
-            "sink": (
-                self._serialize_argument(stmt.sink) if stmt.sink is not None else None
-            ),
-            "kwonlyargs": [self._serialize_argument(arg) for arg in stmt.kwonlyargs],
-            "kw_sink": (
-                self._serialize_argument(stmt.kw_sink)
-                if stmt.kw_sink is not None
-                else None
-            ),
             "returns": self._serialize_optional(stmt.returns),
             "body": self._serialize_list(stmt.body),
         }
 
+    def _serialize_param_spec(self, spec: ParamSpec) -> dict:
+        return {
+            "_type": "ParamSpec",
+            "pos": [self._serialize_func_param(arg) for arg in spec.pos],
+            "mixed": [self._serialize_func_param(arg) for arg in spec.mixed],
+            "kw": [self._serialize_func_param(arg) for arg in spec.kw],
+        }
+
+    def _serialize_func_param(self, param: Function.Parameter) -> dict:
+        return {
+            "name": param.name,
+            "type": self._serialize_optional(param.type),
+            "default": self._serialize_optional(param.default),
+        }
+
     def visit_type_assign(self, stmt: TypeAssign) -> dict:
         return {
             "_type": "TypeAssign",