diff --git a/gen/gen.py b/gen/gen.py index 50c9c9d..9ea59ee 100644 --- a/gen/gen.py +++ b/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 diff --git a/gen/midas.py b/gen/midas.py index fed8626..fb90199 100644 --- a/gen/midas.py +++ b/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] - mixed: list[FunctionType.Argument] - kw: list[FunctionType.Argument] + pos: list[FunctionType.Parameter] + mixed: list[FunctionType.Parameter] + 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 diff --git a/gen/python.py b/gen/python.py index 73f64c7..60a92f7 100644 --- a/gen/python.py +++ b/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] - args: list[Argument] - sink: Optional[Argument] - kwonlyargs: list[Argument] - kw_sink: Optional[Argument] + params: ParamSpec 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 diff --git a/midas/ast/location.py b/midas/ast/location.py index 47fe360..4d1d7bf 100644 --- a/midas/ast/location.py +++ b/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, diff --git a/midas/ast/midas.py b/midas/ast/midas.py index 128b1c0..eae434a 100644 --- a/midas/ast/midas.py +++ b/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] - mixed: list[FunctionType.Argument] - kw: list[FunctionType.Argument] + pos: list[FunctionType.Parameter] + mixed: list[FunctionType.Parameter] + 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 diff --git a/midas/ast/printer.py b/midas/ast/printer.py deleted file mode 100644 index 95c4109..0000000 --- a/midas/ast/printer.py +++ /dev/null @@ -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)}") diff --git a/midas/ast/printer/__init__.py b/midas/ast/printer/__init__.py new file mode 100644 index 0000000..8b89f7a --- /dev/null +++ b/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 diff --git a/midas/ast/printer/base.py b/midas/ast/printer/base.py new file mode 100644 index 0000000..e6d9b37 --- /dev/null +++ b/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) diff --git a/midas/ast/printer/midas.py b/midas/ast/printer/midas.py new file mode 100644 index 0000000..cbb62d7 --- /dev/null +++ b/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)}" diff --git a/midas/ast/printer/midas_ast.py b/midas/ast/printer/midas_ast.py new file mode 100644 index 0000000..69a0b16 --- /dev/null +++ b/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) diff --git a/midas/ast/printer/python_ast.py b/midas/ast/printer/python_ast.py new file mode 100644 index 0000000..b51c8f9 --- /dev/null +++ b/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)}") diff --git a/midas/ast/python.py b/midas/ast/python.py index 6202fcc..65e896c 100644 --- a/midas/ast/python.py +++ b/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] - args: list[Argument] - sink: Optional[Argument] - kwonlyargs: list[Argument] - kw_sink: Optional[Argument] + params: ParamSpec 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) diff --git a/midas/checker/builtins.py b/midas/checker/builtins.py index be6a781..427d414 100644 --- a/midas/checker/builtins.py +++ b/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): diff --git a/midas/checker/checker.py b/midas/checker/checker.py index c64a282..a305ad1 100644 --- a/midas/checker/checker.py +++ b/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() diff --git a/midas/checker/diagnostic.py b/midas/checker/diagnostic.py index 233dd74..33b25e3 100644 --- a/midas/checker/diagnostic.py +++ b/midas/checker/diagnostic.py @@ -14,6 +14,12 @@ 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 + """ + file_path: Optional[str] location: Location type: DiagnosticType @@ -21,6 +27,18 @@ class Diagnostic: @property def location_str(self) -> str: + """Get diagnostic type and location as a human readable string + + The location is formatted as " in from L: to :", + 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.
+ If the location's end is not specified, the formulation "at L:" is used. + + Returns: + str: the formatted type and location string + """ + start_loc: str = f"L{self.location.lineno}:{self.location.col_offset+1}" end_loc: Optional[str] = "" if ( diff --git a/midas/checker/dispatcher.py b/midas/checker/dispatcher.py index 0a421d3..fef6b04 100644 --- a/midas/checker/dispatcher.py +++ b/midas/checker/dispatcher.py @@ -26,22 +26,29 @@ class HasLocation(Protocol): E = TypeVar("E", bound=HasLocation) TypedExpr = tuple[E, Type] +"""An expression and its type""" @dataclass(frozen=True, kw_only=True) class MappedArgument(Generic[E]): - expr: E - type: Type - argument: Function.Argument + """An argument passed in a call and the corresponding parameter""" + + arg_expr: E + arg_type: Type + parameter: Function.Parameter @dataclass(frozen=True, kw_only=True) class OverloadCandidate: + """An overloaded function call candidate with its mapped arguments""" + function: Function mapped: list[MappedArgument] class CallError(StrEnum): + """Reason of a call error""" + INVALID_ARGS = "Invalid arguments" NO_MATCHING_OVERLOAD = "No matching overload" IMPOSSIBLE_UNIFICATION = "Parameters unification failed" @@ -50,16 +57,28 @@ class CallError(StrEnum): @dataclass(frozen=True, kw_only=True) class CallResult: + """The result of a function call + + Holds a return type, an optional error reason and message + """ + error: Optional[CallError] = None + """The reason of the error, if there is one""" + result: Type = UnknownType() + """The result type. `UnknownType()` if the call is invalid""" + message: Optional[str] = None + """An optional error message""" @property def is_valid(self) -> bool: + """Whether the call is valid (i.e. no error)""" return self.error is None @property def error_message(self) -> str: + """A descriptive message for the error, if there is one""" if self.message is not None: return self.message if self.error is not None: @@ -68,6 +87,15 @@ class CallResult: class CallDispatcher(Generic[E]): + """Helper class to handle dispatching calls and mapping arguments + + This class is responsible for mapping call-site arguments to function + parameters, verifying the validity of calls and computing their + return types + + :class:`CallDispatcher` is generic to handle AST nodes from both Midas and Python + """ + def __init__(self, types: TypesRegistry, reporter: FileReporter) -> None: self.types: TypesRegistry = types self.reporter: FileReporter = reporter @@ -86,22 +114,21 @@ class CallDispatcher(Generic[E]): ) -> CallResult: """Get the result type of a function call - If the function has overloads, the function will try to resolve the + If the callee has overloads, this function will try to resolve the appropriate signature. - Argument types are matched to the defined parameters. - The function doesn't take the raw expression as a parameter to accommodate - for desugared calls such as for operators. + Argument types are matched with the defined parameters. + This function doesn't take the raw expression as a parameter to + accommodate for desugared calls such as for operators. Args: location (Location): the call location callee (Type): the called function - positional (list[TypedExpr]): the list positional arguments + positional (list[TypedExpr]): the list of positional arguments keywords (dict[str, TypedExpr]): the map of keyword arguments report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True. Returns: - Type: the return type of the call, or `None` if either - the call is invalid or no overload matched the arguments uniquely + CallResult: the call result, either a type or an error """ match callee: case Function() as function: @@ -179,6 +206,18 @@ class CallDispatcher(Generic[E]): positional: list[TypedExpr[E]], keywords: dict[str, TypedExpr[E]], ) -> Union[tuple[Function, None], tuple[None, CallError]]: + """Unwrap a type to get a callable `Function` + + Args: + callee (Type): the called type + positional (list[TypedExpr[E]]): the list of positional arguments + keywords (dict[str, TypedExpr[E]]): the map of keyword arguments + + Returns: + Union[tuple[Function, None], tuple[None, CallError]]: a tuple + containing the callable `Function` type, or `None` if it could + not be unwrapped, and an error, or `None` if there was none. + """ match callee: case DerivedType(type=base): return self._unwrap_function(base, positional, keywords) @@ -219,11 +258,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, - f"Wrong type for argument '{arg.argument.name}', expected {arg.argument.type}, got {arg.type}", + arg.arg_expr.location, + f"Wrong type for argument '{arg.parameter.name}', expected {arg.parameter.type}, got {arg.arg_type}", ) valid = False return valid @@ -246,8 +285,9 @@ class CallDispatcher(Generic[E]): report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True. Returns: - Optional[Function]: the resolved function signature if it can be - determined unambiguously, or `None`. + Union[tuple[Function, None], tuple[None, str]]: a tuple containing + the resolved function signature if it can be determined + unambiguously, or `None`, and an error message, or `None` """ candidates: list[OverloadCandidate] = [] errors: list[CallError] = [] @@ -345,30 +385,32 @@ class CallDispatcher(Generic[E]): Returns: tuple[bool, list[MappedArgument]]: a boolean reporting whether - the call is valid and the list of mapped arguments + the call is valid and the list of mapped arguments """ - set_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) - mixed_params: list[Function.Argument] = list(function.args) - kw_params: dict[str, Function.Argument] = { - arg.name: arg for arg in function.kw_args + pos_params: list[Function.Parameter] = list(function.params.pos) + mixed_params: list[Function.Parameter] = list(function.params.mixed) + kw_params: dict[str, Function.Parameter] = { + 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 +427,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], - type=arg[1], - argument=param, + arg_expr=arg[0], + arg_type=arg[1], + 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 +456,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], - type=arg[1], - argument=param, + arg_expr=arg[0], + arg_type=arg[1], + parameter=param, ) ) - def join_args(args: list[str]) -> str: - args = list(map(lambda a: f"'{a}'", args)) - if len(args) == 0: + def join_params(params: list[str]) -> str: + params = list(map(lambda p: f"'{p}'", params)) + if len(params) == 0: return "" - if len(args) == 1: - return args[0] - return ", ".join(args[:-1]) + " and " + args[-1] + if len(params) == 1: + return params[0] + 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 @@ -462,8 +504,8 @@ class CallDispatcher(Generic[E]): of `mapped2`. If any of the parameter type in `mapped1` is not a subtype of the corresponding parameter in `mapped2`, `False` is returned. - This is used to check whether a given overload is - a more specific function/ a subtype of another. + This is used to check whether a given overload is a more specific + function / a subtype of another. Args: mapped1 (list[MappedArgument]): the first argument mappings (subtype) @@ -474,11 +516,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 - type1: Type = by_expr[arg.expr] + type2: Type = arg.parameter.type + type1: Type = by_expr[arg.arg_expr] if not self.types.is_subtype(type1, type2): return False return True diff --git a/midas/checker/evaluator.py b/midas/checker/evaluator.py index c21f1c2..f7b946f 100644 --- a/midas/checker/evaluator.py +++ b/midas/checker/evaluator.py @@ -11,10 +11,18 @@ from midas.lexer.token import TokenType @dataclass(frozen=True, kw_only=True) class PartialPredicate(Predicate): + """A partially applied predicate""" + scope: dict[str, Any] + """A dictionary of already applied parameters""" class Evaluator(m.Expr.Visitor[Any]): + """Helper class to evaluate an expression + + This class is used to evaluate constraint types on literals at compile-time. + """ + def __init__(self, types: TypesRegistry, reporter: Optional[FileReporter] = None): self.types: TypesRegistry = types self.reporter: Optional[FileReporter] = reporter @@ -22,16 +30,51 @@ class Evaluator(m.Expr.Visitor[Any]): self.scopes: list[dict[str, Any]] = [{}] def evaluate(self, expr: m.Expr) -> Any: + """Evaluate the given expression + + Args: + expr (m.Expr): the expression to evaluate + + Returns: + Any: the value of the expression + """ value: Any = expr.accept(self) if self.reporter is not None: self.reporter.debug(expr.location, f"Value: {value}") return value def get_value(self, name: str) -> Any: + """Get the value of a variable in the current scope + + Args: + name (str): the name of the variable + + Raises: + KeyError: if the variable is not defined + + Returns: + Any: the value of the variable + """ scope: dict[str, Any] = self.scopes[-1] return scope[name] def set_value(self, name: str, value: Any, force_declare: bool = False): + """Set the value of a variable + + If `force_declare` is `False`, this function first tries to find the + closest scope in which the variable is defined and assign the value in + that scope, if it can find one. + + If `force_declare` is `True` or if the variable is not defined in any + scope, it is declare and assigned in the current scope + + Args: + name (str): the name of the variable + value (Any): the value of the variable + force_declare (bool, optional): if `False` and the variable is + defined in a scope, the value is assigned in that scope (the + closest if there are multiple declarations). Defaults to False. + """ if not force_declare: for scope in reversed(self.scopes): if name in scope: @@ -131,8 +174,21 @@ class Evaluator(m.Expr.Visitor[Any]): return self.get_value("_") def _evaluate_predicate( - self, predicate: Predicate, args: list[Any], kwargs: dict[str, Any] + self, + predicate: Predicate, + args: list[Any], + kwargs: dict[str, Any], ) -> Any: + """Evaluate a predicate function call + + Args: + predicate (Predicate): the predicate to evaluate + args (list[Any]): a list of positional arguments + kwargs (dict[str, Any]): a map of keyword arguments + + Returns: + Any: the value returned by the predicate call + """ res: Any = None if isinstance(predicate, PartialPredicate): self.scopes.append(predicate.scope) @@ -158,15 +214,27 @@ 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 - keywords: dict[str, Function.Argument] = { - arg.name: arg for arg in function.args + function.kw_args + """Map call arguments to a function's parameters and set their values in context + + Each argument is mapped to a parameter of the function, then its value + is set in the context using :func:`set_value` with the parameter's name + + Args: + function (Function): the called function + args (list[Any]): a list of positional arguments + kwargs (dict[str, Any]): a map of keyword arguments + """ + positional: list[Function.Parameter] = ( + function.params.pos + function.params.mixed + ) + 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) diff --git a/midas/checker/frames/column_groupby_methods.py b/midas/checker/frames/column_groupby_methods.py index 7470f46..f84d8df 100644 --- a/midas/checker/frames/column_groupby_methods.py +++ b/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 @@ -15,6 +22,8 @@ if TYPE_CHECKING: @dataclass(frozen=True, kw_only=True) class Call: + """A column group-by method call, implements :class:`utils.MethodCall`""" + location: Location call_expr: p.Expr groupby: ColumnGroupBy @@ -28,6 +37,8 @@ class Call: class ColumnGroupByMethodRegistry(MethodRegistry[Call]): + """The method registry for column group-by types""" + NAMED_ARGS: dict[str, str] = { "numeric_only": "bool", "skipna": "bool", @@ -38,31 +49,46 @@ 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] = [] - for i, arg in enumerate(args): - match arg: + """Compute the result type of an aggregate method call + + Args: + call (Call): the call object + params (list[str | tuple[str, str, bool], optional): a list of extra + mixed parameters. The list can contain strings to include + parameters predefined in `NAMED_ARGS`, or tuples containing the + parameter's name, type and required flag. Defaults to []. + preserve_inner_type (bool, optional): If `True`, the result type + will preserve the column's inner type (e.g. for `min`/`max`), + otherwise the inner type is widened to `TopType`. Defaults to False. + + Returns: + Type: the result type + """ + real_params: list[Function.Parameter] = [] + for i, param in enumerate(params): + 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 diff --git a/midas/checker/frames/column_manager.py b/midas/checker/frames/column_manager.py index f62d025..8adeb1b 100644 --- a/midas/checker/frames/column_manager.py +++ b/midas/checker/frames/column_manager.py @@ -15,6 +15,8 @@ if TYPE_CHECKING: class ColumnManager: + """Helper class to handle methods and subscripts on column types""" + def __init__(self, typer: PythonTyper) -> None: self.typer: PythonTyper = typer self.method_resolver: ColumnMethodRegistry = ColumnMethodRegistry(self.typer) @@ -32,6 +34,20 @@ class ColumnManager: positional: list[TypedExpr], keywords: dict[str, TypedExpr], ) -> Type: + """Compute the result type of a column's method call + + Args: + method (str): the method name + location (Location): the call's location + call_expr (p.Expr): the call expression + column (ColumnType): the column type + column_expr (p.Expr): the column expression + positional (list[TypedExpr]): the list of positional arguments + keywords (dict[str, TypedExpr]): the map of keyword arguments + + Returns: + Type: the result type + """ call: Call = Call( location=location, call_expr=call_expr, @@ -52,6 +68,20 @@ class ColumnManager: positional: list[TypedExpr], keywords: dict[str, TypedExpr], ) -> Type: + """Compute the result type of a column group-by's method call + + Args: + method (str): the method name + location (Location): the call's location + call_expr (p.Expr): the call expression + groupby (ColumnGroupBy): the column group-by object + groupby_expr (p.Expr): the column group-by expression + positional (list[TypedExpr]): the list of positional arguments + keywords (dict[str, TypedExpr]): the map of keyword arguments + + Returns: + Type: the result type + """ call: GroupByCall = GroupByCall( location=location, call_expr=call_expr, @@ -63,6 +93,15 @@ class ColumnManager: return self.groupby_method_resolver.call(method, call) def get_attribute(self, column: ColumnType, name: str) -> Optional[Type]: + """Get the type of a column's attribute + + Args: + column (ColumnType): the column type + name (str): the attribute's name + + Returns: + Optional[Type]: the attribute's type, or `None` if it doesn't exist + """ types: TypesRegistry = self.typer.types match name: case "ndim" | "size": diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 95bce9e..5b65314 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -13,6 +13,7 @@ from midas.checker.types import ( ColumnType, Function, GenericType, + ParamSpec, TopType, Type, TypeVar, @@ -26,6 +27,8 @@ if TYPE_CHECKING: @dataclass(frozen=True, kw_only=True) class Call: + """A column method call, implements :class:`utils.MethodCall`""" + location: Location call_expr: p.Expr column: ColumnType @@ -39,6 +42,8 @@ class Call: class ColumnMethodRegistry(MethodRegistry[Call]): + """The method registry for column types""" + def _element_binary_op(self, call: Call, method: str) -> ColumnType: """Compute the result of an element-wise binary operation @@ -74,6 +79,18 @@ class ColumnMethodRegistry(MethodRegistry[Call]): return new_column def _element_wise(self, call: Call, method: str) -> Type: + """Compute the result of an element-wise method call + + If the call is valid, this method also generates an assertion to check + that both operands have the same length at runtime + + Args: + call (Call): the call object + method (str): the method's name + + Returns: + Type: the result type + """ # TODO: support add with scalar # Build signature with new column type and generic operand @@ -82,14 +99,16 @@ class ColumnMethodRegistry(MethodRegistry[Call]): name="add", params=[param_type], body=Function( - args=[ - Function.Argument( - pos=0, - name="other", - type=ColumnType(type=param_type), - required=True, - ), - ], + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="other", + type=ColumnType(type=param_type), + required=True, + ), + ], + ), returns=self._element_binary_op(call, method), ), ) @@ -163,20 +182,35 @@ class ColumnMethodRegistry(MethodRegistry[Call]): def _aggregate( self, call: Call, - kwargs: list[Function.Argument] = [], + kwargs: list[Function.Parameter] = [], *, preserve_inner_type: bool = False, ) -> Type: + """Compute the result type of an aggregate method call + + Args: + call (Call): the call object + kwargs (list[Function.Parameter], optional): a list of extra + keyword-only parameters. Defaults to []. + preserve_inner_type (bool, optional): If `True`, the result type + will preserve the column's inner type (e.g. for `min`/`max`), + otherwise the inner type is widened to `TopType`. Defaults to False. + + Returns: + Type: the result type + """ signature = Function( - kw_args=[ - Function.Argument( - pos=0, - name="axis", - type=TopType(), - required=False, - ), - *kwargs, - ], + params=ParamSpec( + kw=[ + Function.Parameter( + pos=0, + name="axis", + type=TopType(), + required=False, + ), + *kwargs, + ], + ), returns=call.column if preserve_inner_type else ColumnType(type=TopType()), ) @@ -221,7 +255,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 +273,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 +285,16 @@ class ColumnMethodRegistry(MethodRegistry[Call]): @method() def head(self, call: Call) -> Type: signature = Function( - args=[ - Function.Argument( - pos=0, - name="n", - type=self.types.get_type("int"), - required=False, - ), - ], + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + ), returns=call.column, ) @@ -273,14 +309,16 @@ class ColumnMethodRegistry(MethodRegistry[Call]): @method() def tail(self, call: Call) -> Type: signature = Function( - args=[ - Function.Argument( - pos=0, - name="n", - type=self.types.get_type("int"), - required=False, - ), - ], + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + ), returns=call.column, ) @@ -296,52 +334,33 @@ class ColumnMethodRegistry(MethodRegistry[Call]): def groupby(self, call: Call) -> Type: bool_: Type = self.types.get_type("bool") function: Function = Function( - args=[ - Function.Argument( - pos=0, - name="by", - type=TopType(), - required=False, - ), - Function.Argument( - pos=1, - name="level", - type=TopType(), - required=False, - ), - ], - kw_args=[ - Function.Argument( - pos=2, - name="as_index", - type=bool_, - required=False, - ), - Function.Argument( - pos=3, - name="sort", - 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, - ), - ], + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="by", + type=TopType(), + required=False, + ), + Function.Parameter( + pos=1, + name="level", + type=TopType(), + required=False, + ), + ], + kw=[ + Function.Parameter( + pos=i + 2, + name=name, + type=bool_, + required=False, + ) + for i, name in enumerate( + ["as_index", "sort", "group_keys", "observed", "dropna"] + ) + ], + ), returns=ColumnGroupBy(column=call.column), ) @@ -354,6 +373,14 @@ class ColumnMethodRegistry(MethodRegistry[Call]): return result.result def _assert_same_length(self, call_expr: p.Expr, column1: p.Expr, column2: p.Expr): + """Generate an assertion to check that two columns have the same length + + Args: + call_expr (p.Expr): the call expression, to insert the assertion + at the right place + column1 (p.Expr): the first column expression + column2 (p.Expr): the second column expression + """ func_name: str = "__midas_column_same_length__" # Efficiently compute length diff --git a/midas/checker/frames/frame_groupby_methods.py b/midas/checker/frames/frame_groupby_methods.py index 755cab7..cabf88a 100644 --- a/midas/checker/frames/frame_groupby_methods.py +++ b/midas/checker/frames/frame_groupby_methods.py @@ -21,6 +21,8 @@ if TYPE_CHECKING: @dataclass(frozen=True, kw_only=True) class Call: + """A frame group-by method call, implements :class:`utils.MethodCall`""" + location: Location call_expr: p.Expr groupby: FrameGroupBy @@ -34,14 +36,18 @@ class Call: class FrameGroupByMethodRegistry(MethodRegistry[Call]): - NAMED_ARGS: dict[str, str] = { - "numeric_only": "bool", - "skipna": "bool", - "engine": "str", - "engine_kwargs": "dict", - } + """The method registry for frame group-by types""" def _aggregate(self, call: Call, method: str) -> Type: + """Compute the result type of an aggregate method call + + Args: + call (Call): the call object + method (str): the method's name + + Returns: + Type: the result type + """ new_columns: list[DataFrameType.Column] = [] for column in call.groupby.frame.columns: diff --git a/midas/checker/frames/frame_manager.py b/midas/checker/frames/frame_manager.py index 8b99f14..ca75187 100644 --- a/midas/checker/frames/frame_manager.py +++ b/midas/checker/frames/frame_manager.py @@ -24,10 +24,20 @@ if TYPE_CHECKING: def is_list_of_literals(exprs: list[p.Expr]) -> TypeGuard[list[p.LiteralExpr]]: + """Check whether the given list only contains literal expressions + + Args: + exprs (list[p.Expr]): the list to check + + Returns: + TypeGuard[list[p.LiteralExpr]]: whether `exprs` only contains literal expressions + """ return all(isinstance(expr, p.LiteralExpr) for expr in exprs) class FrameManager: + """Helper class to handle methods and subscripts on frame types""" + def __init__(self, typer: PythonTyper) -> None: self.typer: PythonTyper = typer self.method_resolver: FrameMethodRegistry = FrameMethodRegistry(self.typer) @@ -43,6 +53,18 @@ class FrameManager: index: p.Expr, value_type: Type, ) -> Type: + """Compute the new frame type after assigning a value to an index + + Args: + reporter (FileReporter): the file reporter to use for diagnostics + location (Location): the assignment's location + frame (DataFrameType): the frame type + index (p.Expr): the index expression + value_type (Type): the assigned value + + Returns: + Type: the resulting frame type + """ match index: case p.LiteralExpr(value=str() as name): return self.assign_column(reporter, location, frame, name, value_type) @@ -93,6 +115,18 @@ class FrameManager: name: str, type: Type, ) -> Type: + """Compute the new frame type after assigning a single value to a column + + Args: + reporter (FileReporter): the file reporter to use for diagnostics + location (Location): the assignment's location + frame (DataFrameType): the frame type + name (str): the column name + type (Type): the assigned value type + + Returns: + Type: the resulting frame type + """ if not isinstance(type, ColumnType): reporter.error( location, @@ -108,6 +142,17 @@ class FrameManager: frame: DataFrameType, index: p.Expr, ) -> Type: + """Compute the type of a subscript access + + Args: + reporter (FileReporter): the file reporter to use for diagnostics + location (Location): the subscript's location + frame (DataFrameType): the frame type + index (p.Expr): the index expression + + Returns: + Type: the resulting type + """ match index: case p.LiteralExpr(value=str() as name): column: Optional[ColumnType] = FrameManager._get_column(frame, name) @@ -142,6 +187,17 @@ class FrameManager: groupby: FrameGroupBy, index: p.Expr, ) -> Type: + """Compute the type of a subscript access on a frame group-by object + + Args: + reporter (FileReporter): the file reporter to use for diagnostics + location (Location): the subscript's location + groupby (FrameGroupBy): the group-by object + index (p.Expr): the index expression + + Returns: + Type: the resulting type + """ result: Type = self.get(reporter, location, groupby.frame, index) match result: case ColumnType(): @@ -159,6 +215,16 @@ class FrameManager: def _set_column( cls, frame: DataFrameType, name: str, column: ColumnType ) -> DataFrameType: + """Set a frame's column to the given type + + Args: + frame (DataFrameType): the frame type + name (str): the column's name + column (ColumnType): the new column's type + + Returns: + DataFrameType: the new frame type + """ new_columns: list[DataFrameType.Column] = [] index: int = len(frame.columns) replace: bool = False @@ -185,12 +251,31 @@ class FrameManager: def _set_columns( cls, frame: DataFrameType, names: list[str], columns: list[ColumnType] ) -> DataFrameType: + """Set multiple columns of a frame to the given types + + Args: + frame (DataFrameType): the frame type + names (list[str]): the column names + columns (list[ColumnType]): the new column types + + Returns: + DataFrameType: the new frame type + """ for name, col in zip(names, columns): frame = cls._set_column(frame, name, col) return frame @classmethod def _get_column(cls, frame: DataFrameType, name: str) -> Optional[ColumnType]: + """Get a column's type by name + + Args: + frame (DataFrameType): the frame type + name (str): the column's name + + Returns: + Optional[ColumnType]: the column's type, or `None` if it doesn't exist + """ for col in frame.columns: if col.name == name: return col.type @@ -200,6 +285,15 @@ class FrameManager: def _get_columns( cls, frame: DataFrameType, names: list[str] ) -> list[Optional[ColumnType]]: + """Get multiple column types by name + + Args: + frame (DataFrameType): the frame type + names (list[str]): the column names + + Returns: + list[Optional[ColumnType]]: the column types (see :func:`_get_column`) + """ return [cls._get_column(frame, name) for name in names] def call( @@ -212,6 +306,20 @@ class FrameManager: positional: list[TypedExpr], keywords: dict[str, TypedExpr], ) -> Type: + """Compute the result type of a frame's method call + + Args: + method (str): the method name + location (Location): the call's location + call_expr (p.Expr): the call expression + frame (DataFrameType): the frame type + frame_expr (p.Expr): the frame expression + positional (list[TypedExpr]): the list of positional arguments + keywords (dict[str, TypedExpr]): the map of keyword arguments + + Returns: + Type: the result type + """ call: Call = Call( location=location, call_expr=call_expr, @@ -232,6 +340,20 @@ class FrameManager: positional: list[TypedExpr], keywords: dict[str, TypedExpr], ) -> Type: + """Compute the result type of a frame group-by's method call + + Args: + method (str): the method name + location (Location): the call's location + call_expr (p.Expr): the call expression + groupby (FrameGroupBy): the frame group-by object + groupby_expr (p.Expr): the frame group-by expression + positional (list[TypedExpr]): the list of positional arguments + keywords (dict[str, TypedExpr]): the map of keyword arguments + + Returns: + Type: the result type + """ call: GroupByCall = GroupByCall( location=location, call_expr=call_expr, @@ -243,6 +365,15 @@ class FrameManager: return self.groupby_method_resolver.call(method, call) def get_attribute(self, frame: DataFrameType, name: str) -> Optional[Type]: + """Get the type of a frame's attribute + + Args: + frame (DataFrameType): the frame type + name (str): the attribute's name + + Returns: + Optional[Type]: the attribute's type, or `None` if it doesn't exist + """ types: TypesRegistry = self.typer.types match name: case "ndim" | "size": diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 1be481f..3832865 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -14,6 +14,7 @@ from midas.checker.types import ( FrameGroupBy, Function, OverloadedFunction, + ParamSpec, TopType, Type, UnknownType, @@ -26,6 +27,8 @@ if TYPE_CHECKING: @dataclass(frozen=True, kw_only=True) class Call: + """A frame method call, implements :class:`utils.MethodCall`""" + location: Location call_expr: p.Expr frame: DataFrameType @@ -39,6 +42,8 @@ class Call: class FrameMethodRegistry(MethodRegistry[Call]): + """The method registry for frame types""" + def _get_method_result( self, call: Call, @@ -147,17 +152,31 @@ class FrameMethodRegistry(MethodRegistry[Call]): return DataFrameType(columns=new_columns) def _element_wise(self, call: Call, method: str) -> Type: + """Compute the result of an element-wise method call + + If the call is valid, this method also generates an assertion to check + that both operands have the same length at runtime + + Args: + call (Call): the call object + method (str): the method's name + + Returns: + Type: the result type + """ # TODO: support scalar, sequence, Series, dict operand # Build signature with new schema and generic operand signature = Function( - args=[ - Function.Argument( - pos=0, - name="other", - type=DataFrameType(columns=[]), - required=True, - ), - ], + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="other", + type=DataFrameType(columns=[]), + required=True, + ), + ], + ), returns=self._element_binary_op(call, method), ) @@ -227,29 +246,43 @@ 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: + """Compute the result type of an aggregate method call + + Args: + call (Call): the call object + kwargs (list[Function.Parameter], optional): a list of extra + keyword-only parameters. Defaults to []. + + Returns: + Type: the result type + """ with_axis = Function( - kw_args=[ - Function.Argument( - pos=0, - name="axis", - type=self.types.get_type("int"), - required=False, - ), - *kwargs, - ], + params=ParamSpec( + kw=[ + Function.Parameter( + pos=0, + name="axis", + type=self.types.get_type("int"), + required=False, + ), + *kwargs, + ], + ), returns=ColumnType(type=TopType()), ) without_axis = Function( - kw_args=[ - Function.Argument( - pos=0, - name="axis", - type=self.types.get_type("None"), - required=True, - ), - *kwargs, - ], + params=ParamSpec( + kw=[ + Function.Parameter( + pos=0, + name="axis", + type=self.types.get_type("None"), + required=True, + ), + *kwargs, + ], + ), returns=TopType(), ) overload = OverloadedFunction( @@ -300,7 +333,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 +351,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 +363,16 @@ class FrameMethodRegistry(MethodRegistry[Call]): @method() def head(self, call: Call) -> Type: signature = Function( - args=[ - Function.Argument( - pos=0, - name="n", - type=self.types.get_type("int"), - required=False, - ), - ], + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + ), returns=call.frame, ) @@ -352,14 +387,16 @@ class FrameMethodRegistry(MethodRegistry[Call]): @method() def tail(self, call: Call) -> Type: signature = Function( - args=[ - Function.Argument( - pos=0, - name="n", - type=self.types.get_type("int"), - required=False, - ), - ], + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + ), returns=call.frame, ) @@ -375,52 +412,33 @@ class FrameMethodRegistry(MethodRegistry[Call]): def groupby(self, call: Call) -> Type: bool_: Type = self.types.get_type("bool") function: Function = Function( - args=[ - Function.Argument( - pos=0, - name="by", - type=TopType(), - required=False, - ), - Function.Argument( - pos=1, - name="level", - type=TopType(), - required=False, - ), - ], - kw_args=[ - Function.Argument( - pos=2, - name="as_index", - type=bool_, - required=False, - ), - Function.Argument( - pos=3, - name="sort", - 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, - ), - ], + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="by", + type=TopType(), + required=False, + ), + Function.Parameter( + pos=1, + name="level", + type=TopType(), + required=False, + ), + ], + kw=[ + Function.Parameter( + pos=i + 2, + name=name, + type=bool_, + required=False, + ) + for i, name in enumerate( + ["as_index", "sort", "group_keys", "observed", "dropna"] + ) + ], + ), returns=FrameGroupBy(frame=call.frame), ) @@ -433,6 +451,14 @@ class FrameMethodRegistry(MethodRegistry[Call]): return result.result def _assert_same_length(self, call_expr: p.Expr, frame1: p.Expr, frame2: p.Expr): + """Generate an assertion to check that two frames have the same length + + Args: + call_expr (p.Expr): the call expression, to insert the assertion + at the right place + frame1 (p.Expr): the first frame expression + frame2 (p.Expr): the second frame expression + """ func_name: str = "__midas_frame_same_length__" # Efficiently compute length diff --git a/midas/checker/frames/utils.py b/midas/checker/frames/utils.py index dee4696..044ca8e 100644 --- a/midas/checker/frames/utils.py +++ b/midas/checker/frames/utils.py @@ -24,6 +24,12 @@ if TYPE_CHECKING: class _MethodRegistryMeta(type): + """Meta-class for :class:`MethodRegistry` + + Collects methods marked with the :func:`method` decorator into a dictionary + named `_methods` on the class itself + """ + _methods: dict[str, Callable[..., Type]] = {} def __new__( @@ -42,6 +48,11 @@ class _MethodRegistryMeta(type): class MethodCall(Protocol): + """A method call object + + Must have at least `location`, `call_expr` and `subject` properties + """ + @property def location(self) -> Location: ... @@ -56,6 +67,8 @@ T = TypeVar("T", bound=MethodCall) class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta): + """A registry of methods""" + def __init__(self, typer: PythonTyper) -> None: self.typer: PythonTyper = typer @@ -76,6 +89,15 @@ class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta): return self.typer.assertions def call(self, method: str, call: T) -> Type: + """Compute the result type of a call to the given method + + Args: + method (str): the method's name + call (T): the call + + Returns: + Type: the result type + """ func: Optional[Callable[[Self, T], Type]] = self._methods.get(method) if func is None: self.reporter.warning( @@ -90,6 +112,13 @@ Method = Callable[[_Self, T], Type] def method(*names: str) -> Callable[[Method[_Self, T]], Method[_Self, T]]: + """Simple decorator to mark a method as part of the registry + + Args: + names (str): names by which the method can be called. If left empty, the + Python method's name will be used + """ + def wrapper(func: Method[_Self, T]) -> Method[_Self, T]: names_: tuple[str, ...] = names if len(names_) == 0: diff --git a/midas/checker/midas.py b/midas/checker/midas.py index 32716fc..827ce72 100644 --- a/midas/checker/midas.py +++ b/midas/checker/midas.py @@ -1,5 +1,4 @@ import logging -from dataclasses import dataclass from pathlib import Path from typing import Optional @@ -21,6 +20,7 @@ from midas.checker.types import ( ExtensionType, Function, GenericType, + ParamSpec, Predicate, Type, TypeVar, @@ -32,30 +32,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 - - -@dataclass(frozen=True, kw_only=True) -class MappedArgument: - expr: m.Expr - type: Type - argument: Function.Argument - - -@dataclass(frozen=True, kw_only=True) -class OverloadCandidate: - function: Function - mapped: list[MappedArgument] - - class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type]): """A resolver which evaluates Midas type definitions and build a registry""" @@ -82,10 +58,21 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type self._preamble: Environment = Preamble(self.types) def set_reporter(self, reporter: FileReporter): + """Set the file reporter to use for diagnostics + + Args: + reporter (FileReporter): the file reporter + """ self.reporter = reporter self.dispatcher.set_reporter(reporter) def process(self, source: str, path: Optional[str]): + """Process some Midas source code + + Args: + source (str): the Midas source code + path (Optional[str]): the path of the source file, if known + """ reporter: FileReporter = self.reporter.for_file(path) self.set_reporter(reporter) @@ -98,6 +85,14 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type self.resolve(stmts) def type_of(self, expr: m.Expr) -> Type: + """Compute the type of the given expression + + Args: + expr (m.Expr): the expression to type + + Returns: + Type: the type of the expression + """ type: Type = expr.accept(self) return type @@ -118,6 +113,21 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type return self.types.get_type(name) def get_variable(self, name: str) -> Type: + """Get the type of a variable + + This function will first look into the current predicate's parameters if + we are in a predicate definition. + The the variable is looked up in the preamble (i.e. global environment) + + Args: + name (str): the name of the variable + + Raises: + NameError: if the variable cannot be found + + Returns: + Type: the type of the variable + """ if name in self._predicate_params: return self._predicate_params[name] predicate: Optional[Predicate] = self.types.lookup_predicate(name) @@ -145,6 +155,11 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type self.types._types[name] = inferrer.infer(type) def assert_bool(self, expr: m.Expr): + """Check that the given expression is a subtype of `bool` or report an error + + Args: + expr (m.Expr): the expression to check + """ type: Type = self.type_of(expr) if not self.types.is_subtype(type, self._bool): self.reporter.error(expr.location, f"Must be a boolean but is {type}") @@ -196,9 +211,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] = [ - self._visit_param_spec(spec) for spec in stmt.params - ] + params: list[ParamSpec] = [self._visit_param_spec(spec) for spec in stmt.params] if not self._is_valid_predicate(type): self.reporter.error( @@ -209,9 +222,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, - args=spec.mixed, - kw_args=spec.kw, + params=spec, returns=type, ) self._predicate_params = {} @@ -225,6 +236,16 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type ) def _is_valid_predicate(self, body: Type) -> bool: + """Check whether the given type is valid as a predicate's body + + Accepted types are either subtypes of `bool` or valid predicates + + Args: + body (Type): the potential predicate body + + Returns: + bool: `True` if `body` can be a predicate body, `False` otherwise + """ match body: case Function(returns=returns): return self._is_valid_predicate(returns) @@ -250,7 +271,11 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type return self._visit_binary_expr(expr.location, expr.left, expr.right, method) def _visit_binary_expr( - self, location: Location, left_expr: m.Expr, right_expr: m.Expr, method: str + self, + location: Location, + left_expr: m.Expr, + right_expr: m.Expr, + method: str, ) -> Type: left: Type = self.type_of(left_expr) right: Type = self.type_of(right_expr) @@ -386,30 +411,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, - args=params.mixed, - kw_args=params.kw, + params=self._visit_param_spec(type.params), 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: - return Function.Argument( + def process_param( + 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), - type=arg.type.accept(self), - required=arg.required, + name=param.name.lexeme if param.name is not None else str(i), + type=param.type.accept(self), + required=param.required, ) - return TypedParamSpec( - pos=[process_arg(arg, i) for i, arg in enumerate(spec.pos)], - mixed=[process_arg(arg, i + n_pos) for i, arg in enumerate(spec.mixed)], - kw=[process_arg(arg, i + n_pos + n_mixed) for i, arg in enumerate(spec.kw)], + return ParamSpec( + pos=[process_param(param, i) for i, param in enumerate(spec.pos)], + mixed=[ + 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: diff --git a/midas/checker/preamble.py b/midas/checker/preamble.py index a60dabb..5446787 100644 --- a/midas/checker/preamble.py +++ b/midas/checker/preamble.py @@ -7,6 +7,7 @@ from midas.checker.types import ( Function, GenericType, OverloadedFunction, + ParamSpec, TopType, Type, TypeVar, @@ -22,6 +23,8 @@ class Param: class Preamble(Environment): + """The initial environment containing some of Python's builtin functions""" + def __init__(self, types: TypesRegistry) -> None: super().__init__() self._types: TypesRegistry = types @@ -132,9 +135,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 +147,11 @@ class Preamble(Environment): ] function = Function( - pos_args=map_args(pos, 0), - args=map_args(mixed, len(pos)), - kw_args=map_args(kw, len(pos) + len(mixed)), + params=ParamSpec( + pos=map_params(pos, 0), + mixed=map_params(mixed, len(pos)), + kw=map_params(kw, len(pos) + len(mixed)), + ), returns=returns, ) if len(type_vars) != 0: diff --git a/midas/checker/python.py b/midas/checker/python.py index 9103cf4..779a480 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -1,6 +1,5 @@ import ast import logging -from dataclasses import dataclass from typing import Any, Optional import midas.ast.python as p @@ -31,6 +30,7 @@ from midas.checker.types import ( FrameGroupBy, Function, GenericType, + ParamSpec, TopType, TupleType, Type, @@ -55,19 +55,6 @@ class UndefinedMethodException(Exception): pass -@dataclass(frozen=True, kw_only=True) -class MappedArgument: - expr: p.Expr - type: Type - argument: Function.Argument - - -@dataclass(frozen=True, kw_only=True) -class OverloadCandidate: - function: Function - mapped: list[MappedArgument] - - class PythonTyper( p.Stmt.Visitor[None], p.Expr.Visitor[Type], @@ -96,10 +83,24 @@ class PythonTyper( self.assertions: AssertionCollector = AssertionCollector() def set_reporter(self, reporter: FileReporter): + """Set the file reporter to use for diagnostics + + Args: + reporter (FileReporter): the file reporter + """ self.reporter = reporter self.dispatcher.set_reporter(self.reporter) def process(self, source: str, path: Optional[str]) -> TypedAST: + """Process some Python source code + + Args: + source (str): the Python source code + path (Optional[str]): the path of the source file, if known + + Returns: + TypedAST: all generated typechecking information + """ reporter: FileReporter = self.reporter.for_file(path) self.set_reporter(reporter) @@ -124,7 +125,7 @@ class PythonTyper( ) def judge(self, expr: p.Expr, type: Type): - """Record a typing judgement + """Record a typing judgement for the given expression Args: expr (p.Expr): the judged expression @@ -133,7 +134,7 @@ class PythonTyper( self.judgements.append((expr, type)) def compute_type(self, expr: p.Expr) -> Type: - """Evaluate the type of an expression + """Evaluate the type of the given expression Args: expr (p.Expr): the expression to type @@ -144,7 +145,7 @@ class PythonTyper( return expr.accept(self) def type_of(self, expr: p.Expr) -> Type: - """Evaluate the type of an expression and record the judgement + """Evaluate the type of the given expression and record the judgement Args: expr (p.Expr): the expression to evaluate @@ -157,9 +158,22 @@ class PythonTyper( return type def resolve_type_expr(self, expr: p.MidasType) -> Type: + """Resolve the type of a type expression (annotation) + + Args: + expr (p.MidasType): the type expression + + Returns: + Type: the resolved type + """ return expr.accept(self) def process_stmt(self, stmt: p.Stmt) -> None: + """Type check the given statement + + Args: + stmt (p.Stmt): the statement to type-check + """ stmt.accept(self) def process_block(self, block: list[p.Stmt], env: Environment) -> bool: @@ -223,6 +237,24 @@ class PythonTyper( positional: list[TypedExpr], keywords: dict[str, TypedExpr], ) -> Type: + """Evaluate a method call on an object + + Calls to dataframes and columns types are delegated to the appropriate manager + + Args: + location (Location): the location of the call + call_expr (p.Expr): the call expression + obj (TypedExpr): the object on which the method is called + method_name (str): the method name + positional (list[TypedExpr]): the list of positional arguments + keywords (dict[str, TypedExpr]): the map of keyword arguments + + Raises: + UndefinedMethodException: if the method is not defined + + Returns: + Type: the return type of the call + """ unfolded: Type = unfold_type(obj[1]) match unfolded: case DataFrameType(): @@ -282,6 +314,15 @@ class PythonTyper( return result.result def is_subtype(self, type1: Type, type2: Type) -> bool: + """Check whether `type1` is a subtype of `type2` + + Args: + type1 (Type): the potential "subtype" + type2 (Type): the potential "supertype" + + Returns: + bool: whether `type1` is a subtype of `type2` + """ return self.types.is_subtype(type1, type2) def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None: @@ -289,61 +330,64 @@ class PythonTyper( def visit_function(self, stmt: p.Function) -> None: env: Environment = Environment(self.env) - pos_args: list[Function.Argument] = [] - args: list[Function.Argument] = [] - kw_args: list[Function.Argument] = [] + pos: list[Function.Parameter] = [] + mixed: list[Function.Parameter] = [] + kw: list[Function.Parameter] = [] - def eval_arg_type(arg: p.Function.Argument) -> Type: - if arg.type is not None: - return self.resolve_type_expr(arg.type) - if arg.default is not None: - return self.type_of(arg.default) + def eval_param_type(param: p.Function.Parameter) -> Type: + if param.type is not None: + return self.resolve_type_expr(param.type) + if param.default is not None: + return self.type_of(param.default) return UnknownType() - pos: int = 0 - for arg in stmt.posonlyargs: - pos_args.append( - Function.Argument( - pos=pos, - name=arg.name, - type=eval_arg_type(arg), - required=arg.default is None, + position: int = 0 + for param in stmt.params.pos: + pos.append( + Function.Parameter( + pos=position, + name=param.name, + type=eval_param_type(param), + required=param.default is None, ) ) - pos += 1 - for arg in stmt.args: - args.append( - Function.Argument( - pos=pos, - name=arg.name, - type=eval_arg_type(arg), - required=arg.default is None, + position += 1 + for param in stmt.params.mixed: + mixed.append( + Function.Parameter( + pos=position, + name=param.name, + type=eval_param_type(param), + required=param.default is None, ) ) - pos += 1 - for arg in stmt.kwonlyargs: - kw_args.append( - Function.Argument( - pos=pos, # not relevant - name=arg.name, - type=eval_arg_type(arg), - required=arg.default is None, + position += 1 + for param in stmt.params.kw: + kw.append( + Function.Parameter( + pos=position, # not relevant + name=param.name, + type=eval_param_type(param), + required=param.default is None, ) ) - pos += 1 + position += 1 - all_args: list[Function.Argument] = pos_args + args + kw_args - for arg in all_args: - env.define(arg.name, arg.type) + param_spec: ParamSpec = ParamSpec( + pos=pos, + 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, - args=args, - kw_args=kw_args, + params=param_spec, returns=returns_hint, ) self.env.define(stmt.name, inside_function) @@ -375,13 +419,11 @@ class PythonTyper( # TODO: handle *args and **kwargs sinks function: Type = Function( - pos_args=pos_args, - args=args, - kw_args=kw_args, + params=param_spec, 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: @@ -406,6 +448,15 @@ class PythonTyper( self._assign(stmt.location, target, value_type) def _assign(self, location: Location, target: p.Expr, value_type: Type): + """Handle an assignment to the given target + + Delegate to the appropriate method according to the target type + + Args: + location (Location): the location of the assignment + target (p.Expr): the assignment's target + value_type (Type): the value to be assigned + """ match target: case p.VariableExpr(): self._assign_var(location, target, value_type) @@ -427,6 +478,13 @@ class PythonTyper( ) def _assign_var(self, location: Location, target: p.VariableExpr, value_type: Type): + """Type check assignment to the given target + + Args: + location (Location): the location of the assignment + target (p.VariableExpr): the assignment's target + value_type (Type): the value to be assigned + """ name: str = target.name var_type: Optional[Type] = self.look_up_variable(name, target) @@ -445,6 +503,13 @@ class PythonTyper( def _assign_attr( self, location: Location, object: p.Expr, name: str, value_type: Type ): + """Type check assignment to the given target + + Args: + location (Location): the location of the assignment + target (p.VariableExpr): the assignment's target + value_type (Type): the value to be assigned + """ object_type: Type = self.type_of(object) member: Optional[Type] = self.types.lookup_member(object_type, name) if member is None: @@ -464,6 +529,13 @@ class PythonTyper( index: p.Expr, value_type: Type, ): + """Type check assignment to the given target + + Args: + location (Location): the location of the assignment + target (p.VariableExpr): the assignment's target + value_type (Type): the value to be assigned + """ var_type: Type = self.type_of(var) unfolded_type: Type = unfold_type(var_type) # TODO: what happens if type is an alias of a dataframe type @@ -885,6 +957,15 @@ class PythonTyper( ) def _get_iterator_type(self, expr: p.Expr, type: Type) -> Optional[Type]: + """Get the item type of an iterator type + + Args: + expr (p.Expr): the iterator expression + type (Type): the iterator type + + Returns: + Optional[Type]: the item type, or `None` if it cannot be determined + """ # TODO: lookup __iter__ getitem: Optional[Type] = self.types.lookup_member(type, "__getitem__") if getitem is None: @@ -904,6 +985,16 @@ class PythonTyper( return result.result def define_typevar(self, call: p.CallExpr) -> Optional[TypeVar]: + """Define a type variable from a call to `typing.TypeVar` + + Args: + call (p.CallExpr): the call to `typing.TypeVar` + + Returns: + Optional[TypeVar]: the define type variable, or `None` if the call + is invalid + """ + def is_kw_true(name: str) -> bool: match call.keywords.get(name): case p.LiteralExpr(value=True): @@ -946,6 +1037,19 @@ class PythonTyper( return None def _parse_type_from_expr(self, expr: p.Expr) -> p.MidasType: + """Parse a type expression from a raw expression + + This is useful for expressions inside a `TypeVar`'s `bound` parameter + + Args: + expr (p.Expr): the expression to parse + + Raises: + NotImplementedError: if the expression is not supported + + Returns: + p.MidasType: the parsed type node + """ location: Location = expr.location parser = PythonParser() match expr: @@ -958,6 +1062,16 @@ class PythonTyper( raise NotImplementedError def _get_literal(self, expr: p.Expr) -> tuple[bool, Any]: + """Get the literal value of a literal-like expression + + Args: + expr (p.Expr): the expression + + Returns: + tuple[bool, Any]: a tuple containing a boolean indicating whether + the given expression is literal-like, and the literal value (or + `None` if the first value is `False`) + """ match expr: case p.LiteralExpr(value=value): return True, value @@ -1014,6 +1128,17 @@ class PythonTyper( def _evaluate_cast_statically( self, expr: p.CastExpr, subject_type: Type, target_type: Type, lit_value: Any ) -> bool: + """Evaluate the given cast expression statically + + Args: + expr (p.CastExpr): the cast expression + subject_type (Type): the subject type being casted + target_type (Type): the target type to which the expression is casted + lit_value (Any): the literal value of the expression + + Returns: + bool: whether the cast expression could be evaluated successfully + """ match target_type: case TopType(): return True diff --git a/midas/checker/registry.py b/midas/checker/registry.py index 92b54a6..c30c5e4 100644 --- a/midas/checker/registry.py +++ b/midas/checker/registry.py @@ -29,11 +29,15 @@ from midas.checker.types import ( @dataclass class Member: + """A member of a type (property or method)""" + kind: MemberKind type: Type class TypesRegistry: + """A registry of types, type members and predicates""" + def __init__(self) -> None: self.logger: logging.Logger = logging.getLogger("TypesRegistry") self._types: dict[str, Type] = {} @@ -81,6 +85,25 @@ class TypesRegistry: member_type: Type, kind: MemberKind, ): + """Define a member on a type + + If the member is a method and a member with the same name is already + defined on the given type, the two are combined into an :class:`OverloadedFunction`. + + If the member is a property and a member with the same name is already + defined on the given type, the new definition is dropped and an error + is reported. + + In any case, if a member with the same name but a different kind is + already defined on the given type, the new definition is dropped and + an error is reported. + + Args: + type_name (str): the name of the type on which the member is defined + member_name (str): the name of the new member + member_type (Type): the type of the new member + kind (MemberKind): the kind of member to define (property or method) + """ members: dict[str, Member] = self._members.setdefault(type_name, {}) if member_name in members: current: Member = members[member_name] @@ -109,11 +132,29 @@ class TypesRegistry: members[member_name] = Member(kind=kind, type=member_type) def define_predicate(self, name: str, predicate: Predicate): + """Define a predicate + + Args: + name (str): the name of the new predicate + predicate (Predicate): the predicate to define + + Raises: + ValueError: if a predicate with the same name is already defined + """ if name in self._predicates: raise ValueError(f"Predicate {name} already defined") self._predicates[name] = predicate def is_builtin_subtype(self, name1: str, name2: str) -> bool: + """Check whether a type is a subtype of another base on builtin subtype rules + + Args: + name1 (str): the name of the potential subtype + name2 (str): the name of the potential supertype + + Returns: + bool: _description_ + """ subtypes: set[str] = BUILTIN_SUBTYPES.get(name2, set()) if name1 in subtypes: return True @@ -218,6 +259,15 @@ class TypesRegistry: return False def are_equivalent(self, type1: Type, type2: Type) -> bool: + """Check whether two types are equivalent (T <: S and S <: T) + + Args: + type1 (Type): the first type + type2 (Type): the second type + + Returns: + bool: whether `type1` is a subtype and a supertype of `type2` + """ return self.is_subtype(type1, type2) and self.is_subtype(type2, type1) # TODO: verify the logic in here @@ -234,98 +284,118 @@ class TypesRegistry: if not self.is_subtype(func1.returns, func2.returns): return False - pos1: list[Function.Argument] = func1.pos_args - mixed1: list[Function.Argument] = func1.args - kw1: dict[str, Function.Argument] = {a.name: a for a in func1.kw_args} - pos2: list[Function.Argument] = func2.pos_args - mixed2: list[Function.Argument] = func2.args - kw2: dict[str, Function.Argument] = {a.name: a for a in func2.kw_args} + pos1: list[Function.Parameter] = func1.params.pos + mixed1: list[Function.Parameter] = func1.params.mixed + kw1: dict[str, Function.Parameter] = { + param.name: param for param in func1.params.kw + } + 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_name: dict[str, Function.Argument] = {arg.name: arg for arg in mixed2} + mixed_by_pos: dict[int, Function.Parameter] = { + 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: - arg2: Function.Argument - if arg1.pos < len(pos2): - arg2 = pos2[arg1.pos] - elif arg1.pos in mixed_by_pos: - arg2 = mixed_by_pos[arg1.pos] - elif not arg1.required: + for param1 in pos1: + param2: Function.Parameter + if param1.pos < len(pos2): + param2 = pos2[param1.pos] + elif param1.pos in mixed_by_pos: + param2 = mixed_by_pos[param1.pos] + elif not param1.required: continue else: return False - if not is_arg_subtype(arg2, arg1): + if not is_arg_subtype(param2, param1): return False - for name, arg1 in kw1.items(): - arg2: Function.Argument + for name, param1 in kw1.items(): + param2: Function.Parameter if name in kw2: - arg2 = kw2[name] + param2 = kw2[name] elif name in mixed_by_name: - arg2 = mixed_by_name[name] - elif not arg1.required: + param2 = mixed_by_name[name] + 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: - pos_arg2: Optional[Function.Argument] = None - kw_arg2: Optional[Function.Argument] = None - if arg1.name in kw2: - kw_arg2 = kw2[arg1.name] - elif arg1.name in mixed_by_name: - kw_arg2 = mixed_by_name[arg1.name] - if arg1.pos < len(pos2): - pos_arg2 = pos2[arg1.pos] - elif arg1.pos in mixed_by_pos: - pos_arg2 = mixed_by_pos[arg1.pos] + for param1 in mixed1: + pos_param2: Optional[Function.Parameter] = None + kw_param2: Optional[Function.Parameter] = None + if param1.name in kw2: + kw_param2 = kw2[param1.name] + elif param1.name in mixed_by_name: + kw_param2 = mixed_by_name[param1.name] + if param1.pos < len(pos2): + pos_param2 = pos2[param1.pos] + elif param1.pos in mixed_by_pos: + pos_param2 = mixed_by_pos[param1.pos] # No match in func2 and arg is required - if pos_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_names: set[str] = {a.name for a in mixed1} - for arg2 in pos2: - if not arg2.required: + mixed_positions: set[int] = {param.pos for param in mixed1} + mixed_names: set[str] = {param.name for param in mixed1} + for param2 in pos2: + if not param2.required: 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(): - if not arg2.required: + for name, param2 in kw2.items(): + if not param2.required: continue if name not in kw1 and name not in mixed_names: return False - for arg2 in mixed2: - if arg2.required: + for param2 in mixed2: + if param2.required: continue - pos_match: bool = arg2.pos < len(pos1) or arg2.pos in mixed_positions - kw_match: bool = arg2.name in kw1 or arg2.name in mixed_names + pos_match: bool = param2.pos < len(pos1) or param2.pos in mixed_positions + kw_match: bool = param2.name in kw1 or param2.name in mixed_names if not pos_match or not kw_match: return False return True def apply_generic(self, type: Type, args: list[Type]) -> Type: + """Instantiate a generic type with the given type arguments + + Args: + type (Type): the generic + args (list[Type]): the type arguments + + Raises: + ValueError: if the arguments are invalid (wrong number, bound violation) + + Returns: + Type: the applied generic type + """ match type: case DerivedType(name=name, type=base): return DerivedType(name=name, type=self.apply_generic(base, args)) @@ -391,6 +461,19 @@ class TypesRegistry: return [types[i] for i in keep] def lookup_member(self, type: Type, member_name: str) -> Optional[Type]: + """Lookup a member by name on a given type + + This function first looks up directly on the specified type, then + recurse through supertypes until it finds the member or reaches + the root type + + Args: + type (Type): the type on which to lookup the member + member_name (str): the member's name + + Returns: + Optional[Type]: the member's type, or `None` if it is not defined + """ match type: case BaseType(name=name): if name in self._members: @@ -451,18 +534,54 @@ class TypesRegistry: return None def lookup_predicate(self, name: str) -> Optional[Predicate]: + """Lookup a predicate by name + + Args: + name (str): the name of the predicate + + Returns: + Optional[Predicate]: the predicate, or `None` if is not defined + """ return self._predicates.get(name) def _by_name_or_type(self, name_or_type: str | Type) -> Type: + """Get a type by name or return it as is + + If `name_or_type` is a string, the associated type is looked up and returned. + Otherwise, the type is returned as is. + + Args: + name_or_type (str | Type): the type or type's name + + Returns: + Type: the type + """ if isinstance(name_or_type, str): return self.get_type(name_or_type) return name_or_type def list_of(self, item_type: str | Type) -> Type: + """Helper method to type a list of a given item type + + Args: + item_type (str | Type): the item type + + Returns: + Type: the list type + """ list_ = self.get_type("list") return self.apply_generic(list_, [self._by_name_or_type(item_type)]) def tuple_of(self, *item_types: str | Type) -> Type: + """Helper method to type a tuple of given item types + + Args: + item_type (str | Type): the item types + + Returns: + Type: the tuple type + """ + tuple_ = self.get_type("tuple") return self.apply_generic( tuple_, @@ -470,6 +589,15 @@ class TypesRegistry: ) def dict_of(self, key_type: str | Type, value_type: str | Type) -> Type: + """Helper method to type a dict of given key and value types + + Args: + key_type (str | Type): the key type + value_type (str | Type): the value type + + Returns: + Type: the dict type + """ dict_ = self.get_type("dict") return self.apply_generic( dict_, diff --git a/midas/checker/reporter.py b/midas/checker/reporter.py index b61b8f3..970ef34 100644 --- a/midas/checker/reporter.py +++ b/midas/checker/reporter.py @@ -7,6 +7,8 @@ from midas.checker.diagnostic import Diagnostic, DiagnosticType class Reporter: + """Helper class to store diagnostics""" + def __init__(self): self.diagnostics: list[Diagnostic] = [] @@ -17,6 +19,14 @@ class Reporter: location: Location, message: str, ): + """Create and record a diagnostic + + Args: + path (Optional[str]): the path linked to this diagnostic + type (DiagnosticType): the type of diagnostic + location (Location): the location if the diagnostic in the file + message (str): the diagnostic's message + """ self.diagnostics.append( Diagnostic( file_path=path, @@ -27,21 +37,52 @@ class Reporter: ) def for_file(self, path: Optional[str]) -> FileReporter: + """Create a new file reporter for the given path using this reporter + + Args: + path (Optional[str]): the path for the new file reporter + + Returns: + FileReporter: the new file reporter, linked to this reporter + """ return FileReporter(self, path) class FileReporter: + """Helper class to manage diagnostics for a file""" + def __init__(self, base_reporter: Reporter, path: Optional[str]) -> None: self.base_reporter: Reporter = base_reporter self.path: Optional[str] = path def for_file(self, path: Optional[str]) -> FileReporter: + """Create a new file reporter for the given path with the same base reporter + + Args: + path (Optional[str]): the path for the new file reporter + + Returns: + FileReporter: the file reporter + """ return FileReporter(self.base_reporter, path) def report(self, type: DiagnosticType, location: Location, message: str): + """Report a diagnostic to the base reporter + + Args: + type (DiagnosticType): the type of diagnostic + location (Location): the location of the diagnostic in the file + message (str): the diagnostic's message + """ self.base_reporter.report(self.path, type, location, message) def error(self, location: Location, message: str): + """Report an error diagnostic + + Args: + location (Location): the location of the diagnostic in the file + message (str): the diagnostic's message + """ self.report( type=DiagnosticType.ERROR, location=location, @@ -49,6 +90,12 @@ class FileReporter: ) def warning(self, location: Location, message: str): + """Report a warning diagnostic + + Args: + location (Location): the location of the diagnostic in the file + message (str): the diagnostic's message + """ self.report( type=DiagnosticType.WARNING, location=location, @@ -56,6 +103,12 @@ class FileReporter: ) def info(self, location: Location, message: str): + """Report an info diagnostic + + Args: + location (Location): the location of the diagnostic in the file + message (str): the diagnostic's message + """ self.report( type=DiagnosticType.INFO, location=location, @@ -63,6 +116,12 @@ class FileReporter: ) def debug(self, location: Location, message: str): + """Report a debug diagnostic + + Args: + location (Location): the location of the diagnostic in the file + message (str): the diagnostic's message + """ self.report( type=DiagnosticType.DEBUG, location=location, diff --git a/midas/checker/resolver.py b/midas/checker/resolver.py index 43768bd..d7af800 100644 --- a/midas/checker/resolver.py +++ b/midas/checker/resolver.py @@ -78,6 +78,14 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]): return def is_defined(self, name: str) -> bool: + """Check whether the given variable is defined in any scope + + Args: + name (str): the name of the variable + + Returns: + bool: `True` if the variable is defined in a scope, `False` otherwise + """ for scope in self.scopes: if name in scope: return True @@ -93,7 +101,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) diff --git a/midas/checker/types.py b/midas/checker/types.py index 51658d7..f1b38f4 100644 --- a/midas/checker/types.py +++ b/midas/checker/types.py @@ -10,12 +10,16 @@ from midas.ast.printer import MidasPrinter @dataclass(frozen=True, kw_only=True) class TopType: + """The top type (`Any`)""" + def __str__(self) -> str: return "Any" @dataclass(frozen=True, kw_only=True) class BaseType: + """A base / builtin type""" + name: str def __str__(self) -> str: @@ -24,6 +28,8 @@ class BaseType: @dataclass(frozen=True, kw_only=True) class DerivedType: + """A derived type, i.e. a named subtype of another type""" + name: str type: Type @@ -33,40 +39,32 @@ class DerivedType: @dataclass(frozen=True, kw_only=True) class UnknownType: + """An unknown type""" + def __str__(self) -> str: return "" @dataclass(frozen=True, kw_only=True) class UnitType: + """The unit type (`None`)""" + def __str__(self) -> str: return "None" @dataclass(frozen=True, kw_only=True) class Function: - pos_args: list[Argument] = field(default_factory=list) - args: list[Argument] = field(default_factory=list) - kw_args: list[Argument] = field(default_factory=list) + """A function type""" + + params: ParamSpec returns: Type def __str__(self) -> str: - args: list[str] = [] - 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}" + return f"{self.params} -> {self.returns}" @dataclass(frozen=True, kw_only=True) - class Argument: + class Parameter: pos: int name: str type: Type @@ -77,8 +75,34 @@ class Function: return f"{self.name}: {self.type}{opt}" +@dataclass(frozen=True, kw_only=True) +class ParamSpec: + """A function's parameter spec""" + + pos: list[Function.Parameter] = field(default_factory=list) + mixed: list[Function.Parameter] = field(default_factory=list) + kw: list[Function.Parameter] = field(default_factory=list) + + 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: + """A list of method overloads""" + overloads: list[Type] def __str__(self) -> str: @@ -87,6 +111,8 @@ class OverloadedFunction: @dataclass(frozen=True, kw_only=True) class ComplexType: + """A type with inline members""" + members: dict[str, Type] def __str__(self) -> str: @@ -96,6 +122,8 @@ class ComplexType: @dataclass(frozen=True, kw_only=True) class ExtensionType: + """An extension of a type, adding members through a `ComplexType`""" + base: Type extension: ComplexType @@ -104,6 +132,8 @@ class ExtensionType: class Variance(StrEnum): + """The variance of a :class:`TypeVar`""" + INVARIANT = "INVARIANT" COVARIANT = "COVARIANT" CONTRAVARIANT = "CONTRAVARIANT" @@ -111,6 +141,8 @@ class Variance(StrEnum): @dataclass(frozen=True, kw_only=True) class TypeVar: + """A type variable, often used as type parameters for a generic type""" + name: str bound: Optional[Type] variance: Variance = Variance.INVARIANT @@ -128,6 +160,8 @@ class TypeVar: @dataclass(frozen=True, kw_only=True) class GenericType: + """A generic type, with type parameters and a generic body type""" + name: str params: list[TypeVar] body: Type @@ -138,6 +172,8 @@ class GenericType: @dataclass(frozen=True, kw_only=True) class AppliedType: + """An instance of a :class:`GenericType`, with concrete type arguments substituted in its body""" + name: str args: list[Type] body: Type @@ -148,6 +184,8 @@ class AppliedType: @dataclass(frozen=True, kw_only=True) class ConstraintType: + """A type with a constraint expression""" + type: Type constraint: m.Expr @@ -158,6 +196,8 @@ class ConstraintType: @dataclass(frozen=True, kw_only=True) class TupleType: + """A tuple type, containing any number of ordered item types""" + items: tuple[Type, ...] def __str__(self) -> str: @@ -166,6 +206,8 @@ class TupleType: @dataclass(frozen=True, kw_only=True) class ColumnType: + """A column type containing items of a given unique type""" + type: Type def __str__(self) -> str: @@ -174,6 +216,8 @@ class ColumnType: @dataclass(frozen=True, kw_only=True) class DataFrameType: + """A data-frame type, containing named columns of specific :class:`ColumnType`""" + columns: list[Column] def __str__(self) -> str: @@ -189,6 +233,8 @@ class DataFrameType: @dataclass(frozen=True, kw_only=True) class FrameGroupBy: + """A frame group-by object""" + frame: DataFrameType def __str__(self) -> str: @@ -197,6 +243,8 @@ class FrameGroupBy: @dataclass(frozen=True, kw_only=True) class ColumnGroupBy: + """A column group-by object""" + column: ColumnType def __str__(self) -> str: @@ -204,12 +252,32 @@ class ColumnGroupBy: def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: - def sub_argument(arg: Function.Argument): - return Function.Argument( - pos=arg.pos, - name=arg.name, - type=substitute_typevars(arg.type, substitutions), - required=arg.required, + """Substitute type variables in the given type + + This function is called recursively on inner type structures + + Args: + type (Type): the type in which to substitute type variables + substitutions (dict[str, Type]): a mapping of type variable names to + concrete types + + Returns: + Type: the resulting type with substitutions applied + """ + + def sub_parameter(param: Function.Parameter): + return Function.Parameter( + pos=param.pos, + name=param.name, + type=substitute_typevars(param.type, substitutions), + 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 +303,11 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: ) case Function( - pos_args=pos_args, - args=args, - kw_args=kw_args, + params=params, returns=returns, ): return Function( - pos_args=list(map(sub_argument, pos_args)), - args=list(map(sub_argument, args)), - kw_args=list(map(sub_argument, kw_args)), + params=sub_param_spec(params), returns=substitute_typevars(returns, substitutions), ) @@ -343,6 +407,14 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: def unfold_type(type: Type) -> Type: + """Unfold a chain of :class:`DerivedType` to get the root supertype + + Args: + type (Type): the type to unfold + + Returns: + Type: the root supertype + """ match type: case DerivedType(type=ref_type): return unfold_type(ref_type) @@ -351,14 +423,23 @@ def unfold_type(type: Type) -> Type: def to_annotation(type: Type) -> str: - def _args_annotation(func: Function) -> str: - if len(func.kw_args) != 0: + """Convert the given type to a Python annotation string + + Args: + type (Type): the type to convert + + Returns: + str: the annotation string + """ + + def _params_annotation(spec: ParamSpec) -> str: + if len(spec.kw) != 0: return "..." - args: str = ", ".join( - to_annotation(arg.type) for arg in func.pos_args + func.args + params: str = ", ".join( + to_annotation(param.type) for param in spec.pos + spec.mixed ) - return f"[{args}]" + return f"[{params}]" match type: case TopType(): @@ -376,8 +457,8 @@ def to_annotation(type: Type) -> str: case UnitType(): return "None" - case Function(returns=returns): - params_annot: str = _args_annotation(type) + case Function(params=params, returns=returns): + params_annot: str = _params_annotation(params) return f"Callable[{params_annot}, {to_annotation(returns)}]" case OverloadedFunction(): @@ -419,6 +500,8 @@ def to_annotation(type: Type) -> str: @dataclass(frozen=True, kw_only=True) class Predicate: + """A predicate""" + type: Type body: m.Expr alias: bool diff --git a/midas/checker/unifier.py b/midas/checker/unifier.py index 138a503..24804eb 100644 --- a/midas/checker/unifier.py +++ b/midas/checker/unifier.py @@ -8,6 +8,7 @@ from midas.checker.types import ( DataFrameType, Function, GenericType, + ParamSpec, TopType, Type, TypeVar, @@ -18,6 +19,14 @@ class UnificationError(Exception): ... class Unifier: + """ + Helper class to unify generic types in concrete usages + + This can be used for example when a generic function is called with concrete + arguments, at which point the type parameters of the function signature + should be resolvable + """ + def __init__(self, types: TypesRegistry) -> None: self.types: TypesRegistry = types self.logger: logging.Logger = logging.getLogger("Unifier") @@ -28,26 +37,37 @@ class Unifier: positional: list[Type], keywords: dict[str, Type], ) -> Optional[Type]: + """Try and unify a generic function call given concrete arguments + + Args: + type (GenericType): the generic function type + positional (list[Type]): the list of positional arguments + keywords (dict[str, Type]): the map of keyword arguments + + Returns: + Optional[Type]: the concrete function type if unifiable, or `None` + """ concrete_func: Function = Function( - pos_args=[ - Function.Argument( - pos=i, - name=str(i), - type=arg, - required=True, - ) - for i, arg in enumerate(positional) - ], - args=[], - kw_args=[ - Function.Argument( - pos=len(positional) + i, - name=name, - type=arg, - required=True, - ) - for i, (name, arg) in enumerate(keywords.items()) - ], + params=ParamSpec( + pos=[ + Function.Parameter( + pos=i, + name=str(i), + type=arg, + required=True, + ) + for i, arg in enumerate(positional) + ], + kw=[ + Function.Parameter( + pos=len(positional) + i, + name=name, + type=arg, + required=True, + ) + for i, (name, arg) in enumerate(keywords.items()) + ], + ), returns=TopType(), # TODO: use expected type ) return self.unify_generic(type, concrete_func, match_return=False) @@ -58,6 +78,18 @@ class Unifier: concrete: Type, match_return: bool = True, ) -> Optional[Type]: + """Unify a generic type's parameters given a concrete usage + + Args: + template (GenericType): the generic type + concrete (Type): a concrete usage + match_return (bool, optional): if `template` is a function type, + whether its return type must be matched (see :func:`match`). + Defaults to True. + + Returns: + Optional[Type]: the concrete type if unifiable, or `None` + """ substitutions: dict[str, Type] try: substitutions = self.match(template.body, concrete, match_return) @@ -79,6 +111,22 @@ class Unifier: concrete: Type, match_return: bool = True, ) -> dict[str, Type]: + """Match a generic type with a concrete usage, recording parameter substitutions + + Args: + template (Type): the generic type + concrete (Type): a concrete usage + match_return (bool, optional): if `template` and `concrete` are both + :class:`Function`, whether their return types are also matched. + Defaults to True. + + Raises: + UnificationError: if there is a conflict in parameter substitutions + + Returns: + dict[str, Type]: the parameter substitutions which, + applied to `template`, yield `concrete` + """ # TODO: if concrete is Generic, record bound TypeVar. Then when merging # substitutions, check that the constraint is respected match (template, concrete): @@ -125,7 +173,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] = {} @@ -148,6 +196,18 @@ class Unifier: return {} def merge(self, subs1: dict[str, Type], subs2: dict[str, Type]) -> dict[str, Type]: + """Merge two maps of substitutions and raise an error if incompatible + + Args: + subs1 (dict[str, Type]): the first substitutions + subs2 (dict[str, Type]): the second substitutions + + Raises: + UnificationError: if there is a conflict between the two maps + + Returns: + dict[str, Type]: the merged map of substitutions + """ merged: dict[str, Type] = subs1.copy() for k, v in subs2.items(): @@ -161,19 +221,32 @@ class Unifier: def map_params( self, func1: Function, func2: Function - ) -> list[tuple[Function.Argument, Function.Argument]]: - pos1: list[Function.Argument] = func1.pos_args - mixed1: list[Function.Argument] = func1.args - kw1: list[Function.Argument] = func1.kw_args + ) -> list[tuple[Function.Parameter, Function.Parameter]]: + """Map parameters of two functions - pos2: list[Function.Argument] = func2.pos_args - mixed2: list[Function.Argument] = func2.args - kw2: list[Function.Argument] = func2.kw_args + Args: + func1 (Function): the first function + func2 (Function): the second function - mapped: list[tuple[Function.Argument, Function.Argument]] = [] + Returns: + list[tuple[Function.Parameter, Function.Parameter]]: the list of parameter pairs + """ + pos1: list[Function.Parameter] = func1.params.pos + mixed1: list[Function.Parameter] = func1.params.mixed + kw1: list[Function.Parameter] = func1.params.kw - by_pos2: dict[int, Function.Argument] = {arg.pos: arg for arg in pos2 + mixed2} - by_name2: dict[str, Function.Argument] = {arg.name: arg for arg in mixed2 + kw2} + pos2: list[Function.Parameter] = func2.params.pos + mixed2: list[Function.Parameter] = func2.params.mixed + kw2: list[Function.Parameter] = func2.params.kw + + mapped: list[tuple[Function.Parameter, Function.Parameter]] = [] + + by_pos2: dict[int, Function.Parameter] = { + 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: diff --git a/midas/checker/variance.py b/midas/checker/variance.py index e7f5ac0..bcbdc2b 100644 --- a/midas/checker/variance.py +++ b/midas/checker/variance.py @@ -16,14 +16,27 @@ Polarity = Literal[-1, 0, 1] class Tracker: + """Helper class to track the polarity of type parameter references and computer their variance""" + def __init__(self, vars: list[TypeVar]) -> None: self.vars: list[TypeVar] = vars self.refs: dict[str, set[Polarity]] = {var.name: set() for var in self.vars} def record(self, var: TypeVar, polarity: Polarity): + """Record a polarity of the given type parameter + + Args: + var (TypeVar): the type parameter + polarity (Polarity): the polarity + """ self.refs[var.name].add(polarity) def get_updated_vars(self) -> list[TypeVar]: + """Get a list of the tracked type variables with their recorded variance + + Returns: + list[TypeVar]: the list of update type parameters + """ return [ TypeVar( name=var.name, bound=var.bound, variance=self.get_variance(var.name) @@ -32,6 +45,18 @@ class Tracker: ] def get_variance(self, name: str) -> Variance: + """Get the variance of a type parameter + + If the type parameter is only referenced in positive positions, it is + covariant. If it is only referenced in negative positions, it is + contravariant. Otherwise, it is invariant + + Args: + name (str): the name of the type parameter + + Returns: + Variance: the variance of the type parameter + """ refs: set[Polarity] = self.refs[name] if refs == {-1}: return Variance.CONTRAVARIANT @@ -46,11 +71,22 @@ class Tracker: class VarianceInferrer: + """Helper class to compute type parameter variance""" + def __init__(self, types: TypesRegistry) -> None: self.types: TypesRegistry = types self.tracker: Tracker = Tracker([]) def infer(self, type: GenericType) -> GenericType: + """Infer the variance of a generic type's parameters + + Args: + type (GenericType): the generic type + + Returns: + GenericType: a new generic type with its parameters updated with + their inferred variance + """ self.tracker = Tracker(type.params) self.walk(type.body, 1, type.name) @@ -71,20 +107,36 @@ class VarianceInferrer: base_name: str, path: Optional[list[str]] = None, ): + """Walk the type nodes and record variance + + This function recurses into type substructures (e.g. function parameters, + overloads, constraint type bases, etc.) + + When recursing, the polarity is flipped for consumer positions (e.g. function + parameters) or kept the same for producer positions (e.g. return type) + + Args: + type (Type): the type to visit + polarity (Polarity): the current polarity + base_name (str): the root generic type name (used to detect and + handle cyclic references) + path (Optional[list[str]], optional): the path to reach the current + type from the root generic type (used for debugging). Defaults to None. + """ if path is None: path = [] 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): - all_args: list[Function.Argument] = pos_args + mixed_args + kw_args - for arg in all_args: + case Function(params=spec): + all_params: list[Function.Parameter] = spec.pos + spec.mixed + spec.kw + 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"]) diff --git a/midas/cli/highlighter.py b/midas/cli/highlighter.py index 8d142f1..15f3251 100644 --- a/midas/cli/highlighter.py +++ b/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_function_argument(arg) + self._highlight_param_spec(stmt.params) for body_stmt in stmt.body: body_stmt.accept(self) - def _highlight_function_argument(self, arg: p.Function.Argument) -> None: - self.wrap(arg, "argument") - if arg.type is not None: - arg.type.accept(self) + def _highlight_param_spec(self, spec: p.ParamSpec) -> None: + for param in spec.all: + self._highlight_function_param(param) + + 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) diff --git a/midas/cli/hl_python.css b/midas/cli/hl_python.css index 899c461..b0b29d0 100644 --- a/midas/cli/hl_python.css +++ b/midas/cli/hl_python.css @@ -23,7 +23,7 @@ span { --col: 215, 103, 224; } - &.argument { + &.parameter { --col: 103, 192, 224; } } \ No newline at end of file diff --git a/midas/generator/constraints.py b/midas/generator/constraints.py index e840b42..c936884 100644 --- a/midas/generator/constraints.py +++ b/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=[], - args=[ - Function.Argument( - pos=0, - name="_", - type=self.types.get_type("Any"), - required=True, - ) - ], - kw_args=[], + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="_", + type=self.types.get_type("Any"), + required=True, + ) + ], + ), 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, - annotation=ast.Constant(value=to_annotation(arg.type)), + arg=param.name, + 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, - annotation=ast.Constant(value=to_annotation(arg.type)), + arg=param.name, + annotation=ast.Constant(value=to_annotation(param.type)), ) - for arg in func.args + for param in params.mixed ], kwonlyargs=[ ast.arg( - arg=arg.name, - annotation=ast.Constant(value=to_annotation(arg.type)), + arg=param.name, + 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=[], diff --git a/midas/generator/generator.py b/midas/generator/generator.py index 352a544..07799b8 100644 --- a/midas/generator/generator.py +++ b/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( - 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 - ], - ), + args=self.make_args(stmt.params), body=self._visit_body(stmt.body), decorator_list=[], ) diff --git a/midas/generator/stubs.py b/midas/generator/stubs.py index 07ab0db..7ddcb06 100644 --- a/midas/generator/stubs.py +++ b/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)) - for arg in func.pos_args + ast.arg( + 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)) - for arg in func.args + ast.arg( + 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)) - for arg in func.kw_args + ast.arg( + 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=[], diff --git a/midas/lexer/base.py b/midas/lexer/base.py index c4f4d82..f21205b 100644 --- a/midas/lexer/base.py +++ b/midas/lexer/base.py @@ -16,9 +16,10 @@ class Lexer(ABC): """An abstract lexer which provides methods to easily extend it into a concrete one This implementation is based on the [_Crafting Interpreters_][1] book by Robert Nystrom, - more specifically on my [previous Python implementation](https://git.kb28.ch/HEL/pebble) + more specifically on my [previous Python implementation][2] [1]: https://craftinginterpreters.com/ + [2]: https://git.kb28.ch/HEL/pebble """ def __init__(self, source: str, file: Optional[str] = None) -> None: @@ -168,6 +169,6 @@ class Lexer(ABC): def scan_token(self) -> None: """Scan a token - This function should (at least) consume the current character and produce the appropriate token(s), using `add_token` + This function should (at least) consume the current character and produce the appropriate token(s), using :func:`add_token` """ pass diff --git a/midas/lexer/midas.py b/midas/lexer/midas.py index a530b47..e1f6758 100644 --- a/midas/lexer/midas.py +++ b/midas/lexer/midas.py @@ -81,6 +81,12 @@ class MidasLexer(Lexer): return None def scan_string(self, opening: str): + """Scan the rest of a string and add it as a token + + Args: + opening (str): the opening quote or double quote, to be matched + at the end of the string + """ while self.peek() != opening and not self.is_at_end(): self.advance() @@ -147,6 +153,18 @@ class MidasLexer(Lexer): self.add_token(TokenType.COMMENT) def is_identifier_char(self, char: str, *, start: bool) -> bool: + """Check whether a character is a valid as part of an identifier + + Identifiers can contain any alphanumerical character or underscore. + They cannot start with a digit. + + Args: + char (str): the character to check + start (bool): whether this is the first character of the identifier + + Returns: + bool: `True` if the character is valid, `False` otherwise + """ if char == "_": return True if char.isalpha(): diff --git a/midas/lexer/token.py b/midas/lexer/token.py index f06c21e..d06052c 100644 --- a/midas/lexer/token.py +++ b/midas/lexer/token.py @@ -104,6 +104,15 @@ class Token: ) def location_to(self, to: Token) -> Location: + """Create a new :class:`Location` spanning from this token to another + + Args: + to (Token): the end token + + Returns: + Location: a new :class:`Location` starting at this token and ending + at `to`, both included + """ return Location.span(self.get_location(), to.get_location()) @property diff --git a/midas/parser/base.py b/midas/parser/base.py index 255cd26..e039c02 100644 --- a/midas/parser/base.py +++ b/midas/parser/base.py @@ -16,6 +16,9 @@ class TokenError: def get_report(self) -> str: """Get a detailed error message + The error message is formatted as "() Error at : ". + For example: "(L2:5) Error at '3': Expected ')' after arguments." + Returns: str: the complete error message """ @@ -32,9 +35,10 @@ class Parser(ABC, Generic[T]): """An abstract parser which provides methods to easily extend it into a concrete one This implementation is based on the [_Crafting Interpreters_][1] book by Robert Nystrom, - more specifically on my [previous Python implementation](https://git.kb28.ch/HEL/pebble) + more specifically on my [previous Python implementation][2] [1]: https://craftinginterpreters.com/ + [2]: https://git.kb28.ch/HEL/pebble """ IGNORE: set[TokenType] = { @@ -173,7 +177,7 @@ class Parser(ABC, Generic[T]): error_msg (str): the error message if the token doesn't match Raises: - SyntaxError: if the current token doesn't match the given type + ParsingError: if the current token doesn't match the given type Returns: Token: the current token which matched the given type diff --git a/midas/parser/midas.py b/midas/parser/midas.py index 534f95e..8b3282a 100644 --- a/midas/parser/midas.py +++ b/midas/parser/midas.py @@ -35,10 +35,11 @@ from midas.parser.base import Parser from midas.parser.errors import ParsingError -class MidasParser(Parser): +class MidasParser(Parser[list[Stmt]]): """A simple parser for midas type definitions""" SYNC_BOUNDARY: set[TokenType] = { + TokenType.ALIAS, TokenType.TYPE, TokenType.EXTEND, TokenType.PREDICATE, @@ -73,10 +74,10 @@ class MidasParser(Parser): def declaration(self) -> Optional[Stmt]: """Try and parse a declaration - Any parsing error is caught and None is returned + Any parsing error is caught and `None` is returned Returns: - Optional[Stmt]: the parsed Midas statement, or None if a ParsingError was raised + Optional[Stmt]: the parsed Midas statement, or `None` if a ParsingError was raised """ try: if self.match(TokenType.TYPE): @@ -95,23 +96,14 @@ class MidasParser(Parser): def type_declaration(self) -> TypeStmt: """Parse a type declaration - A type declaration can either be a simple type alias or a new complex type. - In either case, it can have an optional template expression after its name, wrapped in brackets. - A simple type alias is derived from a base type expression, and can have a optional constraint expression preceded by the `where` keyword. - A full simple type alias is thus written: - ``` - type Name[Template](TypeExpr) where Condition - ``` + A type declaration creates a named subtype of a type expression. + It can have an optional template expression after its name, wrapped in brackets, to handle type parameters. - A new complex type has a set of properties which are named, have a type and an optional constraint expression (also preceded by the `where` keyword). - A full complex type definition is thus written: - ``` - type Name[Template] { - prop1: TypeExpr1 where Condition1 - prop2: TypeExpr2 where Condition2 - ... - } - ``` + A type statement consists of: + - the `type` keyword + - a name (identifier) + - (optional) type parameters + - a body, a type expression (see :func:`type_expr`) Returns: TypeStmt: the parsed type declaration statement @@ -165,11 +157,16 @@ class MidasParser(Parser): def alias_declaration(self) -> AliasStmt: """Parse an alias declaration + An alias statement consists of: + - the `alias` keyword + - a name (identifier) + - a body, a type expression (see :func:`type_expr`) + Returns: AliasStmt: the parsed alias declaration statement """ keyword: Token = self.previous() - name: Token = self.consume_identifier("Expected type name") + name: Token = self.consume_identifier("Expected alias name") self.consume(TokenType.EQUAL, "Expected '=' before alias definition") @@ -184,8 +181,8 @@ class MidasParser(Parser): def type_expr(self) -> Type: """Parse a type expression - A type is an identifier, optionally followed by a template expression. - It can also optionally be followed by a '?' to indicate a nullable type + A type expression can either be a function type (see :func:`function`) + or a constraint type (see :func:`constraint_type`) Returns: TypeExpr: the parsed type expression @@ -205,6 +202,15 @@ class MidasParser(Parser): return base def constraint_type(self) -> Type: + """Parse a constraint type expression + + A constraint type consists of a base type (see :func:`base_type`), + optionally followed by the `where` keyword and a constraint + expression (see :func:`constraint`) + + Returns: + Type: the parsed constraint type expression + """ type: Type = self.base_type() if self.match(TokenType.WHERE): constraint: Expr = self.constraint() @@ -216,6 +222,14 @@ class MidasParser(Parser): return type def base_type(self) -> Type: + """Parse a base type expression + + A base type is either a parenthesized type expression (see :func:`type_expr`) + or a generic type (see :func:`generic_type`) + + Returns: + Type: the parsed base type expression + """ if self.match(TokenType.LEFT_PAREN): type: Type = self.type_expr() self.consume(TokenType.RIGHT_PAREN, "Unclosed parenthesis") @@ -227,6 +241,17 @@ class MidasParser(Parser): return self.generic_type() def generic_type(self) -> Type: + """Parse a generic type expression + + A generic type consists of a named type (see :func:`named_type`), + optionally followed by type arguments in brackets. + + The special `Frame` type accepts a frame schema instead of type + arguments (see :func:`frame_type`). + + Returns: + Type: the parsed generic type + """ type: NamedType = self.named_type() if self.check(TokenType.LEFT_BRACKET): if type.name.lexeme == "Frame": @@ -240,6 +265,13 @@ class MidasParser(Parser): return type def type_args(self) -> list[Type]: + """Parse a list of type arguments + + Type arguments are a comma-separated list of type expression wrapped in brackets. + + Returns: + list[Type]: the list of type arguments, if any, or an empty list + """ args: list[Type] = [] self.consume(TokenType.LEFT_BRACKET, "Missing '[' before generic arguments") while not self.is_at_end() and not self.check(TokenType.RIGHT_BRACKET): @@ -250,6 +282,13 @@ class MidasParser(Parser): return args def named_type(self) -> NamedType: + """Parse a named type expression + + A named type is an identifier token + + Returns: + NamedType: the parsed named type expression + """ name: Token = self.consume_identifier("Expected type name") return NamedType( location=name.get_location(), @@ -257,13 +296,13 @@ class MidasParser(Parser): ) def complex_type(self) -> ComplexType: - """Parse a type definition body + """Parse a complex type expression - A type definition body is a set of whitespace-separated - property statements enclosed in curly braces + A complex type consists of zero or more member statements enclosed in + curly braces Returns: - ComplexType: the parsed complex type + ComplexType: the parsed complex type expression """ left: Token = self.consume( TokenType.LEFT_BRACE, "Expected '{' to start type body" @@ -285,6 +324,20 @@ class MidasParser(Parser): ) def frame_type(self) -> FrameType: + """Parse a frame type expression + + A frame type consists of: + - the `Frame` identifier + - an opening bracket `[` + - a list of comma-separated column expression consisting of: + - a name (token) + - a colon `:` + - a type expression (see :func:`type_expr`) + - a closing bracket `]` + + Returns: + FrameType: the parsed frame type + """ keyword: Token = self.previous() self.consume(TokenType.LEFT_BRACKET, "Expected '[' to start frame schema") @@ -311,9 +364,9 @@ class MidasParser(Parser): ) def constraint(self) -> Expr: - """Parse a constraint + """Parse a constraint expression - A constraint is basically a logical predicate + A constraint is an expression (see :func:`expression`) Returns: Expr: the parsed constraint expression @@ -321,10 +374,20 @@ class MidasParser(Parser): return self.expression() def expression(self) -> Expr: + """Parse an expression + + An expression consists of a logical AND expression (see :func:`and_`) + + Returns: + Expr: the parsed expression + """ return self.and_() def and_(self) -> Expr: - """Parse a logical AND expression or a simpler expression + """Parse a logical AND expression + + An AND consists of one or more equality expressions (see :func:`equality`) + separated by logical AND operators (`&`) Returns: Expr: the parsed expression @@ -340,7 +403,10 @@ class MidasParser(Parser): return expr def equality(self) -> Expr: - """Parse a logical equality expression or a simpler expression + """Parse an equality expression + + An equality consists of one or more comparison expressions (see :func:`comparison`) + separated by equality operators (`==`, `!=`) Returns: Expr: the parsed expression @@ -356,7 +422,10 @@ class MidasParser(Parser): return expr def comparison(self) -> Expr: - """Parse a logical comparison expression or a simpler expression + """Parse a comparison expression + + A comparison consists of one or more term expressions (see :func:`term`) + separated by comparison operators (`<`, `<=`, `>`, `>=`) Returns: Expr: the parsed expression @@ -377,6 +446,14 @@ class MidasParser(Parser): return expr def term(self) -> Expr: + """Parse a term expression + + A term consists of one or more factor expressions (see :func:`factor`) + separated by weak arithmetic operators (`+`, `-`) + + Returns: + Expr: the parsed expression + """ expr: Expr = self.factor() while self.match(TokenType.PLUS, TokenType.MINUS): operator: Token = self.previous() @@ -388,6 +465,14 @@ class MidasParser(Parser): return expr def factor(self) -> Expr: + """Parse a factor expression + + A factor consists of one or more unary expressions (see :func:`unary`) + separated by strong arithmetic operators (`*`, `/`) + + Returns: + Expr: the parsed expression + """ expr: Expr = self.unary() while self.match(TokenType.STAR, TokenType.SLASH): operator: Token = self.previous() @@ -399,12 +484,15 @@ class MidasParser(Parser): return expr def unary(self) -> Expr: - """Parse a unary expression or a simpler expression + """Parse a unary expression + + A unary consists of a call expression (see :func:`call`) optionally + preceded by zero or more unary operators (`+`, `-`) Returns: Expr: the parsed expression """ - if self.match(TokenType.MINUS): + if self.match(TokenType.PLUS, TokenType.MINUS): operator: Token = self.previous() right: Expr = self.unary() location: Location = Location.span(operator.get_location(), right.location) @@ -412,12 +500,44 @@ class MidasParser(Parser): return self.call() def call(self) -> Expr: + """Parse a call expression + + A call consists of a reference expression (see :func:`reference`) + optionally followed by zero or more argument groups. + + Argument groups are parenthesize, comma-separated list of arguments (see :func:`finish_call`) + + Returns: + Expr: the parsed expression + """ expr: Expr = self.reference() while self.match(TokenType.LEFT_PAREN): expr = self.finish_call(expr) return expr def finish_call(self, callee: Expr) -> Expr: + """Parse an argument group, i.e. the arguments of a call + + Arguments are either passed positionally or by name (keyword argument). + All positional arguments must come before any keyword argument and + vice-versa. Arguments are separated by commas. + + A positional argument simply consists of an expression (see :func:`expression`) + + A keyword argument consists of and identifier, followed by the equal `=` + token and an expression (see :func:`expression`). + + Args: + callee (Expr): the callee expression + + Raises: + ParsingError: if a positional argument is passed after a keyword + argument or if a keyword argument's name is invalid (i.e. not + an identifier) + + Returns: + Expr: the parsed call expression + """ pos_args: list[Expr] = [] kw_args: dict[str, Expr] = {} keywords: bool = False @@ -437,13 +557,14 @@ class MidasParser(Parser): else: value = self.expression() if self.check(TokenType.EQUAL): + error_msg: str if keywords: - raise self.error(self.peek(), "Invalid keyword argument name") + error_msg = "Invalid keyword argument name" else: - raise self.error( - self.peek(), - "Cannot pass positional arguments after a keyword argument", + error_msg = ( + "Cannot pass positional arguments after a keyword argument" ) + raise self.error(self.peek(), error_msg) pos_args.append(value) if not self.match(TokenType.COMMA): @@ -460,7 +581,12 @@ class MidasParser(Parser): ) def reference(self) -> Expr: - """Parse an attribute access expression or a simpler expression + """Parse a reference expression + + A reference consists of a primary expression (see :func:`primary`) + optionally followed by zero or more attribute accesses. + + An attribute access consists of a dot `.` token followed by an identifier Returns: Expr: the parsed expression @@ -475,7 +601,12 @@ class MidasParser(Parser): def primary(self) -> Expr: """Parse a primary expression - This includes literals (booleans, numbers, etc.), wildcards, identifiers and grouped expressions + This includes literals (booleans, numbers, etc.), wildcards, identifiers + and grouped expressions + + Raises: + ParsingError: if a primary expressions cannot be parsed from the + following tokens Returns: Expr: the parsed expression @@ -508,14 +639,41 @@ class MidasParser(Parser): raise self.error(self.peek(), "Expected expression") def consume_identifier(self, message: str = "Expected identifier") -> Token: + """Consume the current token if it is a valid identifier or raise an error (see :func:`check_identifier`) + + If the current token is not a valid identifier, an error is raised + with the provided message + + Args: + message (str, optional): the error message. Defaults to "Expected identifier". + + Raises: + ParsingError: if the current token is not a valid identifier + + Returns: + Token: the current token which is a valid identifier + """ if not self.match_identifier(): raise self.error(self.peek(), message) return self.previous() def match_identifier(self) -> bool: + """Consume the next token if it is a valid identifier (see :func:`check_identifier`) + + Returns: + bool: whether a token was matched and consumed + """ return self.match(TokenType.IDENTIFIER, *KEYWORDS.values()) def check_identifier(self) -> bool: + """Check whether the current token is a valid identifier + + A valid identifier is either an identifier token or a keyword token. + This function always returns False if the parser is at the EOF token + + Returns: + bool: True if the current token is a valid identifier and not EOF + """ for tt in [TokenType.IDENTIFIER, *KEYWORDS.values()]: if self.check(tt): return True @@ -524,7 +682,14 @@ class MidasParser(Parser): def member_stmt(self) -> MemberStmt: """Parse a member statement - A type member statement is written `prop name: Type` or `def name: Type` + A member statement is written consists of: + - the `prop` (for a property) or `def` (for a method) keyword + - an name (identifier) + - a colon `:` + - a type expression (see :func:`type_expr`) + + Raises: + ParsingError: if the first token is neither `prop` nor `def` Returns: MemberStmt: the parsed member statement @@ -551,7 +716,13 @@ class MidasParser(Parser): def extend_declaration(self) -> ExtendStmt: """Parse an extension definition - An extension is written `extend Type { operations }` or `extend[S <: T, U] Type { operations }` + An extension statement consists of: + - the `extend` keyword + - a type name (identifier) + - (optional) type parameters (see :func:`type_params`) + - an opening brace `{` + - zero or more member statements (see :func:`member_stmt`) + - a closing brace `}` Returns: ExtendStmt: the parsed extension statement @@ -576,7 +747,12 @@ class MidasParser(Parser): def predicate_declaration(self) -> PredicateStmt: """Parse a predicate declaration - A predicate is written `predicate Name(subject: Type) = constraint_expression` + A predicate statement consists of: + - the `predicate` keyword + - a name (identifier) + - (optional) zero or more parameter specs (see :func:`function_params`) + - an equal sign `=` + - a body, a constraint expression (see :func:`constraint`) Returns: PredicateStmt: the parsed predicate declaration statement @@ -587,7 +763,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 +775,18 @@ class MidasParser(Parser): ) def function(self) -> FunctionType: - params: ParamSpec = self.function_args() + """Parse a function type expression + + A function consists of: + - the `fn` keyword + - a parameter spec (see :func:`function_params`) + - the arrow keyword `->` + - a result type expression (see :func:`type_expr`) + + Returns: + FunctionType: the parsed function type expression + """ + params: ParamSpec = self.function_params() self.consume(TokenType.ARROW, "Expected '->' before result type") result: Type = self.type_expr() @@ -610,36 +797,53 @@ class MidasParser(Parser): returns=result, ) - def function_args(self) -> ParamSpec: + def function_params(self) -> ParamSpec: + """Parse a parameter spec + + A parameter spec consists of zero or more comma-separated parameters, + wrapped in parentheses. + + Like in Python, it can contain positional-only, mixed and keyword-only + parameters (separated by `/` and `*`). + + Each parameter has a type (see :func:`type_expr`), + preceded by a name (identifier) and a colon `:` (not required for + positional-only parameters). + + Returns: + ParamSpec: the parsed parameter spec + """ l_paren: Token = self.consume( TokenType.LEFT_PAREN, "Expected '(' before function parameters" ) - pos_args: list[FunctionType.Argument] = [] - args: list[FunctionType.Argument] = [] - kw_args: list[FunctionType.Argument] = [] + pos: list[FunctionType.Parameter] = [] + mixed: list[FunctionType.Parameter] = [] + 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 - args = [] - args_first_tokens = [] + pos = mixed + mixed = [] + 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 +851,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): - if arg.name is None: + for param, token in zip(mixed, mixed_first_tokens): + 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) diff --git a/midas/parser/python.py b/midas/parser/python.py index a6c9a78..966ab16 100644 --- a/midas/parser/python.py +++ b/midas/parser/python.py @@ -23,6 +23,7 @@ from midas.ast.python import ( LiteralExpr, LogicalExpr, MidasType, + ParamSpec, RawExpr, RawStmt, ReturnStmt, @@ -49,6 +50,8 @@ class UnsupportedSyntaxError(Exception): class PythonParser: + """A parser to convert raw Python `ast` nodes in custom IR nodes""" + CAST_FUNCTION = "cast" UNSAFE_CAST_FUNCTION = "unsafe_cast" @@ -212,27 +215,10 @@ class PythonParser: match node: case ast.FunctionDef( name=name, - args=ast.arguments( - posonlyargs=posonlyargs, - args=args, - vararg=sink, - kwonlyargs=kwonlyargs, - kwarg=kw_sink, - defaults=defaults, - kw_defaults=kw_defaults, - ), + args=args, 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 +227,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), - 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 - ), + params=self._parse_param_spec(args), 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, diff --git a/tests/cases/checker/03_functions.py.ref.json b/tests/cases/checker/03_functions.py.ref.json index 820a7ea..bd2f5a9 100644 --- a/tests/cases/checker/03_functions.py.ref.json +++ b/tests/cases/checker/03_functions.py.ref.json @@ -124,7 +124,7 @@ 22 ] }, - "message": "Multiple values for argument 'b'" + "message": "Multiple values for parameter 'b'" }, { "type": "Error", @@ -152,7 +152,7 @@ 12 ] }, - "message": "Unknown keyword argument 'a'" + "message": "Unknown keyword parameter 'a'" }, { "type": "Error", @@ -194,7 +194,7 @@ 17 ] }, - "message": "Unknown keyword argument 'g'" + "message": "Unknown keyword parameter 'g'" }, { "type": "Error", @@ -277,36 +277,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -351,36 +353,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -443,36 +447,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -539,36 +545,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -649,36 +657,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -762,36 +772,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -850,36 +862,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -929,36 +943,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -1034,36 +1050,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -1150,36 +1168,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -1266,36 +1286,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } @@ -1382,36 +1404,38 @@ "name": "foo" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "int" - }, - "required": true - } - ], - "args": [ - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [ - { - "pos": 2, - "name": "c", - "type": { - "name": "str" - }, - "required": true - } - ], + "params": { + "pos": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "int" + }, + "required": true + } + ], + "mixed": [ + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [ + { + "pos": 2, + "name": "c", + "type": { + "name": "str" + }, + "required": true + } + ] + }, "returns": { "name": "bool" } diff --git a/tests/cases/checker/06_subtyping.py.ref.json b/tests/cases/checker/06_subtyping.py.ref.json index 03194f1..fc56069 100644 --- a/tests/cases/checker/06_subtyping.py.ref.json +++ b/tests/cases/checker/06_subtyping.py.ref.json @@ -136,26 +136,28 @@ "name": "maximum" }, "type": { - "pos_args": [], - "args": [ - { - "pos": 0, - "name": "a", - "type": { - "name": "float" + "params": { + "pos": [], + "mixed": [ + { + "pos": 0, + "name": "a", + "type": { + "name": "float" + }, + "required": true }, - "required": true - }, - { - "pos": 1, - "name": "b", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [], + { + "pos": 1, + "name": "b", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [] + }, "returns": { "name": "float" } diff --git a/tests/cases/checker/07_variance.py.ref.json b/tests/cases/checker/07_variance.py.ref.json index abfebf3..b34c2e6 100644 --- a/tests/cases/checker/07_variance.py.ref.json +++ b/tests/cases/checker/07_variance.py.ref.json @@ -312,16 +312,18 @@ "name": "print" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "object", - "type": {}, - "required": false - } - ], - "args": [], - "kw_args": [], + "params": { + "pos": [ + { + "pos": 0, + "name": "object", + "type": {}, + "required": false + } + ], + "mixed": [], + "kw": [] + }, "returns": {} } }, diff --git a/tests/cases/checker/08_unification.py.ref.json b/tests/cases/checker/08_unification.py.ref.json index c384d8a..27309d4 100644 --- a/tests/cases/checker/08_unification.py.ref.json +++ b/tests/cases/checker/08_unification.py.ref.json @@ -120,16 +120,18 @@ "name": "bool" }, "type": { - "pos_args": [ - { - "pos": 0, - "name": "object", - "type": {}, - "required": false - } - ], - "args": [], - "kw_args": [], + "params": { + "pos": [ + { + "pos": 0, + "name": "object", + "type": {}, + "required": false + } + ], + "mixed": [], + "kw": [] + }, "returns": { "name": "bool" } @@ -377,18 +379,20 @@ "name": "double" }, "type": { - "pos_args": [], - "args": [ - { - "pos": 0, - "name": "value", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [], + "params": { + "pos": [], + "mixed": [ + { + "pos": 0, + "name": "value", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [] + }, "returns": { "name": "float" } @@ -439,54 +443,58 @@ } ], "body": { - "pos_args": [ - { - "pos": 0, - "name": "transform", - "type": { - "pos_args": [ - { - "pos": 0, - "name": "v", - "type": { - "name": "T", - "bound": null, - "variance": "INVARIANT" - }, - "required": true - } - ], - "args": [], - "kw_args": [], - "returns": { - "name": "U", - "bound": null, - "variance": "INVARIANT" - } - }, - "required": true - }, - { - "pos": 1, - "name": "iterable", - "type": { - "name": "list", - "args": [ - { - "name": "T", + "params": { + "pos": [ + { + "pos": 0, + "name": "transform", + "type": { + "params": { + "pos": [ + { + "pos": 0, + "name": "v", + "type": { + "name": "T", + "bound": null, + "variance": "INVARIANT" + }, + "required": true + } + ], + "mixed": [], + "kw": [] + }, + "returns": { + "name": "U", "bound": null, "variance": "INVARIANT" } - ], - "body": { - "name": "list" - } + }, + "required": true }, - "required": true - } - ], - "args": [], - "kw_args": [], + { + "pos": 1, + "name": "iterable", + "type": { + "name": "list", + "args": [ + { + "name": "T", + "bound": null, + "variance": "INVARIANT" + } + ], + "body": { + "name": "list" + } + }, + "required": true + } + ], + "mixed": [], + "kw": [] + }, "returns": { "name": "list", "args": [ @@ -548,18 +556,20 @@ "name": "double" }, "type": { - "pos_args": [], - "args": [ - { - "pos": 0, - "name": "value", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [], + "params": { + "pos": [], + "mixed": [ + { + "pos": 0, + "name": "value", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw": [] + }, "returns": { "name": "float" } @@ -610,54 +620,58 @@ } ], "body": { - "pos_args": [ - { - "pos": 0, - "name": "transform", - "type": { - "pos_args": [ - { - "pos": 0, - "name": "v", - "type": { - "name": "T", - "bound": null, - "variance": "INVARIANT" - }, - "required": true - } - ], - "args": [], - "kw_args": [], - "returns": { - "name": "U", - "bound": null, - "variance": "INVARIANT" - } - }, - "required": true - }, - { - "pos": 1, - "name": "iterable", - "type": { - "name": "list", - "args": [ - { - "name": "T", + "params": { + "pos": [ + { + "pos": 0, + "name": "transform", + "type": { + "params": { + "pos": [ + { + "pos": 0, + "name": "v", + "type": { + "name": "T", + "bound": null, + "variance": "INVARIANT" + }, + "required": true + } + ], + "mixed": [], + "kw": [] + }, + "returns": { + "name": "U", "bound": null, "variance": "INVARIANT" } - ], - "body": { - "name": "list" - } + }, + "required": true }, - "required": true - } - ], - "args": [], - "kw_args": [], + { + "pos": 1, + "name": "iterable", + "type": { + "name": "list", + "args": [ + { + "name": "T", + "bound": null, + "variance": "INVARIANT" + } + ], + "body": { + "name": "list" + } + }, + "required": true + } + ], + "mixed": [], + "kw": [] + }, "returns": { "name": "list", "args": [ @@ -709,18 +723,20 @@ "name": "is_odd" }, "type": { - "pos_args": [], - "args": [ - { - "pos": 0, - "name": "value", - "type": { - "name": "int" - }, - "required": true - } - ], - "kw_args": [], + "params": { + "pos": [], + "mixed": [ + { + "pos": 0, + "name": "value", + "type": { + "name": "int" + }, + "required": true + } + ], + "kw": [] + }, "returns": { "name": "bool" } @@ -771,54 +787,58 @@ } ], "body": { - "pos_args": [ - { - "pos": 0, - "name": "transform", - "type": { - "pos_args": [ - { - "pos": 0, - "name": "v", - "type": { - "name": "T", - "bound": null, - "variance": "INVARIANT" - }, - "required": true - } - ], - "args": [], - "kw_args": [], - "returns": { - "name": "U", - "bound": null, - "variance": "INVARIANT" - } - }, - "required": true - }, - { - "pos": 1, - "name": "iterable", - "type": { - "name": "list", - "args": [ - { - "name": "T", + "params": { + "pos": [ + { + "pos": 0, + "name": "transform", + "type": { + "params": { + "pos": [ + { + "pos": 0, + "name": "v", + "type": { + "name": "T", + "bound": null, + "variance": "INVARIANT" + }, + "required": true + } + ], + "mixed": [], + "kw": [] + }, + "returns": { + "name": "U", "bound": null, "variance": "INVARIANT" } - ], - "body": { - "name": "list" - } + }, + "required": true }, - "required": true - } - ], - "args": [], - "kw_args": [], + { + "pos": 1, + "name": "iterable", + "type": { + "name": "list", + "args": [ + { + "name": "T", + "bound": null, + "variance": "INVARIANT" + } + ], + "body": { + "name": "list" + } + }, + "required": true + } + ], + "mixed": [], + "kw": [] + }, "returns": { "name": "list", "args": [ diff --git a/tests/cases/python-parser/03_functions.py.ref.json b/tests/cases/python-parser/03_functions.py.ref.json index 2288b82..64e8c00 100644 --- a/tests/cases/python-parser/03_functions.py.ref.json +++ b/tests/cases/python-parser/03_functions.py.ref.json @@ -7,50 +7,51 @@ { "_type": "Function", "name": "func", - "posonlyargs": [], - "args": [ - { - "name": "col1", - "type": { - "_type": "BaseType", - "base": "Column", - "args": [ - { - "_type": "ConstraintType", - "type": { - "_type": "BaseType", - "base": "float", - "args": [] - }, - "constraint": "0 <= _ <= 1" - } - ] + "params": { + "_type": "ParamSpec", + "pos": [], + "mixed": [ + { + "name": "col1", + "type": { + "_type": "BaseType", + "base": "Column", + "args": [ + { + "_type": "ConstraintType", + "type": { + "_type": "BaseType", + "base": "float", + "args": [] + }, + "constraint": "0 <= _ <= 1" + } + ] + }, + "default": null }, - "default": null - }, - { - "name": "col2", - "type": { - "_type": "BaseType", - "base": "Column", - "args": [ - { - "_type": "ConstraintType", - "type": { - "_type": "BaseType", - "base": "float", - "args": [] - }, - "constraint": "0 <= _ <= 1" - } - ] - }, - "default": null - } - ], - "sink": null, - "kwonlyargs": [], - "kw_sink": null, + { + "name": "col2", + "type": { + "_type": "BaseType", + "base": "Column", + "args": [ + { + "_type": "ConstraintType", + "type": { + "_type": "BaseType", + "base": "float", + "args": [] + }, + "constraint": "0 <= _ <= 1" + } + ] + }, + "default": null + } + ], + "kw": [] + }, "returns": { "_type": "BaseType", "base": "Column", @@ -119,41 +120,42 @@ { "_type": "Function", "name": "func2", - "posonlyargs": [ - { - "name": "a", - "type": { - "_type": "BaseType", - "base": "int", - "args": [] - }, - "default": null - } - ], - "args": [ - { - "name": "b", - "type": { - "_type": "BaseType", - "base": "float", - "args": [] - }, - "default": null - } - ], - "sink": null, - "kwonlyargs": [ - { - "name": "c", - "type": { - "_type": "BaseType", - "base": "str", - "args": [] - }, - "default": null - } - ], - "kw_sink": null, + "params": { + "_type": "ParamSpec", + "pos": [ + { + "name": "a", + "type": { + "_type": "BaseType", + "base": "int", + "args": [] + }, + "default": null + } + ], + "mixed": [ + { + "name": "b", + "type": { + "_type": "BaseType", + "base": "float", + "args": [] + }, + "default": null + } + ], + "kw": [ + { + "name": "c", + "type": { + "_type": "BaseType", + "base": "str", + "args": [] + }, + "default": null + } + ] + }, "returns": null, "body": [] } diff --git a/tests/serializer/midas.py b/tests/serializer/midas.py index 081059c..13fcbaa 100644 --- a/tests/serializer/midas.py +++ b/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], - "mixed": [self._serialize_func_arg(arg) for arg in spec.mixed], - "kw": [self._serialize_func_arg(arg) for arg in spec.kw], + "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_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, - "type": arg.type.accept(self), - "required": arg.required, + "name": param.name.lexeme if param.name is not None else None, + "type": param.type.accept(self), + "required": param.required, } def visit_extension_type(self, type: ExtensionType) -> dict: diff --git a/tests/serializer/python.py b/tests/serializer/python.py index 4e531cb..afb5521 100644 --- a/tests/serializer/python.py +++ b/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], - "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 - ), + "params": self._serialize_param_spec(stmt.params), "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",