HEL/midas
1
0
Fork 0
Code Issues Pull Requests Actions Releases 2 Activity

Merge pull request 'Add docstrings, rebrand parameters and refactor printers' (#29) from feat/add-docstrings into main

Browse Source 
Reviewed-on: #29
This commit was merged in pull request #29.
This commit is contained in:
Louis Heredero
2026-07-06 09:29:43 +00:00
parent 03bc32400b aae3073744
commit a418fe3eaf
51 changed files with 3482 additions and 2252 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
gen/gen.py
View File

@@ -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 import re
from pathlib import Path from pathlib import Path

8
gen/midas.py
View File

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

27
gen/python.py
View File

@@ -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 ###> MidasType | Type annotations | node
class BaseType: class BaseType:
base: str base: str
@@ -42,25 +57,17 @@ class ExpressionStmt:
class Function: class Function:
name: str name: str
posonlyargs: list[Argument] params: ParamSpec
args: list[Argument]
sink: Optional[Argument]
kwonlyargs: list[Argument]
kw_sink: Optional[Argument]
returns: Optional[MidasType] returns: Optional[MidasType]
body: list[Stmt] body: list[Stmt]
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class Argument: class Parameter:
location: Optional[Location] = None location: Optional[Location] = None
name: str name: str
type: Optional[MidasType] type: Optional[MidasType]
default: Optional[Expr] default: Optional[Expr]
@property
def all_args(self) -> list[Argument]:
return self.posonlyargs + self.args + self.kwonlyargs
class TypeAssign: class TypeAssign:
name: str name: str

12
midas/ast/location.py
View File

@@ -13,6 +13,8 @@ class HasLocation(Protocol):
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class Location: class Location:
"""Information about the location of an AST node"""
lineno: int lineno: int
col_offset: int col_offset: int
end_lineno: Optional[int] end_lineno: Optional[int]
@@ -29,6 +31,16 @@ class Location:
@staticmethod @staticmethod
def span(start: Location, end: Location) -> Location: 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( return Location(
lineno=start.lineno, lineno=start.lineno,
col_offset=start.col_offset, col_offset=start.col_offset,

8
midas/ast/midas.py
View File

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

896
midas/ast/printer.py
View File

@@ -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)}")

3
midas/ast/printer/__init__.py Normal file
View File

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

103
midas/ast/printer/base.py Normal file
View File

@@ -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)

183
midas/ast/printer/midas.py Normal file
View File

@@ -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)}"

253
midas/ast/printer/midas_ast.py Normal file
View File

@@ -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)

285
midas/ast/printer/python_ast.py Normal file
View File

@@ -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)}")

22
midas/ast/python.py
View File

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

7
midas/checker/builtins.py
View File

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

5
midas/checker/checker.py
View File

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

18
midas/checker/diagnostic.py
View File

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

144
midas/checker/dispatcher.py
View File

@@ -26,22 +26,29 @@ class HasLocation(Protocol):
E = TypeVar("E", bound=HasLocation) E = TypeVar("E", bound=HasLocation)
TypedExpr = tuple[E, Type] TypedExpr = tuple[E, Type]
"""An expression and its type"""
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class MappedArgument(Generic[E]): class MappedArgument(Generic[E]):
expr: E """An argument passed in a call and the corresponding parameter"""
type: Type
argument: Function.Argument arg_expr: E
arg_type: Type
parameter: Function.Parameter
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class OverloadCandidate: class OverloadCandidate:
"""An overloaded function call candidate with its mapped arguments"""
function: Function function: Function
mapped: list[MappedArgument] mapped: list[MappedArgument]
class CallError(StrEnum): class CallError(StrEnum):
"""Reason of a call error"""
INVALID_ARGS = "Invalid arguments" INVALID_ARGS = "Invalid arguments"
NO_MATCHING_OVERLOAD = "No matching overload" NO_MATCHING_OVERLOAD = "No matching overload"
IMPOSSIBLE_UNIFICATION = "Parameters unification failed" IMPOSSIBLE_UNIFICATION = "Parameters unification failed"
@@ -50,16 +57,28 @@ class CallError(StrEnum):
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class CallResult: class CallResult:
"""The result of a function call
Holds a return type, an optional error reason and message
"""
error: Optional[CallError] = None error: Optional[CallError] = None
"""The reason of the error, if there is one"""
result: Type = UnknownType() result: Type = UnknownType()
"""The result type. `UnknownType()` if the call is invalid"""
message: Optional[str] = None message: Optional[str] = None
"""An optional error message"""
@property @property
def is_valid(self) -> bool: def is_valid(self) -> bool:
"""Whether the call is valid (i.e. no error)"""
return self.error is None return self.error is None
@property @property
def error_message(self) -> str: def error_message(self) -> str:
"""A descriptive message for the error, if there is one"""
if self.message is not None: if self.message is not None:
return self.message return self.message
if self.error is not None: if self.error is not None:
@@ -68,6 +87,15 @@ class CallResult:
class CallDispatcher(Generic[E]): 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: def __init__(self, types: TypesRegistry, reporter: FileReporter) -> None:
self.types: TypesRegistry = types self.types: TypesRegistry = types
self.reporter: FileReporter = reporter self.reporter: FileReporter = reporter
@@ -86,22 +114,21 @@ class CallDispatcher(Generic[E]):
) -> CallResult: ) -> CallResult:
"""Get the result type of a function call """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. appropriate signature.
Argument types are matched to the defined parameters. Argument types are matched with the defined parameters.
The function doesn't take the raw expression as a parameter to accommodate This function doesn't take the raw expression as a parameter to
for desugared calls such as for operators. accommodate for desugared calls such as for operators.
Args: Args:
location (Location): the call location location (Location): the call location
callee (Type): the called function 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 keywords (dict[str, TypedExpr]): the map of keyword arguments
report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True. report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
Returns: Returns:
Type: the return type of the call, or `None` if either CallResult: the call result, either a type or an error
the call is invalid or no overload matched the arguments uniquely
""" """
match callee: match callee:
case Function() as function: case Function() as function:
@@ -179,6 +206,18 @@ class CallDispatcher(Generic[E]):
positional: list[TypedExpr[E]], positional: list[TypedExpr[E]],
keywords: dict[str, TypedExpr[E]], keywords: dict[str, TypedExpr[E]],
) -> Union[tuple[Function, None], tuple[None, CallError]]: ) -> 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: match callee:
case DerivedType(type=base): case DerivedType(type=base):
return self._unwrap_function(base, positional, keywords) return self._unwrap_function(base, positional, keywords)
@@ -219,11 +258,11 @@ class CallDispatcher(Generic[E]):
""" """
valid: bool = True valid: bool = True
for arg in arguments: 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: if report_errors:
self.reporter.error( self.reporter.error(
arg.expr.location, arg.arg_expr.location,
f"Wrong type for argument '{arg.argument.name}', expected {arg.argument.type}, got {arg.type}", f"Wrong type for argument '{arg.parameter.name}', expected {arg.parameter.type}, got {arg.arg_type}",
) )
valid = False valid = False
return valid 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. report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
Returns: Returns:
Optional[Function]: the resolved function signature if it can be Union[tuple[Function, None], tuple[None, str]]: a tuple containing
determined unambiguously, or `None`. the resolved function signature if it can be determined
unambiguously, or `None`, and an error message, or `None`
""" """
candidates: list[OverloadCandidate] = [] candidates: list[OverloadCandidate] = []
errors: list[CallError] = [] errors: list[CallError] = []
@@ -347,28 +387,30 @@ class CallDispatcher(Generic[E]):
tuple[bool, list[MappedArgument]]: a boolean reporting whether 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] = [ 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] = [ 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]] = [] mapped: list[MappedArgument[E]] = []
pos_params: list[Function.Argument] = list(function.pos_args) pos_params: list[Function.Parameter] = list(function.params.pos)
mixed_params: list[Function.Argument] = list(function.args) mixed_params: list[Function.Parameter] = list(function.params.mixed)
kw_params: dict[str, Function.Argument] = { kw_params: dict[str, Function.Parameter] = {
arg.name: arg for arg in function.kw_args param.name: param for param in function.params.kw
} }
valid_call: bool = True valid_call: bool = True
# TODO: handle *args and **kwargs sinks # TODO: handle *args and **kwargs sinks
for arg in positional: for arg in positional:
param: Function.Argument param: Function.Parameter
if len(pos_params) != 0: if len(pos_params) != 0:
param = pos_params.pop(0) param = pos_params.pop(0)
elif len(mixed_params) != 0: elif len(mixed_params) != 0:
@@ -385,27 +427,27 @@ class CallDispatcher(Generic[E]):
required_positional.remove(name) required_positional.remove(name)
if name in required_keyword: if name in required_keyword:
required_keyword.remove(name) required_keyword.remove(name)
set_args.add(name) set_params.add(name)
mapped.append( mapped.append(
MappedArgument( MappedArgument(
expr=arg[0], arg_expr=arg[0],
type=arg[1], arg_type=arg[1],
argument=param, parameter=param,
) )
) )
kw_params.update({arg.name: arg for arg in mixed_params}) kw_params.update({param.name: param for param in mixed_params})
for name, arg in keywords.items(): for name, arg in keywords.items():
param: Function.Argument param: Function.Parameter
if name not in kw_params: if name not in kw_params:
if report_errors: if report_errors:
if name in set_args: if name in set_params:
self.reporter.error( self.reporter.error(
arg[0].location, f"Multiple values for argument '{name}'" arg[0].location, f"Multiple values for parameter '{name}'"
) )
else: else:
self.reporter.error( self.reporter.error(
arg[0].location, f"Unknown keyword argument '{name}'" arg[0].location, f"Unknown keyword parameter '{name}'"
) )
valid_call = False valid_call = False
continue continue
@@ -414,40 +456,40 @@ class CallDispatcher(Generic[E]):
required_positional.remove(name) required_positional.remove(name)
if name in required_keyword: if name in required_keyword:
required_keyword.remove(name) required_keyword.remove(name)
set_args.add(name) set_params.add(name)
mapped.append( mapped.append(
MappedArgument( MappedArgument(
expr=arg[0], arg_expr=arg[0],
type=arg[1], arg_type=arg[1],
argument=param, parameter=param,
) )
) )
def join_args(args: list[str]) -> str: def join_params(params: list[str]) -> str:
args = list(map(lambda a: f"'{a}'", args)) params = list(map(lambda p: f"'{p}'", params))
if len(args) == 0: if len(params) == 0:
return "" return ""
if len(args) == 1: if len(params) == 1:
return args[0] return params[0]
return ", ".join(args[:-1]) + " and " + args[-1] return ", ".join(params[:-1]) + " and " + params[-1]
if len(required_positional) != 0: if len(required_positional) != 0:
plural: str = "" if len(required_positional) == 1 else "s" plural: str = "" if len(required_positional) == 1 else "s"
args: str = join_args(required_positional) params: str = join_params(required_positional)
if report_errors: if report_errors:
self.reporter.error( self.reporter.error(
location, location,
f"Missing required positional argument{plural}: {args}", f"Missing required positional argument{plural}: {params}",
) )
valid_call = False valid_call = False
if len(required_keyword) != 0: if len(required_keyword) != 0:
plural: str = "" if len(required_keyword) == 1 else "s" plural: str = "" if len(required_keyword) == 1 else "s"
args: str = join_args(required_keyword) params: str = join_params(required_keyword)
if report_errors: if report_errors:
self.reporter.error( self.reporter.error(
location, location,
f"Missing required keyword argument{plural}: {args}", f"Missing required keyword argument{plural}: {params}",
) )
valid_call = False 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 `mapped2`. If any of the parameter type in `mapped1` is not a subtype
of the corresponding parameter in `mapped2`, `False` is returned. of the corresponding parameter in `mapped2`, `False` is returned.
This is used to check whether a given overload is This is used to check whether a given overload is a more specific
a more specific function/ a subtype of another. function / a subtype of another.
Args: Args:
mapped1 (list[MappedArgument]): the first argument mappings (subtype) mapped1 (list[MappedArgument]): the first argument mappings (subtype)
@@ -474,11 +516,11 @@ class CallDispatcher(Generic[E]):
""" """
by_expr: dict[E, Type] = {} by_expr: dict[E, Type] = {}
for arg in mapped1: for arg in mapped1:
by_expr[arg.expr] = arg.argument.type by_expr[arg.arg_expr] = arg.parameter.type
for arg in mapped2: for arg in mapped2:
type2: Type = arg.argument.type type2: Type = arg.parameter.type
type1: Type = by_expr[arg.expr] type1: Type = by_expr[arg.arg_expr]
if not self.types.is_subtype(type1, type2): if not self.types.is_subtype(type1, type2):
return False return False
return True return True

80
midas/checker/evaluator.py
View File

@@ -11,10 +11,18 @@ from midas.lexer.token import TokenType
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class PartialPredicate(Predicate): class PartialPredicate(Predicate):
"""A partially applied predicate"""
scope: dict[str, Any] scope: dict[str, Any]
"""A dictionary of already applied parameters"""
class Evaluator(m.Expr.Visitor[Any]): 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): def __init__(self, types: TypesRegistry, reporter: Optional[FileReporter] = None):
self.types: TypesRegistry = types self.types: TypesRegistry = types
self.reporter: Optional[FileReporter] = reporter self.reporter: Optional[FileReporter] = reporter
@@ -22,16 +30,51 @@ class Evaluator(m.Expr.Visitor[Any]):
self.scopes: list[dict[str, Any]] = [{}] self.scopes: list[dict[str, Any]] = [{}]
def evaluate(self, expr: m.Expr) -> 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) value: Any = expr.accept(self)
if self.reporter is not None: if self.reporter is not None:
self.reporter.debug(expr.location, f"Value: {value}") self.reporter.debug(expr.location, f"Value: {value}")
return value return value
def get_value(self, name: str) -> Any: 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] scope: dict[str, Any] = self.scopes[-1]
return scope[name] return scope[name]
def set_value(self, name: str, value: Any, force_declare: bool = False): 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: if not force_declare:
for scope in reversed(self.scopes): for scope in reversed(self.scopes):
if name in scope: if name in scope:
@@ -131,8 +174,21 @@ class Evaluator(m.Expr.Visitor[Any]):
return self.get_value("_") return self.get_value("_")
def _evaluate_predicate( def _evaluate_predicate(
self, predicate: Predicate, args: list[Any], kwargs: dict[str, Any] self,
predicate: Predicate,
args: list[Any],
kwargs: dict[str, Any],
) -> 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 res: Any = None
if isinstance(predicate, PartialPredicate): if isinstance(predicate, PartialPredicate):
self.scopes.append(predicate.scope) self.scopes.append(predicate.scope)
@@ -158,15 +214,27 @@ class Evaluator(m.Expr.Visitor[Any]):
return res return res
def _map_args(self, function: Function, args: list[Any], kwargs: dict[str, Any]): def _map_args(self, function: Function, args: list[Any], kwargs: dict[str, Any]):
positional: list[Function.Argument] = function.pos_args + function.args """Map call arguments to a function's parameters and set their values in context
keywords: dict[str, Function.Argument] = {
arg.name: arg for arg in function.args + function.kw_args 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): for i, arg in enumerate(args):
param: Function.Argument = positional[i] param: Function.Parameter = positional[i]
self.set_value(param.name, arg) self.set_value(param.name, arg)
for name, arg in kwargs.items(): for name, arg in kwargs.items():
param: Function.Argument = keywords[name] param: Function.Parameter = keywords[name]
self.set_value(param.name, arg) self.set_value(param.name, arg)

44
midas/checker/frames/column_groupby_methods.py
View File

@@ -7,7 +7,14 @@ import midas.ast.python as p
from midas.ast.location import Location from midas.ast.location import Location
from midas.checker.dispatcher import CallResult from midas.checker.dispatcher import CallResult
from midas.checker.frames.utils import MethodRegistry, method 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: if TYPE_CHECKING:
from midas.checker.python import TypedExpr from midas.checker.python import TypedExpr
@@ -15,6 +22,8 @@ if TYPE_CHECKING:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class Call: class Call:
"""A column group-by method call, implements :class:`utils.MethodCall`"""
location: Location location: Location
call_expr: p.Expr call_expr: p.Expr
groupby: ColumnGroupBy groupby: ColumnGroupBy
@@ -28,6 +37,8 @@ class Call:
class ColumnGroupByMethodRegistry(MethodRegistry[Call]): class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
"""The method registry for column group-by types"""
NAMED_ARGS: dict[str, str] = { NAMED_ARGS: dict[str, str] = {
"numeric_only": "bool", "numeric_only": "bool",
"skipna": "bool", "skipna": "bool",
@@ -38,31 +49,46 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
def _aggregate( def _aggregate(
self, self,
call: Call, call: Call,
args: list[str | tuple[str, str, bool]] = [], params: list[str | tuple[str, str, bool]] = [],
*, *,
preserve_inner_type: bool = False, preserve_inner_type: bool = False,
) -> Type: ) -> Type:
real_args: list[Function.Argument] = [] """Compute the result type of an aggregate method call
for i, arg in enumerate(args):
match arg: 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: case str() as name:
arg = Function.Argument( param = Function.Parameter(
pos=i, pos=i,
name=name, name=name,
type=self.types.get_type(self.NAMED_ARGS[name]), type=self.types.get_type(self.NAMED_ARGS[name]),
required=False, required=False,
) )
case (name, type, required): case (name, type, required):
arg = Function.Argument( param = Function.Parameter(
pos=i, pos=i,
name=name, name=name,
type=self.types.get_type(type), type=self.types.get_type(type),
required=required, required=required,
) )
real_args.append(arg) real_params.append(param)
signature = Function( signature = Function(
args=real_args, params=ParamSpec(mixed=real_params),
returns=( returns=(
call.groupby.column call.groupby.column
if preserve_inner_type if preserve_inner_type

39
midas/checker/frames/column_manager.py
View File

@@ -15,6 +15,8 @@ if TYPE_CHECKING:
class ColumnManager: class ColumnManager:
"""Helper class to handle methods and subscripts on column types"""
def __init__(self, typer: PythonTyper) -> None: def __init__(self, typer: PythonTyper) -> None:
self.typer: PythonTyper = typer self.typer: PythonTyper = typer
self.method_resolver: ColumnMethodRegistry = ColumnMethodRegistry(self.typer) self.method_resolver: ColumnMethodRegistry = ColumnMethodRegistry(self.typer)
@@ -32,6 +34,20 @@ class ColumnManager:
positional: list[TypedExpr], positional: list[TypedExpr],
keywords: dict[str, TypedExpr], keywords: dict[str, TypedExpr],
) -> Type: ) -> 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( call: Call = Call(
location=location, location=location,
call_expr=call_expr, call_expr=call_expr,
@@ -52,6 +68,20 @@ class ColumnManager:
positional: list[TypedExpr], positional: list[TypedExpr],
keywords: dict[str, TypedExpr], keywords: dict[str, TypedExpr],
) -> Type: ) -> 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( call: GroupByCall = GroupByCall(
location=location, location=location,
call_expr=call_expr, call_expr=call_expr,
@@ -63,6 +93,15 @@ class ColumnManager:
return self.groupby_method_resolver.call(method, call) return self.groupby_method_resolver.call(method, call)
def get_attribute(self, column: ColumnType, name: str) -> Optional[Type]: 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 types: TypesRegistry = self.typer.types
match name: match name:
case "ndim" | "size": case "ndim" | "size":

113
midas/checker/frames/column_methods.py
View File

@@ -13,6 +13,7 @@ from midas.checker.types import (
ColumnType, ColumnType,
Function, Function,
GenericType, GenericType,
ParamSpec,
TopType, TopType,
Type, Type,
TypeVar, TypeVar,
@@ -26,6 +27,8 @@ if TYPE_CHECKING:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class Call: class Call:
"""A column method call, implements :class:`utils.MethodCall`"""
location: Location location: Location
call_expr: p.Expr call_expr: p.Expr
column: ColumnType column: ColumnType
@@ -39,6 +42,8 @@ class Call:
class ColumnMethodRegistry(MethodRegistry[Call]): class ColumnMethodRegistry(MethodRegistry[Call]):
"""The method registry for column types"""
def _element_binary_op(self, call: Call, method: str) -> ColumnType: def _element_binary_op(self, call: Call, method: str) -> ColumnType:
"""Compute the result of an element-wise binary operation """Compute the result of an element-wise binary operation
@@ -74,6 +79,18 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
return new_column return new_column
def _element_wise(self, call: Call, method: str) -> Type: 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 # TODO: support add with scalar
# Build signature with new column type and generic operand # Build signature with new column type and generic operand
@@ -82,14 +99,16 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
name="add", name="add",
params=[param_type], params=[param_type],
body=Function( body=Function(
args=[ params=ParamSpec(
Function.Argument( mixed=[
Function.Parameter(
pos=0, pos=0,
name="other", name="other",
type=ColumnType(type=param_type), type=ColumnType(type=param_type),
required=True, required=True,
), ),
], ],
),
returns=self._element_binary_op(call, method), returns=self._element_binary_op(call, method),
), ),
) )
@@ -163,13 +182,27 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
def _aggregate( def _aggregate(
self, self,
call: Call, call: Call,
kwargs: list[Function.Argument] = [], kwargs: list[Function.Parameter] = [],
*, *,
preserve_inner_type: bool = False, preserve_inner_type: bool = False,
) -> Type: ) -> 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( signature = Function(
kw_args=[ params=ParamSpec(
Function.Argument( kw=[
Function.Parameter(
pos=0, pos=0,
name="axis", name="axis",
type=TopType(), type=TopType(),
@@ -177,6 +210,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
), ),
*kwargs, *kwargs,
], ],
),
returns=call.column if preserve_inner_type else ColumnType(type=TopType()), returns=call.column if preserve_inner_type else ColumnType(type=TopType()),
) )
@@ -221,7 +255,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
return self._aggregate( return self._aggregate(
call, call,
[ [
Function.Argument( Function.Parameter(
pos=1, pos=1,
name="ddof", name="ddof",
type=self.types.get_type("int"), type=self.types.get_type("int"),
@@ -239,7 +273,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
return self._aggregate( return self._aggregate(
call, call,
[ [
Function.Argument( Function.Parameter(
pos=1, pos=1,
name="var", name="var",
type=self.types.get_type("int"), type=self.types.get_type("int"),
@@ -251,14 +285,16 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
@method() @method()
def head(self, call: Call) -> Type: def head(self, call: Call) -> Type:
signature = Function( signature = Function(
args=[ params=ParamSpec(
Function.Argument( mixed=[
Function.Parameter(
pos=0, pos=0,
name="n", name="n",
type=self.types.get_type("int"), type=self.types.get_type("int"),
required=False, required=False,
), ),
], ],
),
returns=call.column, returns=call.column,
) )
@@ -273,14 +309,16 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
@method() @method()
def tail(self, call: Call) -> Type: def tail(self, call: Call) -> Type:
signature = Function( signature = Function(
args=[ params=ParamSpec(
Function.Argument( mixed=[
Function.Parameter(
pos=0, pos=0,
name="n", name="n",
type=self.types.get_type("int"), type=self.types.get_type("int"),
required=False, required=False,
), ),
], ],
),
returns=call.column, returns=call.column,
) )
@@ -296,52 +334,33 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
def groupby(self, call: Call) -> Type: def groupby(self, call: Call) -> Type:
bool_: Type = self.types.get_type("bool") bool_: Type = self.types.get_type("bool")
function: Function = Function( function: Function = Function(
args=[ params=ParamSpec(
Function.Argument( mixed=[
Function.Parameter(
pos=0, pos=0,
name="by", name="by",
type=TopType(), type=TopType(),
required=False, required=False,
), ),
Function.Argument( Function.Parameter(
pos=1, pos=1,
name="level", name="level",
type=TopType(), type=TopType(),
required=False, required=False,
), ),
], ],
kw_args=[ kw=[
Function.Argument( Function.Parameter(
pos=2, pos=i + 2,
name="as_index", name=name,
type=bool_, type=bool_,
required=False, required=False,
), )
Function.Argument( for i, name in enumerate(
pos=3, ["as_index", "sort", "group_keys", "observed", "dropna"]
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,
),
], ],
),
returns=ColumnGroupBy(column=call.column), returns=ColumnGroupBy(column=call.column),
) )
@@ -354,6 +373,14 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
return result.result return result.result
def _assert_same_length(self, call_expr: p.Expr, column1: p.Expr, column2: p.Expr): 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__" func_name: str = "__midas_column_same_length__"
# Efficiently compute length # Efficiently compute length

18
midas/checker/frames/frame_groupby_methods.py
View File

@@ -21,6 +21,8 @@ if TYPE_CHECKING:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class Call: class Call:
"""A frame group-by method call, implements :class:`utils.MethodCall`"""
location: Location location: Location
call_expr: p.Expr call_expr: p.Expr
groupby: FrameGroupBy groupby: FrameGroupBy
@@ -34,14 +36,18 @@ class Call:
class FrameGroupByMethodRegistry(MethodRegistry[Call]): class FrameGroupByMethodRegistry(MethodRegistry[Call]):
NAMED_ARGS: dict[str, str] = { """The method registry for frame group-by types"""
"numeric_only": "bool",
"skipna": "bool",
"engine": "str",
"engine_kwargs": "dict",
}
def _aggregate(self, call: Call, method: str) -> Type: 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] = [] new_columns: list[DataFrameType.Column] = []
for column in call.groupby.frame.columns: for column in call.groupby.frame.columns:

131
midas/checker/frames/frame_manager.py
View File

@@ -24,10 +24,20 @@ if TYPE_CHECKING:
def is_list_of_literals(exprs: list[p.Expr]) -> TypeGuard[list[p.LiteralExpr]]: 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) return all(isinstance(expr, p.LiteralExpr) for expr in exprs)
class FrameManager: class FrameManager:
"""Helper class to handle methods and subscripts on frame types"""
def __init__(self, typer: PythonTyper) -> None: def __init__(self, typer: PythonTyper) -> None:
self.typer: PythonTyper = typer self.typer: PythonTyper = typer
self.method_resolver: FrameMethodRegistry = FrameMethodRegistry(self.typer) self.method_resolver: FrameMethodRegistry = FrameMethodRegistry(self.typer)
@@ -43,6 +53,18 @@ class FrameManager:
index: p.Expr, index: p.Expr,
value_type: Type, value_type: 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: match index:
case p.LiteralExpr(value=str() as name): case p.LiteralExpr(value=str() as name):
return self.assign_column(reporter, location, frame, name, value_type) return self.assign_column(reporter, location, frame, name, value_type)
@@ -93,6 +115,18 @@ class FrameManager:
name: str, name: str,
type: Type, type: 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): if not isinstance(type, ColumnType):
reporter.error( reporter.error(
location, location,
@@ -108,6 +142,17 @@ class FrameManager:
frame: DataFrameType, frame: DataFrameType,
index: p.Expr, index: p.Expr,
) -> Type: ) -> 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: match index:
case p.LiteralExpr(value=str() as name): case p.LiteralExpr(value=str() as name):
column: Optional[ColumnType] = FrameManager._get_column(frame, name) column: Optional[ColumnType] = FrameManager._get_column(frame, name)
@@ -142,6 +187,17 @@ class FrameManager:
groupby: FrameGroupBy, groupby: FrameGroupBy,
index: p.Expr, index: p.Expr,
) -> Type: ) -> 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) result: Type = self.get(reporter, location, groupby.frame, index)
match result: match result:
case ColumnType(): case ColumnType():
@@ -159,6 +215,16 @@ class FrameManager:
def _set_column( def _set_column(
cls, frame: DataFrameType, name: str, column: ColumnType cls, frame: DataFrameType, name: str, column: ColumnType
) -> DataFrameType: ) -> 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] = [] new_columns: list[DataFrameType.Column] = []
index: int = len(frame.columns) index: int = len(frame.columns)
replace: bool = False replace: bool = False
@@ -185,12 +251,31 @@ class FrameManager:
def _set_columns( def _set_columns(
cls, frame: DataFrameType, names: list[str], columns: list[ColumnType] cls, frame: DataFrameType, names: list[str], columns: list[ColumnType]
) -> DataFrameType: ) -> 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): for name, col in zip(names, columns):
frame = cls._set_column(frame, name, col) frame = cls._set_column(frame, name, col)
return frame return frame
@classmethod @classmethod
def _get_column(cls, frame: DataFrameType, name: str) -> Optional[ColumnType]: 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: for col in frame.columns:
if col.name == name: if col.name == name:
return col.type return col.type
@@ -200,6 +285,15 @@ class FrameManager:
def _get_columns( def _get_columns(
cls, frame: DataFrameType, names: list[str] cls, frame: DataFrameType, names: list[str]
) -> list[Optional[ColumnType]]: ) -> 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] return [cls._get_column(frame, name) for name in names]
def call( def call(
@@ -212,6 +306,20 @@ class FrameManager:
positional: list[TypedExpr], positional: list[TypedExpr],
keywords: dict[str, TypedExpr], keywords: dict[str, TypedExpr],
) -> Type: ) -> 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( call: Call = Call(
location=location, location=location,
call_expr=call_expr, call_expr=call_expr,
@@ -232,6 +340,20 @@ class FrameManager:
positional: list[TypedExpr], positional: list[TypedExpr],
keywords: dict[str, TypedExpr], keywords: dict[str, TypedExpr],
) -> Type: ) -> 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( call: GroupByCall = GroupByCall(
location=location, location=location,
call_expr=call_expr, call_expr=call_expr,
@@ -243,6 +365,15 @@ class FrameManager:
return self.groupby_method_resolver.call(method, call) return self.groupby_method_resolver.call(method, call)
def get_attribute(self, frame: DataFrameType, name: str) -> Optional[Type]: 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 types: TypesRegistry = self.typer.types
match name: match name:
case "ndim" | "size": case "ndim" | "size":

116
midas/checker/frames/frame_methods.py
View File

@@ -14,6 +14,7 @@ from midas.checker.types import (
FrameGroupBy, FrameGroupBy,
Function, Function,
OverloadedFunction, OverloadedFunction,
ParamSpec,
TopType, TopType,
Type, Type,
UnknownType, UnknownType,
@@ -26,6 +27,8 @@ if TYPE_CHECKING:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class Call: class Call:
"""A frame method call, implements :class:`utils.MethodCall`"""
location: Location location: Location
call_expr: p.Expr call_expr: p.Expr
frame: DataFrameType frame: DataFrameType
@@ -39,6 +42,8 @@ class Call:
class FrameMethodRegistry(MethodRegistry[Call]): class FrameMethodRegistry(MethodRegistry[Call]):
"""The method registry for frame types"""
def _get_method_result( def _get_method_result(
self, self,
call: Call, call: Call,
@@ -147,17 +152,31 @@ class FrameMethodRegistry(MethodRegistry[Call]):
return DataFrameType(columns=new_columns) return DataFrameType(columns=new_columns)
def _element_wise(self, call: Call, method: str) -> Type: 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 # TODO: support scalar, sequence, Series, dict operand
# Build signature with new schema and generic operand # Build signature with new schema and generic operand
signature = Function( signature = Function(
args=[ params=ParamSpec(
Function.Argument( mixed=[
Function.Parameter(
pos=0, pos=0,
name="other", name="other",
type=DataFrameType(columns=[]), type=DataFrameType(columns=[]),
required=True, required=True,
), ),
], ],
),
returns=self._element_binary_op(call, method), returns=self._element_binary_op(call, method),
) )
@@ -227,10 +246,21 @@ class FrameMethodRegistry(MethodRegistry[Call]):
def eq(self, call: Call) -> Type: def eq(self, call: Call) -> Type:
return self._element_wise(call, "__eq__") 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( with_axis = Function(
kw_args=[ params=ParamSpec(
Function.Argument( kw=[
Function.Parameter(
pos=0, pos=0,
name="axis", name="axis",
type=self.types.get_type("int"), type=self.types.get_type("int"),
@@ -238,11 +268,13 @@ class FrameMethodRegistry(MethodRegistry[Call]):
), ),
*kwargs, *kwargs,
], ],
),
returns=ColumnType(type=TopType()), returns=ColumnType(type=TopType()),
) )
without_axis = Function( without_axis = Function(
kw_args=[ params=ParamSpec(
Function.Argument( kw=[
Function.Parameter(
pos=0, pos=0,
name="axis", name="axis",
type=self.types.get_type("None"), type=self.types.get_type("None"),
@@ -250,6 +282,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
), ),
*kwargs, *kwargs,
], ],
),
returns=TopType(), returns=TopType(),
) )
overload = OverloadedFunction( overload = OverloadedFunction(
@@ -300,7 +333,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
return self._aggregate( return self._aggregate(
call, call,
[ [
Function.Argument( Function.Parameter(
pos=1, pos=1,
name="ddof", name="ddof",
type=self.types.get_type("int"), type=self.types.get_type("int"),
@@ -318,7 +351,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
return self._aggregate( return self._aggregate(
call, call,
[ [
Function.Argument( Function.Parameter(
pos=1, pos=1,
name="var", name="var",
type=self.types.get_type("int"), type=self.types.get_type("int"),
@@ -330,14 +363,16 @@ class FrameMethodRegistry(MethodRegistry[Call]):
@method() @method()
def head(self, call: Call) -> Type: def head(self, call: Call) -> Type:
signature = Function( signature = Function(
args=[ params=ParamSpec(
Function.Argument( mixed=[
Function.Parameter(
pos=0, pos=0,
name="n", name="n",
type=self.types.get_type("int"), type=self.types.get_type("int"),
required=False, required=False,
), ),
], ],
),
returns=call.frame, returns=call.frame,
) )
@@ -352,14 +387,16 @@ class FrameMethodRegistry(MethodRegistry[Call]):
@method() @method()
def tail(self, call: Call) -> Type: def tail(self, call: Call) -> Type:
signature = Function( signature = Function(
args=[ params=ParamSpec(
Function.Argument( mixed=[
Function.Parameter(
pos=0, pos=0,
name="n", name="n",
type=self.types.get_type("int"), type=self.types.get_type("int"),
required=False, required=False,
), ),
], ],
),
returns=call.frame, returns=call.frame,
) )
@@ -375,52 +412,33 @@ class FrameMethodRegistry(MethodRegistry[Call]):
def groupby(self, call: Call) -> Type: def groupby(self, call: Call) -> Type:
bool_: Type = self.types.get_type("bool") bool_: Type = self.types.get_type("bool")
function: Function = Function( function: Function = Function(
args=[ params=ParamSpec(
Function.Argument( mixed=[
Function.Parameter(
pos=0, pos=0,
name="by", name="by",
type=TopType(), type=TopType(),
required=False, required=False,
), ),
Function.Argument( Function.Parameter(
pos=1, pos=1,
name="level", name="level",
type=TopType(), type=TopType(),
required=False, required=False,
), ),
], ],
kw_args=[ kw=[
Function.Argument( Function.Parameter(
pos=2, pos=i + 2,
name="as_index", name=name,
type=bool_, type=bool_,
required=False, required=False,
), )
Function.Argument( for i, name in enumerate(
pos=3, ["as_index", "sort", "group_keys", "observed", "dropna"]
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,
),
], ],
),
returns=FrameGroupBy(frame=call.frame), returns=FrameGroupBy(frame=call.frame),
) )
@@ -433,6 +451,14 @@ class FrameMethodRegistry(MethodRegistry[Call]):
return result.result return result.result
def _assert_same_length(self, call_expr: p.Expr, frame1: p.Expr, frame2: p.Expr): 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__" func_name: str = "__midas_frame_same_length__"
# Efficiently compute length # Efficiently compute length

29
midas/checker/frames/utils.py
View File

@@ -24,6 +24,12 @@ if TYPE_CHECKING:
class _MethodRegistryMeta(type): 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]] = {} _methods: dict[str, Callable[..., Type]] = {}
def __new__( def __new__(
@@ -42,6 +48,11 @@ class _MethodRegistryMeta(type):
class MethodCall(Protocol): class MethodCall(Protocol):
"""A method call object
Must have at least `location`, `call_expr` and `subject` properties
"""
@property @property
def location(self) -> Location: ... def location(self) -> Location: ...
@@ -56,6 +67,8 @@ T = TypeVar("T", bound=MethodCall)
class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta): class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta):
"""A registry of methods"""
def __init__(self, typer: PythonTyper) -> None: def __init__(self, typer: PythonTyper) -> None:
self.typer: PythonTyper = typer self.typer: PythonTyper = typer
@@ -76,6 +89,15 @@ class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta):
return self.typer.assertions return self.typer.assertions
def call(self, method: str, call: T) -> Type: 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) func: Optional[Callable[[Self, T], Type]] = self._methods.get(method)
if func is None: if func is None:
self.reporter.warning( self.reporter.warning(
@@ -90,6 +112,13 @@ Method = Callable[[_Self, T], Type]
def method(*names: str) -> Callable[[Method[_Self, T]], Method[_Self, T]]: 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]: def wrapper(func: Method[_Self, T]) -> Method[_Self, T]:
names_: tuple[str, ...] = names names_: tuple[str, ...] = names
if len(names_) == 0: if len(names_) == 0:

121
midas/checker/midas.py
View File

@@ -1,5 +1,4 @@
import logging import logging
from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Optional from typing import Optional
@@ -21,6 +20,7 @@ from midas.checker.types import (
ExtensionType, ExtensionType,
Function, Function,
GenericType, GenericType,
ParamSpec,
Predicate, Predicate,
Type, Type,
TypeVar, TypeVar,
@@ -32,30 +32,6 @@ from midas.lexer.token import Token
from midas.parser.midas import MidasParser 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]): 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""" """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) self._preamble: Environment = Preamble(self.types)
def set_reporter(self, reporter: FileReporter): def set_reporter(self, reporter: FileReporter):
"""Set the file reporter to use for diagnostics
Args:
reporter (FileReporter): the file reporter
"""
self.reporter = reporter self.reporter = reporter
self.dispatcher.set_reporter(reporter) self.dispatcher.set_reporter(reporter)
def process(self, source: str, path: Optional[str]): 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) reporter: FileReporter = self.reporter.for_file(path)
self.set_reporter(reporter) 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) self.resolve(stmts)
def type_of(self, expr: m.Expr) -> Type: 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) type: Type = expr.accept(self)
return type 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) return self.types.get_type(name)
def get_variable(self, name: str) -> Type: 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: if name in self._predicate_params:
return self._predicate_params[name] return self._predicate_params[name]
predicate: Optional[Predicate] = self.types.lookup_predicate(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) self.types._types[name] = inferrer.infer(type)
def assert_bool(self, expr: m.Expr): 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) type: Type = self.type_of(expr)
if not self.types.is_subtype(type, self._bool): if not self.types.is_subtype(type, self._bool):
self.reporter.error(expr.location, f"Must be a boolean but is {type}") 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) self._predicate_params[param.name.lexeme] = param.type.accept(self)
type: Type = self.type_of(stmt.body) type: Type = self.type_of(stmt.body)
params: list[TypedParamSpec] = [ params: list[ParamSpec] = [self._visit_param_spec(spec) for spec in stmt.params]
self._visit_param_spec(spec) for spec in stmt.params
]
if not self._is_valid_predicate(type): if not self._is_valid_predicate(type):
self.reporter.error( self.reporter.error(
@@ -209,9 +222,7 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type
type = self._bool type = self._bool
for spec in reversed(params): for spec in reversed(params):
type = Function( type = Function(
pos_args=spec.pos, params=spec,
args=spec.mixed,
kw_args=spec.kw,
returns=type, returns=type,
) )
self._predicate_params = {} 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: 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: match body:
case Function(returns=returns): case Function(returns=returns):
return self._is_valid_predicate(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) return self._visit_binary_expr(expr.location, expr.left, expr.right, method)
def _visit_binary_expr( 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: ) -> Type:
left: Type = self.type_of(left_expr) left: Type = self.type_of(left_expr)
right: Type = self.type_of(right_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: def visit_function_type(self, type: m.FunctionType) -> Type:
params: TypedParamSpec = self._visit_param_spec(type.params)
return Function( return Function(
pos_args=params.pos, params=self._visit_param_spec(type.params),
args=params.mixed,
kw_args=params.kw,
returns=type.returns.accept(self), 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_pos: int = len(spec.pos)
n_mixed: int = len(spec.mixed) n_mixed: int = len(spec.mixed)
def process_arg(arg: m.FunctionType.Argument, i: int) -> Function.Argument: def process_param(
return Function.Argument( param: m.FunctionType.Parameter, i: int
) -> Function.Parameter:
return Function.Parameter(
pos=i, pos=i,
name=arg.name.lexeme if arg.name is not None else str(i), name=param.name.lexeme if param.name is not None else str(i),
type=arg.type.accept(self), type=param.type.accept(self),
required=arg.required, required=param.required,
) )
return TypedParamSpec( return ParamSpec(
pos=[process_arg(arg, i) for i, arg in enumerate(spec.pos)], pos=[process_param(param, i) for i, param in enumerate(spec.pos)],
mixed=[process_arg(arg, i + n_pos) for i, arg in enumerate(spec.mixed)], mixed=[
kw=[process_arg(arg, i + n_pos + n_mixed) for i, arg in enumerate(spec.kw)], process_param(param, i + n_pos) for i, param in enumerate(spec.mixed)
],
kw=[
process_param(param, i + n_pos + n_mixed)
for i, param in enumerate(spec.kw)
],
) )
def visit_frame_type(self, type: m.FrameType) -> Type: def visit_frame_type(self, type: m.FrameType) -> Type:

15
midas/checker/preamble.py
View File

@@ -7,6 +7,7 @@ from midas.checker.types import (
Function, Function,
GenericType, GenericType,
OverloadedFunction, OverloadedFunction,
ParamSpec,
TopType, TopType,
Type, Type,
TypeVar, TypeVar,
@@ -22,6 +23,8 @@ class Param:
class Preamble(Environment): class Preamble(Environment):
"""The initial environment containing some of Python's builtin functions"""
def __init__(self, types: TypesRegistry) -> None: def __init__(self, types: TypesRegistry) -> None:
super().__init__() super().__init__()
self._types: TypesRegistry = types self._types: TypesRegistry = types
@@ -132,9 +135,9 @@ class Preamble(Environment):
returns: Type = UnitType(), returns: Type = UnitType(),
type_vars: list[TypeVar] = [], type_vars: list[TypeVar] = [],
) -> Type: ) -> Type:
def map_args(params: list[Param], offset: int) -> list[Function.Argument]: def map_params(params: list[Param], offset: int) -> list[Function.Parameter]:
return [ return [
Function.Argument( Function.Parameter(
pos=i + offset, pos=i + offset,
name=param.name, name=param.name,
type=param.type, type=param.type,
@@ -144,9 +147,11 @@ class Preamble(Environment):
] ]
function = Function( function = Function(
pos_args=map_args(pos, 0), params=ParamSpec(
args=map_args(mixed, len(pos)), pos=map_params(pos, 0),
kw_args=map_args(kw, len(pos) + len(mixed)), mixed=map_params(mixed, len(pos)),
kw=map_params(kw, len(pos) + len(mixed)),
),
returns=returns, returns=returns,
) )
if len(type_vars) != 0: if len(type_vars) != 0:

245
midas/checker/python.py
View File

@@ -1,6 +1,5 @@
import ast import ast
import logging import logging
from dataclasses import dataclass
from typing import Any, Optional from typing import Any, Optional
import midas.ast.python as p import midas.ast.python as p
@@ -31,6 +30,7 @@ from midas.checker.types import (
FrameGroupBy, FrameGroupBy,
Function, Function,
GenericType, GenericType,
ParamSpec,
TopType, TopType,
TupleType, TupleType,
Type, Type,
@@ -55,19 +55,6 @@ class UndefinedMethodException(Exception):
pass 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( class PythonTyper(
p.Stmt.Visitor[None], p.Stmt.Visitor[None],
p.Expr.Visitor[Type], p.Expr.Visitor[Type],
@@ -96,10 +83,24 @@ class PythonTyper(
self.assertions: AssertionCollector = AssertionCollector() self.assertions: AssertionCollector = AssertionCollector()
def set_reporter(self, reporter: FileReporter): def set_reporter(self, reporter: FileReporter):
"""Set the file reporter to use for diagnostics
Args:
reporter (FileReporter): the file reporter
"""
self.reporter = reporter self.reporter = reporter
self.dispatcher.set_reporter(self.reporter) self.dispatcher.set_reporter(self.reporter)
def process(self, source: str, path: Optional[str]) -> TypedAST: 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) reporter: FileReporter = self.reporter.for_file(path)
self.set_reporter(reporter) self.set_reporter(reporter)
@@ -124,7 +125,7 @@ class PythonTyper(
) )
def judge(self, expr: p.Expr, type: Type): def judge(self, expr: p.Expr, type: Type):
"""Record a typing judgement """Record a typing judgement for the given expression
Args: Args:
expr (p.Expr): the judged expression expr (p.Expr): the judged expression
@@ -133,7 +134,7 @@ class PythonTyper(
self.judgements.append((expr, type)) self.judgements.append((expr, type))
def compute_type(self, expr: p.Expr) -> Type: def compute_type(self, expr: p.Expr) -> Type:
"""Evaluate the type of an expression """Evaluate the type of the given expression
Args: Args:
expr (p.Expr): the expression to type expr (p.Expr): the expression to type
@@ -144,7 +145,7 @@ class PythonTyper(
return expr.accept(self) return expr.accept(self)
def type_of(self, expr: p.Expr) -> Type: 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: Args:
expr (p.Expr): the expression to evaluate expr (p.Expr): the expression to evaluate
@@ -157,9 +158,22 @@ class PythonTyper(
return type return type
def resolve_type_expr(self, expr: p.MidasType) -> 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) return expr.accept(self)
def process_stmt(self, stmt: p.Stmt) -> None: 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) stmt.accept(self)
def process_block(self, block: list[p.Stmt], env: Environment) -> bool: def process_block(self, block: list[p.Stmt], env: Environment) -> bool:
@@ -223,6 +237,24 @@ class PythonTyper(
positional: list[TypedExpr], positional: list[TypedExpr],
keywords: dict[str, TypedExpr], keywords: dict[str, TypedExpr],
) -> Type: ) -> 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]) unfolded: Type = unfold_type(obj[1])
match unfolded: match unfolded:
case DataFrameType(): case DataFrameType():
@@ -282,6 +314,15 @@ class PythonTyper(
return result.result return result.result
def is_subtype(self, type1: Type, type2: Type) -> bool: 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) return self.types.is_subtype(type1, type2)
def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None: def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> None:
@@ -289,61 +330,64 @@ class PythonTyper(
def visit_function(self, stmt: p.Function) -> None: def visit_function(self, stmt: p.Function) -> None:
env: Environment = Environment(self.env) env: Environment = Environment(self.env)
pos_args: list[Function.Argument] = [] pos: list[Function.Parameter] = []
args: list[Function.Argument] = [] mixed: list[Function.Parameter] = []
kw_args: list[Function.Argument] = [] kw: list[Function.Parameter] = []
def eval_arg_type(arg: p.Function.Argument) -> Type: def eval_param_type(param: p.Function.Parameter) -> Type:
if arg.type is not None: if param.type is not None:
return self.resolve_type_expr(arg.type) return self.resolve_type_expr(param.type)
if arg.default is not None: if param.default is not None:
return self.type_of(arg.default) return self.type_of(param.default)
return UnknownType() return UnknownType()
pos: int = 0 position: int = 0
for arg in stmt.posonlyargs: for param in stmt.params.pos:
pos_args.append( pos.append(
Function.Argument( Function.Parameter(
pos=pos, pos=position,
name=arg.name, name=param.name,
type=eval_arg_type(arg), type=eval_param_type(param),
required=arg.default is None, required=param.default is None,
) )
) )
pos += 1 position += 1
for arg in stmt.args: for param in stmt.params.mixed:
args.append( mixed.append(
Function.Argument( Function.Parameter(
pos=pos, pos=position,
name=arg.name, name=param.name,
type=eval_arg_type(arg), type=eval_param_type(param),
required=arg.default is None, required=param.default is None,
) )
) )
pos += 1 position += 1
for arg in stmt.kwonlyargs: for param in stmt.params.kw:
kw_args.append( kw.append(
Function.Argument( Function.Parameter(
pos=pos, # not relevant pos=position, # not relevant
name=arg.name, name=param.name,
type=eval_arg_type(arg), type=eval_param_type(param),
required=arg.default is None, required=param.default is None,
) )
) )
pos += 1 position += 1
all_args: list[Function.Argument] = pos_args + args + kw_args param_spec: ParamSpec = ParamSpec(
for arg in all_args: pos=pos,
env.define(arg.name, arg.type) mixed=mixed,
kw=kw,
)
all_params: list[Function.Parameter] = pos + mixed + kw
for param in all_params:
env.define(param.name, param.type)
returns_hint: Optional[Type] = None returns_hint: Optional[Type] = None
if stmt.returns is not None: if stmt.returns is not None:
returns_hint = self.resolve_type_expr(stmt.returns) returns_hint = self.resolve_type_expr(stmt.returns)
# Early define to handle simple fully-typed recursion # Early define to handle simple fully-typed recursion
inside_function: Function = Function( inside_function: Function = Function(
pos_args=pos_args, params=param_spec,
args=args,
kw_args=kw_args,
returns=returns_hint, returns=returns_hint,
) )
self.env.define(stmt.name, inside_function) self.env.define(stmt.name, inside_function)
@@ -375,13 +419,11 @@ class PythonTyper(
# TODO: handle *args and **kwargs sinks # TODO: handle *args and **kwargs sinks
function: Type = Function( function: Type = Function(
pos_args=pos_args, params=param_spec,
args=args,
kw_args=kw_args,
returns=returns, returns=returns,
) )
generic_params: list[TypeVar] = [] 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: for type in all_types:
if isinstance(type, TypeVar): if isinstance(type, TypeVar):
if type not in generic_params: if type not in generic_params:
@@ -406,6 +448,15 @@ class PythonTyper(
self._assign(stmt.location, target, value_type) self._assign(stmt.location, target, value_type)
def _assign(self, location: Location, target: p.Expr, value_type: 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: match target:
case p.VariableExpr(): case p.VariableExpr():
self._assign_var(location, target, value_type) 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): 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 name: str = target.name
var_type: Optional[Type] = self.look_up_variable(name, target) var_type: Optional[Type] = self.look_up_variable(name, target)
@@ -445,6 +503,13 @@ class PythonTyper(
def _assign_attr( def _assign_attr(
self, location: Location, object: p.Expr, name: str, value_type: Type 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) object_type: Type = self.type_of(object)
member: Optional[Type] = self.types.lookup_member(object_type, name) member: Optional[Type] = self.types.lookup_member(object_type, name)
if member is None: if member is None:
@@ -464,6 +529,13 @@ class PythonTyper(
index: p.Expr, index: p.Expr,
value_type: Type, 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) var_type: Type = self.type_of(var)
unfolded_type: Type = unfold_type(var_type) unfolded_type: Type = unfold_type(var_type)
# TODO: what happens if type is an alias of a dataframe 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]: 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__ # TODO: lookup __iter__
getitem: Optional[Type] = self.types.lookup_member(type, "__getitem__") getitem: Optional[Type] = self.types.lookup_member(type, "__getitem__")
if getitem is None: if getitem is None:
@@ -904,6 +985,16 @@ class PythonTyper(
return result.result return result.result
def define_typevar(self, call: p.CallExpr) -> Optional[TypeVar]: 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: def is_kw_true(name: str) -> bool:
match call.keywords.get(name): match call.keywords.get(name):
case p.LiteralExpr(value=True): case p.LiteralExpr(value=True):
@@ -946,6 +1037,19 @@ class PythonTyper(
return None return None
def _parse_type_from_expr(self, expr: p.Expr) -> p.MidasType: 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 location: Location = expr.location
parser = PythonParser() parser = PythonParser()
match expr: match expr:
@@ -958,6 +1062,16 @@ class PythonTyper(
raise NotImplementedError raise NotImplementedError
def _get_literal(self, expr: p.Expr) -> tuple[bool, Any]: 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: match expr:
case p.LiteralExpr(value=value): case p.LiteralExpr(value=value):
return True, value return True, value
@@ -1014,6 +1128,17 @@ class PythonTyper(
def _evaluate_cast_statically( def _evaluate_cast_statically(
self, expr: p.CastExpr, subject_type: Type, target_type: Type, lit_value: Any self, expr: p.CastExpr, subject_type: Type, target_type: Type, lit_value: Any
) -> bool: ) -> 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: match target_type:
case TopType(): case TopType():
return True return True

224
midas/checker/registry.py
View File

@@ -29,11 +29,15 @@ from midas.checker.types import (
@dataclass @dataclass
class Member: class Member:
"""A member of a type (property or method)"""
kind: MemberKind kind: MemberKind
type: Type type: Type
class TypesRegistry: class TypesRegistry:
"""A registry of types, type members and predicates"""
def __init__(self) -> None: def __init__(self) -> None:
self.logger: logging.Logger = logging.getLogger("TypesRegistry") self.logger: logging.Logger = logging.getLogger("TypesRegistry")
self._types: dict[str, Type] = {} self._types: dict[str, Type] = {}
@@ -81,6 +85,25 @@ class TypesRegistry:
member_type: Type, member_type: Type,
kind: MemberKind, 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, {}) members: dict[str, Member] = self._members.setdefault(type_name, {})
if member_name in members: if member_name in members:
current: Member = members[member_name] current: Member = members[member_name]
@@ -109,11 +132,29 @@ class TypesRegistry:
members[member_name] = Member(kind=kind, type=member_type) members[member_name] = Member(kind=kind, type=member_type)
def define_predicate(self, name: str, predicate: Predicate): 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: if name in self._predicates:
raise ValueError(f"Predicate {name} already defined") raise ValueError(f"Predicate {name} already defined")
self._predicates[name] = predicate self._predicates[name] = predicate
def is_builtin_subtype(self, name1: str, name2: str) -> bool: 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()) subtypes: set[str] = BUILTIN_SUBTYPES.get(name2, set())
if name1 in subtypes: if name1 in subtypes:
return True return True
@@ -218,6 +259,15 @@ class TypesRegistry:
return False return False
def are_equivalent(self, type1: Type, type2: Type) -> bool: 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) return self.is_subtype(type1, type2) and self.is_subtype(type2, type1)
# TODO: verify the logic in here # TODO: verify the logic in here
@@ -234,98 +284,118 @@ class TypesRegistry:
if not self.is_subtype(func1.returns, func2.returns): if not self.is_subtype(func1.returns, func2.returns):
return False return False
pos1: list[Function.Argument] = func1.pos_args pos1: list[Function.Parameter] = func1.params.pos
mixed1: list[Function.Argument] = func1.args mixed1: list[Function.Parameter] = func1.params.mixed
kw1: dict[str, Function.Argument] = {a.name: a for a in func1.kw_args} kw1: dict[str, Function.Parameter] = {
pos2: list[Function.Argument] = func2.pos_args param.name: param for param in func1.params.kw
mixed2: list[Function.Argument] = func2.args }
kw2: dict[str, Function.Argument] = {a.name: a for a in func2.kw_args} pos2: list[Function.Parameter] = func2.params.pos
mixed2: list[Function.Parameter] = func2.params.mixed
kw2: dict[str, Function.Parameter] = {
param.name: param for param in func2.params.kw
}
mixed_by_pos: dict[int, Function.Argument] = {arg.pos: arg for arg in mixed2} mixed_by_pos: dict[int, Function.Parameter] = {
mixed_by_name: dict[str, Function.Argument] = {arg.name: arg for arg in mixed2} param.pos: param for param in mixed2
}
mixed_by_name: dict[str, Function.Parameter] = {
param.name: param for param in mixed2
}
def is_arg_subtype(sub: Function.Argument, sup: Function.Argument) -> bool: def is_arg_subtype(sub: Function.Parameter, sup: Function.Parameter) -> bool:
if not self.is_subtype(sub.type, sup.type): if not self.is_subtype(sub.type, sup.type):
return False return False
if not sup.required and sub.required: if not sup.required and sub.required:
return False return False
return True return True
for arg1 in pos1: for param1 in pos1:
arg2: Function.Argument param2: Function.Parameter
if arg1.pos < len(pos2): if param1.pos < len(pos2):
arg2 = pos2[arg1.pos] param2 = pos2[param1.pos]
elif arg1.pos in mixed_by_pos: elif param1.pos in mixed_by_pos:
arg2 = mixed_by_pos[arg1.pos] param2 = mixed_by_pos[param1.pos]
elif not arg1.required: elif not param1.required:
continue continue
else: else:
return False return False
if not is_arg_subtype(arg2, arg1): if not is_arg_subtype(param2, param1):
return False return False
for name, arg1 in kw1.items(): for name, param1 in kw1.items():
arg2: Function.Argument param2: Function.Parameter
if name in kw2: if name in kw2:
arg2 = kw2[name] param2 = kw2[name]
elif name in mixed_by_name: elif name in mixed_by_name:
arg2 = mixed_by_name[name] param2 = mixed_by_name[name]
elif not arg1.required: elif not param1.required:
continue continue
else: else:
return False return False
if not is_arg_subtype(arg2, arg1): if not is_arg_subtype(param2, param1):
return False return False
for arg1 in mixed1: for param1 in mixed1:
pos_arg2: Optional[Function.Argument] = None pos_param2: Optional[Function.Parameter] = None
kw_arg2: Optional[Function.Argument] = None kw_param2: Optional[Function.Parameter] = None
if arg1.name in kw2: if param1.name in kw2:
kw_arg2 = kw2[arg1.name] kw_param2 = kw2[param1.name]
elif arg1.name in mixed_by_name: elif param1.name in mixed_by_name:
kw_arg2 = mixed_by_name[arg1.name] kw_param2 = mixed_by_name[param1.name]
if arg1.pos < len(pos2): if param1.pos < len(pos2):
pos_arg2 = pos2[arg1.pos] pos_param2 = pos2[param1.pos]
elif arg1.pos in mixed_by_pos: elif param1.pos in mixed_by_pos:
pos_arg2 = mixed_by_pos[arg1.pos] pos_param2 = mixed_by_pos[param1.pos]
# No match in func2 and arg is required # 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 return False
# Matching keyword argument # 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 return False
# Matching positional argument # 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 return False
mixed_positions: set[int] = {a.pos for a in mixed1} mixed_positions: set[int] = {param.pos for param in mixed1}
mixed_names: set[str] = {a.name for a in mixed1} mixed_names: set[str] = {param.name for param in mixed1}
for arg2 in pos2: for param2 in pos2:
if not arg2.required: if not param2.required:
continue 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 return False
for name, arg2 in kw2.items(): for name, param2 in kw2.items():
if not arg2.required: if not param2.required:
continue continue
if name not in kw1 and name not in mixed_names: if name not in kw1 and name not in mixed_names:
return False return False
for arg2 in mixed2: for param2 in mixed2:
if arg2.required: if param2.required:
continue continue
pos_match: bool = arg2.pos < len(pos1) or arg2.pos in mixed_positions pos_match: bool = param2.pos < len(pos1) or param2.pos in mixed_positions
kw_match: bool = arg2.name in kw1 or arg2.name in mixed_names kw_match: bool = param2.name in kw1 or param2.name in mixed_names
if not pos_match or not kw_match: if not pos_match or not kw_match:
return False return False
return True return True
def apply_generic(self, type: Type, args: list[Type]) -> Type: 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: match type:
case DerivedType(name=name, type=base): case DerivedType(name=name, type=base):
return DerivedType(name=name, type=self.apply_generic(base, args)) return DerivedType(name=name, type=self.apply_generic(base, args))
@@ -391,6 +461,19 @@ class TypesRegistry:
return [types[i] for i in keep] return [types[i] for i in keep]
def lookup_member(self, type: Type, member_name: str) -> Optional[Type]: 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: match type:
case BaseType(name=name): case BaseType(name=name):
if name in self._members: if name in self._members:
@@ -451,18 +534,54 @@ class TypesRegistry:
return None return None
def lookup_predicate(self, name: str) -> Optional[Predicate]: 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) return self._predicates.get(name)
def _by_name_or_type(self, name_or_type: str | Type) -> Type: 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): if isinstance(name_or_type, str):
return self.get_type(name_or_type) return self.get_type(name_or_type)
return name_or_type return name_or_type
def list_of(self, item_type: str | Type) -> 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") list_ = self.get_type("list")
return self.apply_generic(list_, [self._by_name_or_type(item_type)]) return self.apply_generic(list_, [self._by_name_or_type(item_type)])
def tuple_of(self, *item_types: str | Type) -> 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") tuple_ = self.get_type("tuple")
return self.apply_generic( return self.apply_generic(
tuple_, tuple_,
@@ -470,6 +589,15 @@ class TypesRegistry:
) )
def dict_of(self, key_type: str | Type, value_type: str | Type) -> Type: 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") dict_ = self.get_type("dict")
return self.apply_generic( return self.apply_generic(
dict_, dict_,

59
midas/checker/reporter.py
View File

@@ -7,6 +7,8 @@ from midas.checker.diagnostic import Diagnostic, DiagnosticType
class Reporter: class Reporter:
"""Helper class to store diagnostics"""
def __init__(self): def __init__(self):
self.diagnostics: list[Diagnostic] = [] self.diagnostics: list[Diagnostic] = []
@@ -17,6 +19,14 @@ class Reporter:
location: Location, location: Location,
message: str, 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( self.diagnostics.append(
Diagnostic( Diagnostic(
file_path=path, file_path=path,
@@ -27,21 +37,52 @@ class Reporter:
) )
def for_file(self, path: Optional[str]) -> FileReporter: 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) return FileReporter(self, path)
class FileReporter: class FileReporter:
"""Helper class to manage diagnostics for a file"""
def __init__(self, base_reporter: Reporter, path: Optional[str]) -> None: def __init__(self, base_reporter: Reporter, path: Optional[str]) -> None:
self.base_reporter: Reporter = base_reporter self.base_reporter: Reporter = base_reporter
self.path: Optional[str] = path self.path: Optional[str] = path
def for_file(self, path: Optional[str]) -> FileReporter: 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) return FileReporter(self.base_reporter, path)
def report(self, type: DiagnosticType, location: Location, message: str): 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) self.base_reporter.report(self.path, type, location, message)
def error(self, location: Location, message: str): 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( self.report(
type=DiagnosticType.ERROR, type=DiagnosticType.ERROR,
location=location, location=location,
@@ -49,6 +90,12 @@ class FileReporter:
) )
def warning(self, location: Location, message: str): 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( self.report(
type=DiagnosticType.WARNING, type=DiagnosticType.WARNING,
location=location, location=location,
@@ -56,6 +103,12 @@ class FileReporter:
) )
def info(self, location: Location, message: str): 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( self.report(
type=DiagnosticType.INFO, type=DiagnosticType.INFO,
location=location, location=location,
@@ -63,6 +116,12 @@ class FileReporter:
) )
def debug(self, location: Location, message: str): 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( self.report(
type=DiagnosticType.DEBUG, type=DiagnosticType.DEBUG,
location=location, location=location,

10
midas/checker/resolver.py
View File

@@ -78,6 +78,14 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
return return
def is_defined(self, name: str) -> bool: 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: for scope in self.scopes:
if name in scope: if name in scope:
return True return True
@@ -93,7 +101,7 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
function (p.Function): the function to resolve function (p.Function): the function to resolve
""" """
self.begin_scope() self.begin_scope()
for param in function.all_args: for param in function.params.all:
self.declare(param.name) self.declare(param.name)
self.define(param.name) self.define(param.name)
self.resolve(*function.body) self.resolve(*function.body)

155
midas/checker/types.py
View File

@@ -10,12 +10,16 @@ from midas.ast.printer import MidasPrinter
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class TopType: class TopType:
"""The top type (`Any`)"""
def __str__(self) -> str: def __str__(self) -> str:
return "Any" return "Any"
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class BaseType: class BaseType:
"""A base / builtin type"""
name: str name: str
def __str__(self) -> str: def __str__(self) -> str:
@@ -24,6 +28,8 @@ class BaseType:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class DerivedType: class DerivedType:
"""A derived type, i.e. a named subtype of another type"""
name: str name: str
type: Type type: Type
@@ -33,40 +39,32 @@ class DerivedType:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class UnknownType: class UnknownType:
"""An unknown type"""
def __str__(self) -> str: def __str__(self) -> str:
return "<Unknown>" return "<Unknown>"
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class UnitType: class UnitType:
"""The unit type (`None`)"""
def __str__(self) -> str: def __str__(self) -> str:
return "None" return "None"
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class Function: class Function:
pos_args: list[Argument] = field(default_factory=list) """A function type"""
args: list[Argument] = field(default_factory=list)
kw_args: list[Argument] = field(default_factory=list) params: ParamSpec
returns: Type returns: Type
def __str__(self) -> str: def __str__(self) -> str:
args: list[str] = [] return f"{self.params} -> {self.returns}"
if len(self.pos_args) != 0:
args += list(map(str, self.pos_args))
args.append("/")
if len(self.args) != 0:
args += list(map(str, self.args))
if len(self.kw_args) != 0:
args.append("*")
args += list(map(str, self.kw_args))
return f"({', '.join(args)}) -> {self.returns}"
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class Argument: class Parameter:
pos: int pos: int
name: str name: str
type: Type type: Type
@@ -77,8 +75,34 @@ class Function:
return f"{self.name}: {self.type}{opt}" 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) @dataclass(frozen=True, kw_only=True)
class OverloadedFunction: class OverloadedFunction:
"""A list of method overloads"""
overloads: list[Type] overloads: list[Type]
def __str__(self) -> str: def __str__(self) -> str:
@@ -87,6 +111,8 @@ class OverloadedFunction:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class ComplexType: class ComplexType:
"""A type with inline members"""
members: dict[str, Type] members: dict[str, Type]
def __str__(self) -> str: def __str__(self) -> str:
@@ -96,6 +122,8 @@ class ComplexType:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class ExtensionType: class ExtensionType:
"""An extension of a type, adding members through a `ComplexType`"""
base: Type base: Type
extension: ComplexType extension: ComplexType
@@ -104,6 +132,8 @@ class ExtensionType:
class Variance(StrEnum): class Variance(StrEnum):
"""The variance of a :class:`TypeVar`"""
INVARIANT = "INVARIANT" INVARIANT = "INVARIANT"
COVARIANT = "COVARIANT" COVARIANT = "COVARIANT"
CONTRAVARIANT = "CONTRAVARIANT" CONTRAVARIANT = "CONTRAVARIANT"
@@ -111,6 +141,8 @@ class Variance(StrEnum):
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class TypeVar: class TypeVar:
"""A type variable, often used as type parameters for a generic type"""
name: str name: str
bound: Optional[Type] bound: Optional[Type]
variance: Variance = Variance.INVARIANT variance: Variance = Variance.INVARIANT
@@ -128,6 +160,8 @@ class TypeVar:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class GenericType: class GenericType:
"""A generic type, with type parameters and a generic body type"""
name: str name: str
params: list[TypeVar] params: list[TypeVar]
body: Type body: Type
@@ -138,6 +172,8 @@ class GenericType:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class AppliedType: class AppliedType:
"""An instance of a :class:`GenericType`, with concrete type arguments substituted in its body"""
name: str name: str
args: list[Type] args: list[Type]
body: Type body: Type
@@ -148,6 +184,8 @@ class AppliedType:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class ConstraintType: class ConstraintType:
"""A type with a constraint expression"""
type: Type type: Type
constraint: m.Expr constraint: m.Expr
@@ -158,6 +196,8 @@ class ConstraintType:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class TupleType: class TupleType:
"""A tuple type, containing any number of ordered item types"""
items: tuple[Type, ...] items: tuple[Type, ...]
def __str__(self) -> str: def __str__(self) -> str:
@@ -166,6 +206,8 @@ class TupleType:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class ColumnType: class ColumnType:
"""A column type containing items of a given unique type"""
type: Type type: Type
def __str__(self) -> str: def __str__(self) -> str:
@@ -174,6 +216,8 @@ class ColumnType:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class DataFrameType: class DataFrameType:
"""A data-frame type, containing named columns of specific :class:`ColumnType`"""
columns: list[Column] columns: list[Column]
def __str__(self) -> str: def __str__(self) -> str:
@@ -189,6 +233,8 @@ class DataFrameType:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class FrameGroupBy: class FrameGroupBy:
"""A frame group-by object"""
frame: DataFrameType frame: DataFrameType
def __str__(self) -> str: def __str__(self) -> str:
@@ -197,6 +243,8 @@ class FrameGroupBy:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class ColumnGroupBy: class ColumnGroupBy:
"""A column group-by object"""
column: ColumnType column: ColumnType
def __str__(self) -> str: def __str__(self) -> str:
@@ -204,12 +252,32 @@ class ColumnGroupBy:
def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
def sub_argument(arg: Function.Argument): """Substitute type variables in the given type
return Function.Argument(
pos=arg.pos, This function is called recursively on inner type structures
name=arg.name,
type=substitute_typevars(arg.type, substitutions), Args:
required=arg.required, 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): def sub_column(col: DataFrameType.Column):
@@ -235,15 +303,11 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
) )
case Function( case Function(
pos_args=pos_args, params=params,
args=args,
kw_args=kw_args,
returns=returns, returns=returns,
): ):
return Function( return Function(
pos_args=list(map(sub_argument, pos_args)), params=sub_param_spec(params),
args=list(map(sub_argument, args)),
kw_args=list(map(sub_argument, kw_args)),
returns=substitute_typevars(returns, substitutions), 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: 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: match type:
case DerivedType(type=ref_type): case DerivedType(type=ref_type):
return unfold_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 to_annotation(type: Type) -> str:
def _args_annotation(func: Function) -> str: """Convert the given type to a Python annotation string
if len(func.kw_args) != 0:
Args:
type (Type): the type to convert
Returns:
str: the annotation string
"""
def _params_annotation(spec: ParamSpec) -> str:
if len(spec.kw) != 0:
return "..." return "..."
args: str = ", ".join( params: str = ", ".join(
to_annotation(arg.type) for arg in func.pos_args + func.args to_annotation(param.type) for param in spec.pos + spec.mixed
) )
return f"[{args}]" return f"[{params}]"
match type: match type:
case TopType(): case TopType():
@@ -376,8 +457,8 @@ def to_annotation(type: Type) -> str:
case UnitType(): case UnitType():
return "None" return "None"
case Function(returns=returns): case Function(params=params, returns=returns):
params_annot: str = _args_annotation(type) params_annot: str = _params_annotation(params)
return f"Callable[{params_annot}, {to_annotation(returns)}]" return f"Callable[{params_annot}, {to_annotation(returns)}]"
case OverloadedFunction(): case OverloadedFunction():
@@ -419,6 +500,8 @@ def to_annotation(type: Type) -> str:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class Predicate: class Predicate:
"""A predicate"""
type: Type type: Type
body: m.Expr body: m.Expr
alias: bool alias: bool

105
midas/checker/unifier.py
View File

@@ -8,6 +8,7 @@ from midas.checker.types import (
DataFrameType, DataFrameType,
Function, Function,
GenericType, GenericType,
ParamSpec,
TopType, TopType,
Type, Type,
TypeVar, TypeVar,
@@ -18,6 +19,14 @@ class UnificationError(Exception): ...
class Unifier: 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: def __init__(self, types: TypesRegistry) -> None:
self.types: TypesRegistry = types self.types: TypesRegistry = types
self.logger: logging.Logger = logging.getLogger("Unifier") self.logger: logging.Logger = logging.getLogger("Unifier")
@@ -28,9 +37,20 @@ class Unifier:
positional: list[Type], positional: list[Type],
keywords: dict[str, Type], keywords: dict[str, Type],
) -> Optional[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( concrete_func: Function = Function(
pos_args=[ params=ParamSpec(
Function.Argument( pos=[
Function.Parameter(
pos=i, pos=i,
name=str(i), name=str(i),
type=arg, type=arg,
@@ -38,9 +58,8 @@ class Unifier:
) )
for i, arg in enumerate(positional) for i, arg in enumerate(positional)
], ],
args=[], kw=[
kw_args=[ Function.Parameter(
Function.Argument(
pos=len(positional) + i, pos=len(positional) + i,
name=name, name=name,
type=arg, type=arg,
@@ -48,6 +67,7 @@ class Unifier:
) )
for i, (name, arg) in enumerate(keywords.items()) for i, (name, arg) in enumerate(keywords.items())
], ],
),
returns=TopType(), # TODO: use expected type returns=TopType(), # TODO: use expected type
) )
return self.unify_generic(type, concrete_func, match_return=False) return self.unify_generic(type, concrete_func, match_return=False)
@@ -58,6 +78,18 @@ class Unifier:
concrete: Type, concrete: Type,
match_return: bool = True, match_return: bool = True,
) -> Optional[Type]: ) -> 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] substitutions: dict[str, Type]
try: try:
substitutions = self.match(template.body, concrete, match_return) substitutions = self.match(template.body, concrete, match_return)
@@ -79,6 +111,22 @@ class Unifier:
concrete: Type, concrete: Type,
match_return: bool = True, match_return: bool = True,
) -> dict[str, Type]: ) -> 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 # TODO: if concrete is Generic, record bound TypeVar. Then when merging
# substitutions, check that the constraint is respected # substitutions, check that the constraint is respected
match (template, concrete): match (template, concrete):
@@ -125,7 +173,7 @@ class Unifier:
return self.match(template_column, concrete_column) return self.match(template_column, concrete_column)
case (Function(), Function()): case (Function(), Function()):
mapped: list[tuple[Function.Argument, Function.Argument]] = ( mapped: list[tuple[Function.Parameter, Function.Parameter]] = (
self.map_params(template, concrete) self.map_params(template, concrete)
) )
substitutions: dict[str, Type] = {} substitutions: dict[str, Type] = {}
@@ -148,6 +196,18 @@ class Unifier:
return {} return {}
def merge(self, subs1: dict[str, Type], subs2: dict[str, Type]) -> dict[str, Type]: 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() merged: dict[str, Type] = subs1.copy()
for k, v in subs2.items(): for k, v in subs2.items():
@@ -161,19 +221,32 @@ class Unifier:
def map_params( def map_params(
self, func1: Function, func2: Function self, func1: Function, func2: Function
) -> list[tuple[Function.Argument, Function.Argument]]: ) -> list[tuple[Function.Parameter, Function.Parameter]]:
pos1: list[Function.Argument] = func1.pos_args """Map parameters of two functions
mixed1: list[Function.Argument] = func1.args
kw1: list[Function.Argument] = func1.kw_args
pos2: list[Function.Argument] = func2.pos_args Args:
mixed2: list[Function.Argument] = func2.args func1 (Function): the first function
kw2: list[Function.Argument] = func2.kw_args 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} pos2: list[Function.Parameter] = func2.params.pos
by_name2: dict[str, Function.Argument] = {arg.name: arg for arg in mixed2 + kw2} 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: for arg1 in pos1:
if (arg2 := by_pos2.get(arg1.pos)) is not None: if (arg2 := by_pos2.get(arg1.pos)) is not None:

62
midas/checker/variance.py
View File

@@ -16,14 +16,27 @@ Polarity = Literal[-1, 0, 1]
class Tracker: class Tracker:
"""Helper class to track the polarity of type parameter references and computer their variance"""
def __init__(self, vars: list[TypeVar]) -> None: def __init__(self, vars: list[TypeVar]) -> None:
self.vars: list[TypeVar] = vars self.vars: list[TypeVar] = vars
self.refs: dict[str, set[Polarity]] = {var.name: set() for var in self.vars} self.refs: dict[str, set[Polarity]] = {var.name: set() for var in self.vars}
def record(self, var: TypeVar, polarity: Polarity): 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) self.refs[var.name].add(polarity)
def get_updated_vars(self) -> list[TypeVar]: 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 [ return [
TypeVar( TypeVar(
name=var.name, bound=var.bound, variance=self.get_variance(var.name) 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: 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] refs: set[Polarity] = self.refs[name]
if refs == {-1}: if refs == {-1}:
return Variance.CONTRAVARIANT return Variance.CONTRAVARIANT
@@ -46,11 +71,22 @@ class Tracker:
class VarianceInferrer: class VarianceInferrer:
"""Helper class to compute type parameter variance"""
def __init__(self, types: TypesRegistry) -> None: def __init__(self, types: TypesRegistry) -> None:
self.types: TypesRegistry = types self.types: TypesRegistry = types
self.tracker: Tracker = Tracker([]) self.tracker: Tracker = Tracker([])
def infer(self, type: GenericType) -> GenericType: 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.tracker = Tracker(type.params)
self.walk(type.body, 1, type.name) self.walk(type.body, 1, type.name)
@@ -71,20 +107,36 @@ class VarianceInferrer:
base_name: str, base_name: str,
path: Optional[list[str]] = None, 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: if path is None:
path = [] path = []
match type: match type:
# Arguments are negative positions -> flip polarity # Arguments are negative positions -> flip polarity
# Return is positive position -> keep polarity # Return is positive position -> keep polarity
case Function(pos_args=pos_args, args=mixed_args, kw_args=kw_args): case Function(params=spec):
all_args: list[Function.Argument] = pos_args + mixed_args + kw_args all_params: list[Function.Parameter] = spec.pos + spec.mixed + spec.kw
for arg in all_args: for param in all_params:
self.walk( self.walk(
arg.type, param.type,
-polarity, -polarity,
base_name, base_name,
path + [f"arg:'{arg.name}'"], path + [f"param:'{param.name}'"],
) )
self.walk(type.returns, polarity, base_name, path + ["return"]) self.walk(type.returns, polarity, base_name, path + ["return"])

15
midas/cli/highlighter.py
View File

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

2
midas/cli/hl_python.css
View File

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

37
midas/generator/constraints.py
View File

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

33
midas/generator/generator.py
View File

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

32
midas/generator/stubs.py
View File

@@ -17,6 +17,7 @@ from midas.checker.types import (
Function, Function,
GenericType, GenericType,
OverloadedFunction, OverloadedFunction,
ParamSpec,
TopType, TopType,
TupleType, TupleType,
Type, Type,
@@ -328,7 +329,7 @@ class StubsGenerator:
return [ return [
ast.FunctionDef( ast.FunctionDef(
name=name, 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), returns=self.dump_type(method.returns),
body=[ast.Expr(value=Empty)], body=[ast.Expr(value=Empty)],
decorator_list=[ast.Name(id="overload")] if overloaded else [], 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] = [ pos: list[ast.arg] = [
ast.arg(arg=f"_{arg.pos}", annotation=self.dump_type(arg.type)) ast.arg(
for arg in func.pos_args arg=f"_{param.pos}",
annotation=self.dump_type(param.type),
)
for param in params.pos
] ]
mixed: list[ast.arg] = [ mixed: list[ast.arg] = [
ast.arg(arg=arg.name, annotation=self.dump_type(arg.type)) ast.arg(
for arg in func.args arg=param.name,
annotation=self.dump_type(param.type),
)
for param in params.mixed
] ]
kw: list[ast.arg] = [ kw: list[ast.arg] = [
ast.arg(arg=arg.name, annotation=self.dump_type(arg.type)) ast.arg(
for arg in func.kw_args arg=param.name,
annotation=self.dump_type(param.type),
)
for param in params.kw
] ]
defaults: list[ast.expr] = [ 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]] = [ 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: if with_self:
arg = ast.arg(arg="self", annotation=None) arg = ast.arg(arg="self", annotation=None)
@@ -391,7 +401,7 @@ class StubsGenerator:
body=[ body=[
ast.FunctionDef( ast.FunctionDef(
name="__call__", 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), returns=self.dump_type(func.returns),
body=[ast.Expr(value=Empty)], body=[ast.Expr(value=Empty)],
decorator_list=[], decorator_list=[],

5
midas/lexer/base.py
View File

@@ -16,9 +16,10 @@ class Lexer(ABC):
"""An abstract lexer which provides methods to easily extend it into a concrete one """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, 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/ [1]: https://craftinginterpreters.com/
[2]: https://git.kb28.ch/HEL/pebble
""" """
def __init__(self, source: str, file: Optional[str] = None) -> None: def __init__(self, source: str, file: Optional[str] = None) -> None:
@@ -168,6 +169,6 @@ class Lexer(ABC):
def scan_token(self) -> None: def scan_token(self) -> None:
"""Scan a token """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 pass

18
midas/lexer/midas.py
View File

@@ -81,6 +81,12 @@ class MidasLexer(Lexer):
return None return None
def scan_string(self, opening: str): 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(): while self.peek() != opening and not self.is_at_end():
self.advance() self.advance()
@@ -147,6 +153,18 @@ class MidasLexer(Lexer):
self.add_token(TokenType.COMMENT) self.add_token(TokenType.COMMENT)
def is_identifier_char(self, char: str, *, start: bool) -> bool: 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 == "_": if char == "_":
return True return True
if char.isalpha(): if char.isalpha():

9
midas/lexer/token.py
View File

@@ -104,6 +104,15 @@ class Token:
) )
def location_to(self, to: Token) -> Location: 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()) return Location.span(self.get_location(), to.get_location())
@property @property

8
midas/parser/base.py
View File

@@ -16,6 +16,9 @@ class TokenError:
def get_report(self) -> str: def get_report(self) -> str:
"""Get a detailed error message """Get a detailed error message
The error message is formatted as "(<position>) Error at <token>: <message>".
For example: "(L2:5) Error at '3': Expected ')' after arguments."
Returns: Returns:
str: the complete error message 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 """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, 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/ [1]: https://craftinginterpreters.com/
[2]: https://git.kb28.ch/HEL/pebble
""" """
IGNORE: set[TokenType] = { IGNORE: set[TokenType] = {
@@ -173,7 +177,7 @@ class Parser(ABC, Generic[T]):
error_msg (str): the error message if the token doesn't match error_msg (str): the error message if the token doesn't match
Raises: Raises:
SyntaxError: if the current token doesn't match the given type ParsingError: if the current token doesn't match the given type
Returns: Returns:
Token: the current token which matched the given type Token: the current token which matched the given type

330
midas/parser/midas.py
View File

@@ -35,10 +35,11 @@ from midas.parser.base import Parser
from midas.parser.errors import ParsingError from midas.parser.errors import ParsingError
class MidasParser(Parser): class MidasParser(Parser[list[Stmt]]):
"""A simple parser for midas type definitions""" """A simple parser for midas type definitions"""
SYNC_BOUNDARY: set[TokenType] = { SYNC_BOUNDARY: set[TokenType] = {
TokenType.ALIAS,
TokenType.TYPE, TokenType.TYPE,
TokenType.EXTEND, TokenType.EXTEND,
TokenType.PREDICATE, TokenType.PREDICATE,
@@ -73,10 +74,10 @@ class MidasParser(Parser):
def declaration(self) -> Optional[Stmt]: def declaration(self) -> Optional[Stmt]:
"""Try and parse a declaration """Try and parse a declaration
Any parsing error is caught and None is returned Any parsing error is caught and `None` is returned
Returns: 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: try:
if self.match(TokenType.TYPE): if self.match(TokenType.TYPE):
@@ -95,23 +96,14 @@ class MidasParser(Parser):
def type_declaration(self) -> TypeStmt: def type_declaration(self) -> TypeStmt:
"""Parse a type declaration """Parse a type declaration
A type declaration can either be a simple type alias or a new complex type. A type declaration creates a named subtype of a type expression.
In either case, it can have an optional template expression after its name, wrapped in brackets. It can have an optional template expression after its name, wrapped in brackets, to handle type parameters.
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 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 type statement consists of:
A full complex type definition is thus written: - the `type` keyword
``` - a name (identifier)
type Name[Template] { - (optional) type parameters
prop1: TypeExpr1 where Condition1 - a body, a type expression (see :func:`type_expr`)
prop2: TypeExpr2 where Condition2
...
}
```
Returns: Returns:
TypeStmt: the parsed type declaration statement TypeStmt: the parsed type declaration statement
@@ -165,11 +157,16 @@ class MidasParser(Parser):
def alias_declaration(self) -> AliasStmt: def alias_declaration(self) -> AliasStmt:
"""Parse an alias declaration """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: Returns:
AliasStmt: the parsed alias declaration statement AliasStmt: the parsed alias declaration statement
""" """
keyword: Token = self.previous() 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") self.consume(TokenType.EQUAL, "Expected '=' before alias definition")
@@ -184,8 +181,8 @@ class MidasParser(Parser):
def type_expr(self) -> Type: def type_expr(self) -> Type:
"""Parse a type expression """Parse a type expression
A type is an identifier, optionally followed by a template expression. A type expression can either be a function type (see :func:`function`)
It can also optionally be followed by a '?' to indicate a nullable type or a constraint type (see :func:`constraint_type`)
Returns: Returns:
TypeExpr: the parsed type expression TypeExpr: the parsed type expression
@@ -205,6 +202,15 @@ class MidasParser(Parser):
return base return base
def constraint_type(self) -> Type: 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() type: Type = self.base_type()
if self.match(TokenType.WHERE): if self.match(TokenType.WHERE):
constraint: Expr = self.constraint() constraint: Expr = self.constraint()
@@ -216,6 +222,14 @@ class MidasParser(Parser):
return type return type
def base_type(self) -> 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): if self.match(TokenType.LEFT_PAREN):
type: Type = self.type_expr() type: Type = self.type_expr()
self.consume(TokenType.RIGHT_PAREN, "Unclosed parenthesis") self.consume(TokenType.RIGHT_PAREN, "Unclosed parenthesis")
@@ -227,6 +241,17 @@ class MidasParser(Parser):
return self.generic_type() return self.generic_type()
def generic_type(self) -> 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() type: NamedType = self.named_type()
if self.check(TokenType.LEFT_BRACKET): if self.check(TokenType.LEFT_BRACKET):
if type.name.lexeme == "Frame": if type.name.lexeme == "Frame":
@@ -240,6 +265,13 @@ class MidasParser(Parser):
return type return type
def type_args(self) -> list[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] = [] args: list[Type] = []
self.consume(TokenType.LEFT_BRACKET, "Missing '[' before generic arguments") self.consume(TokenType.LEFT_BRACKET, "Missing '[' before generic arguments")
while not self.is_at_end() and not self.check(TokenType.RIGHT_BRACKET): while not self.is_at_end() and not self.check(TokenType.RIGHT_BRACKET):
@@ -250,6 +282,13 @@ class MidasParser(Parser):
return args return args
def named_type(self) -> NamedType: 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") name: Token = self.consume_identifier("Expected type name")
return NamedType( return NamedType(
location=name.get_location(), location=name.get_location(),
@@ -257,13 +296,13 @@ class MidasParser(Parser):
) )
def complex_type(self) -> ComplexType: 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 A complex type consists of zero or more member statements enclosed in
property statements enclosed in curly braces curly braces
Returns: Returns:
ComplexType: the parsed complex type ComplexType: the parsed complex type expression
""" """
left: Token = self.consume( left: Token = self.consume(
TokenType.LEFT_BRACE, "Expected '{' to start type body" TokenType.LEFT_BRACE, "Expected '{' to start type body"
@@ -285,6 +324,20 @@ class MidasParser(Parser):
) )
def frame_type(self) -> FrameType: 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() keyword: Token = self.previous()
self.consume(TokenType.LEFT_BRACKET, "Expected '[' to start frame schema") self.consume(TokenType.LEFT_BRACKET, "Expected '[' to start frame schema")
@@ -311,9 +364,9 @@ class MidasParser(Parser):
) )
def constraint(self) -> Expr: 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: Returns:
Expr: the parsed constraint expression Expr: the parsed constraint expression
@@ -321,10 +374,20 @@ class MidasParser(Parser):
return self.expression() return self.expression()
def expression(self) -> Expr: 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_() return self.and_()
def and_(self) -> Expr: 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: Returns:
Expr: the parsed expression Expr: the parsed expression
@@ -340,7 +403,10 @@ class MidasParser(Parser):
return expr return expr
def equality(self) -> 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: Returns:
Expr: the parsed expression Expr: the parsed expression
@@ -356,7 +422,10 @@ class MidasParser(Parser):
return expr return expr
def comparison(self) -> 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: Returns:
Expr: the parsed expression Expr: the parsed expression
@@ -377,6 +446,14 @@ class MidasParser(Parser):
return expr return expr
def term(self) -> 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() expr: Expr = self.factor()
while self.match(TokenType.PLUS, TokenType.MINUS): while self.match(TokenType.PLUS, TokenType.MINUS):
operator: Token = self.previous() operator: Token = self.previous()
@@ -388,6 +465,14 @@ class MidasParser(Parser):
return expr return expr
def factor(self) -> 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() expr: Expr = self.unary()
while self.match(TokenType.STAR, TokenType.SLASH): while self.match(TokenType.STAR, TokenType.SLASH):
operator: Token = self.previous() operator: Token = self.previous()
@@ -399,12 +484,15 @@ class MidasParser(Parser):
return expr return expr
def unary(self) -> 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: Returns:
Expr: the parsed expression Expr: the parsed expression
""" """
if self.match(TokenType.MINUS): if self.match(TokenType.PLUS, TokenType.MINUS):
operator: Token = self.previous() operator: Token = self.previous()
right: Expr = self.unary() right: Expr = self.unary()
location: Location = Location.span(operator.get_location(), right.location) location: Location = Location.span(operator.get_location(), right.location)
@@ -412,12 +500,44 @@ class MidasParser(Parser):
return self.call() return self.call()
def call(self) -> Expr: 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() expr: Expr = self.reference()
while self.match(TokenType.LEFT_PAREN): while self.match(TokenType.LEFT_PAREN):
expr = self.finish_call(expr) expr = self.finish_call(expr)
return expr return expr
def finish_call(self, callee: Expr) -> 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] = [] pos_args: list[Expr] = []
kw_args: dict[str, Expr] = {} kw_args: dict[str, Expr] = {}
keywords: bool = False keywords: bool = False
@@ -437,13 +557,14 @@ class MidasParser(Parser):
else: else:
value = self.expression() value = self.expression()
if self.check(TokenType.EQUAL): if self.check(TokenType.EQUAL):
error_msg: str
if keywords: if keywords:
raise self.error(self.peek(), "Invalid keyword argument name") error_msg = "Invalid keyword argument name"
else: else:
raise self.error( error_msg = (
self.peek(), "Cannot pass positional arguments after a keyword argument"
"Cannot pass positional arguments after a keyword argument",
) )
raise self.error(self.peek(), error_msg)
pos_args.append(value) pos_args.append(value)
if not self.match(TokenType.COMMA): if not self.match(TokenType.COMMA):
@@ -460,7 +581,12 @@ class MidasParser(Parser):
) )
def reference(self) -> Expr: 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: Returns:
Expr: the parsed expression Expr: the parsed expression
@@ -475,7 +601,12 @@ class MidasParser(Parser):
def primary(self) -> Expr: def primary(self) -> Expr:
"""Parse a primary expression """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: Returns:
Expr: the parsed expression Expr: the parsed expression
@@ -508,14 +639,41 @@ class MidasParser(Parser):
raise self.error(self.peek(), "Expected expression") raise self.error(self.peek(), "Expected expression")
def consume_identifier(self, message: str = "Expected identifier") -> Token: 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(): if not self.match_identifier():
raise self.error(self.peek(), message) raise self.error(self.peek(), message)
return self.previous() return self.previous()
def match_identifier(self) -> bool: 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()) return self.match(TokenType.IDENTIFIER, *KEYWORDS.values())
def check_identifier(self) -> bool: 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()]: for tt in [TokenType.IDENTIFIER, *KEYWORDS.values()]:
if self.check(tt): if self.check(tt):
return True return True
@@ -524,7 +682,14 @@ class MidasParser(Parser):
def member_stmt(self) -> MemberStmt: def member_stmt(self) -> MemberStmt:
"""Parse a member statement """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: Returns:
MemberStmt: the parsed member statement MemberStmt: the parsed member statement
@@ -551,7 +716,13 @@ class MidasParser(Parser):
def extend_declaration(self) -> ExtendStmt: def extend_declaration(self) -> ExtendStmt:
"""Parse an extension definition """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: Returns:
ExtendStmt: the parsed extension statement ExtendStmt: the parsed extension statement
@@ -576,7 +747,12 @@ class MidasParser(Parser):
def predicate_declaration(self) -> PredicateStmt: def predicate_declaration(self) -> PredicateStmt:
"""Parse a predicate declaration """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: Returns:
PredicateStmt: the parsed predicate declaration statement PredicateStmt: the parsed predicate declaration statement
@@ -587,7 +763,7 @@ class MidasParser(Parser):
params: list[ParamSpec] = [] params: list[ParamSpec] = []
while self.check(TokenType.LEFT_PAREN): while self.check(TokenType.LEFT_PAREN):
params.append(self.function_args()) params.append(self.function_params())
self.consume(TokenType.EQUAL, "Expected '=' after predicate subject") self.consume(TokenType.EQUAL, "Expected '=' after predicate subject")
body: Expr = self.constraint() body: Expr = self.constraint()
@@ -599,7 +775,18 @@ class MidasParser(Parser):
) )
def function(self) -> FunctionType: 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") self.consume(TokenType.ARROW, "Expected '->' before result type")
result: Type = self.type_expr() result: Type = self.type_expr()
@@ -610,36 +797,53 @@ class MidasParser(Parser):
returns=result, 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( l_paren: Token = self.consume(
TokenType.LEFT_PAREN, "Expected '(' before function parameters" TokenType.LEFT_PAREN, "Expected '(' before function parameters"
) )
pos_args: list[FunctionType.Argument] = [] pos: list[FunctionType.Parameter] = []
args: list[FunctionType.Argument] = [] mixed: list[FunctionType.Parameter] = []
kw_args: list[FunctionType.Argument] = [] kw: list[FunctionType.Parameter] = []
args_first_tokens: list[Token] = [] mixed_first_tokens: list[Token] = []
section: int = 0 section: int = 0
while not self.is_at_end() and not self.check(TokenType.RIGHT_PAREN): while not self.is_at_end() and not self.check(TokenType.RIGHT_PAREN):
match section: match section:
case 0 if self.match(TokenType.SLASH): case 0 if self.match(TokenType.SLASH):
pos_args = args pos = mixed
args = [] mixed = []
args_first_tokens = [] mixed_first_tokens = []
section = 1 section = 1
case 0 | 1 if self.match(TokenType.STAR): case 0 | 1 if self.match(TokenType.STAR):
section = 2 section = 2
case _: case _:
# Record first token of mixed argument for errors if unnamed # Record first token of mixed parameters for errors if unnamed
if section != 2: if section != 2:
args_first_tokens.append(self.peek()) mixed_first_tokens.append(self.peek())
name: Optional[Token] = None name: Optional[Token] = None
if section == 2: if section == 2:
name = self.consume_identifier("Expected keyword argument name") name = self.consume_identifier(
"Expected keyword parameter name"
)
self.consume( self.consume(
TokenType.COLON, "Expected ':' after argument name" TokenType.COLON, "Expected ':' after parameter name"
) )
elif self.check_identifier() and self.check_next(TokenType.COLON): elif self.check_identifier() and self.check_next(TokenType.COLON):
name = self.advance() name = self.advance()
@@ -647,24 +851,24 @@ class MidasParser(Parser):
type: Type = self.type_expr() type: Type = self.type_expr()
optional: bool = self.match(TokenType.QMARK) optional: bool = self.match(TokenType.QMARK)
arg = FunctionType.Argument( param = FunctionType.Parameter(
location=None, location=None,
name=name, name=name,
type=type, type=type,
required=not optional, required=not optional,
) )
if section == 2: if section == 2:
kw_args.append(arg) kw.append(param)
else: else:
args.append(arg) mixed.append(param)
if not self.match(TokenType.COMMA): if not self.match(TokenType.COMMA):
break break
for arg, token in zip(args, args_first_tokens): for param, token in zip(mixed, mixed_first_tokens):
if arg.name is None: if param.name is None:
# Not raised because we can keep parsing # 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") 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)

90
midas/parser/python.py
View File

@@ -23,6 +23,7 @@ from midas.ast.python import (
LiteralExpr, LiteralExpr,
LogicalExpr, LogicalExpr,
MidasType, MidasType,
ParamSpec,
RawExpr, RawExpr,
RawStmt, RawStmt,
ReturnStmt, ReturnStmt,
@@ -49,6 +50,8 @@ class UnsupportedSyntaxError(Exception):
class PythonParser: class PythonParser:
"""A parser to convert raw Python `ast` nodes in custom IR nodes"""
CAST_FUNCTION = "cast" CAST_FUNCTION = "cast"
UNSAFE_CAST_FUNCTION = "unsafe_cast" UNSAFE_CAST_FUNCTION = "unsafe_cast"
@@ -212,27 +215,10 @@ class PythonParser:
match node: match node:
case ast.FunctionDef( case ast.FunctionDef(
name=name, name=name,
args=ast.arguments(
posonlyargs=posonlyargs,
args=args, args=args,
vararg=sink,
kwonlyargs=kwonlyargs,
kwarg=kw_sink,
defaults=defaults,
kw_defaults=kw_defaults,
),
returns=returns, returns=returns,
body=raw_body, 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] = [] body: list[Stmt] = []
for stmt in raw_body: for stmt in raw_body:
stmts = self.parse_stmt(stmt) stmts = self.parse_stmt(stmt)
@@ -241,54 +227,58 @@ class PythonParser:
elif stmts is not None: elif stmts is not None:
body.extend(stmts) 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( return Function(
location=loc, location=loc,
name=name, name=name,
posonlyargs=parse_args(posonlyargs, posargs_defaults), params=self._parse_param_spec(args),
args=parse_args(args, args_defaults),
sink=(
self._parse_function_argument(sink, None)
if sink is not None
else None
),
kwonlyargs=parse_args(kwonlyargs, kwargs_defaults),
kw_sink=(
self._parse_function_argument(kw_sink, None)
if kw_sink is not None
else None
),
returns=self._parse_type(returns) if returns is not None else None, returns=self._parse_type(returns) if returns is not None else None,
body=body, body=body,
) )
case _: case _:
print(f"Unsupported function definition: {ast.unparse(node)}") 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] self, arg: ast.arg, default: Optional[Expr]
) -> Function.Argument: ) -> Function.Parameter:
loc: Location = Location.from_ast(arg) loc: Location = Location.from_ast(arg)
name: str = arg.arg name: str = arg.arg
type: Optional[MidasType] = None type: Optional[MidasType] = None
if arg.annotation is not None: if arg.annotation is not None:
type = self._parse_type(arg.annotation) type = self._parse_type(arg.annotation)
return Function.Argument( return Function.Parameter(
location=loc, location=loc,
name=name, name=name,
type=type, type=type,

126
tests/cases/checker/03_functions.py.ref.json
View File

@@ -124,7 +124,7 @@
22 22
] ]
}, },
"message": "Multiple values for argument 'b'" "message": "Multiple values for parameter 'b'"
}, },
{ {
"type": "Error", "type": "Error",
@@ -152,7 +152,7 @@
12 12
] ]
}, },
"message": "Unknown keyword argument 'a'" "message": "Unknown keyword parameter 'a'"
}, },
{ {
"type": "Error", "type": "Error",
@@ -194,7 +194,7 @@
17 17
] ]
}, },
"message": "Unknown keyword argument 'g'" "message": "Unknown keyword parameter 'g'"
}, },
{ {
"type": "Error", "type": "Error",
@@ -277,7 +277,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -287,7 +288,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -297,7 +298,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -306,7 +307,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -351,7 +353,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -361,7 +364,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -371,7 +374,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -380,7 +383,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -443,7 +447,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -453,7 +458,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -463,7 +468,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -472,7 +477,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -539,7 +545,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -549,7 +556,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -559,7 +566,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -568,7 +575,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -649,7 +657,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -659,7 +668,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -669,7 +678,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -678,7 +687,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -762,7 +772,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -772,7 +783,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -782,7 +793,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -791,7 +802,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -850,7 +862,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -860,7 +873,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -870,7 +883,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -879,7 +892,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -929,7 +943,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -939,7 +954,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -949,7 +964,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -958,7 +973,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -1034,7 +1050,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -1044,7 +1061,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -1054,7 +1071,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -1063,7 +1080,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -1150,7 +1168,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -1160,7 +1179,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -1170,7 +1189,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -1179,7 +1198,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -1266,7 +1286,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -1276,7 +1297,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -1286,7 +1307,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -1295,7 +1316,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -1382,7 +1404,8 @@
"name": "foo" "name": "foo"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -1392,7 +1415,7 @@
"required": true "required": true
} }
], ],
"args": [ "mixed": [
{ {
"pos": 1, "pos": 1,
"name": "b", "name": "b",
@@ -1402,7 +1425,7 @@
"required": true "required": true
} }
], ],
"kw_args": [ "kw": [
{ {
"pos": 2, "pos": 2,
"name": "c", "name": "c",
@@ -1411,7 +1434,8 @@
}, },
"required": true "required": true
} }
], ]
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }

8
tests/cases/checker/06_subtyping.py.ref.json
View File

@@ -136,8 +136,9 @@
"name": "maximum" "name": "maximum"
}, },
"type": { "type": {
"pos_args": [], "params": {
"args": [ "pos": [],
"mixed": [
{ {
"pos": 0, "pos": 0,
"name": "a", "name": "a",
@@ -155,7 +156,8 @@
"required": true "required": true
} }
], ],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "float" "name": "float"
} }

8
tests/cases/checker/07_variance.py.ref.json
View File

@@ -312,7 +312,8 @@
"name": "print" "name": "print"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "object", "name": "object",
@@ -320,8 +321,9 @@
"required": false "required": false
} }
], ],
"args": [], "mixed": [],
"kw_args": [], "kw": []
},
"returns": {} "returns": {}
} }
}, },

80
tests/cases/checker/08_unification.py.ref.json
View File

@@ -120,7 +120,8 @@
"name": "bool" "name": "bool"
}, },
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "object", "name": "object",
@@ -128,8 +129,9 @@
"required": false "required": false
} }
], ],
"args": [], "mixed": [],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -377,8 +379,9 @@
"name": "double" "name": "double"
}, },
"type": { "type": {
"pos_args": [], "params": {
"args": [ "pos": [],
"mixed": [
{ {
"pos": 0, "pos": 0,
"name": "value", "name": "value",
@@ -388,7 +391,8 @@
"required": true "required": true
} }
], ],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "float" "name": "float"
} }
@@ -439,12 +443,14 @@
} }
], ],
"body": { "body": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "transform", "name": "transform",
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "v", "name": "v",
@@ -456,8 +462,9 @@
"required": true "required": true
} }
], ],
"args": [], "mixed": [],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "U", "name": "U",
"bound": null, "bound": null,
@@ -485,8 +492,9 @@
"required": true "required": true
} }
], ],
"args": [], "mixed": [],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "list", "name": "list",
"args": [ "args": [
@@ -548,8 +556,9 @@
"name": "double" "name": "double"
}, },
"type": { "type": {
"pos_args": [], "params": {
"args": [ "pos": [],
"mixed": [
{ {
"pos": 0, "pos": 0,
"name": "value", "name": "value",
@@ -559,7 +568,8 @@
"required": true "required": true
} }
], ],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "float" "name": "float"
} }
@@ -610,12 +620,14 @@
} }
], ],
"body": { "body": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "transform", "name": "transform",
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "v", "name": "v",
@@ -627,8 +639,9 @@
"required": true "required": true
} }
], ],
"args": [], "mixed": [],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "U", "name": "U",
"bound": null, "bound": null,
@@ -656,8 +669,9 @@
"required": true "required": true
} }
], ],
"args": [], "mixed": [],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "list", "name": "list",
"args": [ "args": [
@@ -709,8 +723,9 @@
"name": "is_odd" "name": "is_odd"
}, },
"type": { "type": {
"pos_args": [], "params": {
"args": [ "pos": [],
"mixed": [
{ {
"pos": 0, "pos": 0,
"name": "value", "name": "value",
@@ -720,7 +735,8 @@
"required": true "required": true
} }
], ],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "bool" "name": "bool"
} }
@@ -771,12 +787,14 @@
} }
], ],
"body": { "body": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "transform", "name": "transform",
"type": { "type": {
"pos_args": [ "params": {
"pos": [
{ {
"pos": 0, "pos": 0,
"name": "v", "name": "v",
@@ -788,8 +806,9 @@
"required": true "required": true
} }
], ],
"args": [], "mixed": [],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "U", "name": "U",
"bound": null, "bound": null,
@@ -817,8 +836,9 @@
"required": true "required": true
} }
], ],
"args": [], "mixed": [],
"kw_args": [], "kw": []
},
"returns": { "returns": {
"name": "list", "name": "list",
"args": [ "args": [

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

@@ -7,8 +7,10 @@
{ {
"_type": "Function", "_type": "Function",
"name": "func", "name": "func",
"posonlyargs": [], "params": {
"args": [ "_type": "ParamSpec",
"pos": [],
"mixed": [
{ {
"name": "col1", "name": "col1",
"type": { "type": {
@@ -48,9 +50,8 @@
"default": null "default": null
} }
], ],
"sink": null, "kw": []
"kwonlyargs": [], },
"kw_sink": null,
"returns": { "returns": {
"_type": "BaseType", "_type": "BaseType",
"base": "Column", "base": "Column",
@@ -119,7 +120,9 @@
{ {
"_type": "Function", "_type": "Function",
"name": "func2", "name": "func2",
"posonlyargs": [ "params": {
"_type": "ParamSpec",
"pos": [
{ {
"name": "a", "name": "a",
"type": { "type": {
@@ -130,7 +133,7 @@
"default": null "default": null
} }
], ],
"args": [ "mixed": [
{ {
"name": "b", "name": "b",
"type": { "type": {
@@ -141,8 +144,7 @@
"default": null "default": null
} }
], ],
"sink": null, "kw": [
"kwonlyargs": [
{ {
"name": "c", "name": "c",
"type": { "type": {
@@ -152,8 +154,8 @@
}, },
"default": null "default": null
} }
], ]
"kw_sink": null, },
"returns": null, "returns": null,
"body": [] "body": []
} }

14
tests/serializer/midas.py
View File

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

35
tests/serializer/python.py
View File

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