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

refactor: rebrand function parameters and unify spec

Browse Source 
rename function arguments to parameters where it was wrong, and add ParamSpec for Python AST, like for Midas
This commit is contained in:
Louis Heredero
2026-07-03 19:24:30 +02:00
parent 6b7a682dc5
commit 9229f00375
29 changed files with 691 additions and 654 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
gen/midas.py
View File

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

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

8
midas/ast/midas.py
View File

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

32
midas/ast/printer/midas.py
View File

@@ -142,26 +142,26 @@ class MidasPrinter(
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
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_args) != 0:
args.append("/")
args += mixed_args
if len(kw_args) != 0:
args.append("*")
args += kw_args
return f"({', '.join(args)})"
if len(pos) != 0:
params.append("/")
params += mixed
if len(kw) != 0:
params.append("*")
params += kw
return f"({', '.join(params)})"
def _print_arg(self, arg: m.FunctionType.Argument) -> str:
def _print_param(self, param: m.FunctionType.Parameter) -> str:
res: str = ""
if arg.name is not None:
res += arg.name.lexeme
if param.name is not None:
res += param.name.lexeme
res += ": "
res += arg.type.accept(self)
if not arg.required:
res += param.type.accept(self)
if not param.required:
res += "?"
return res

18
midas/ast/printer/midas_ast.py
View File

@@ -209,31 +209,31 @@ class MidasAstPrinter(
self._write_sequence(
"pos",
spec.pos,
print_func=self._print_function_arg,
print_func=self._print_param,
)
self._write_sequence(
"mixed",
spec.mixed,
print_func=self._print_function_arg,
print_func=self._print_param,
)
self._write_sequence(
"kw",
spec.kw,
print_func=self._print_function_arg,
print_func=self._print_param,
last=True,
)
def _print_function_arg(self, arg: m.FunctionType.Argument) -> None:
self._write_line("Argument")
def _print_param(self, param: m.FunctionType.Parameter) -> None:
self._write_line("Parameter")
with self._child_level():
name: str = "None"
if arg.name is not None:
name = f'"{arg.name.lexeme}"'
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):
arg.type.accept(self)
self._write_line(f"required: {arg.required}", last=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")

47
midas/ast/printer/python_ast.py
View File

@@ -46,31 +46,38 @@ class PythonAstPrinter(
self._write_line("Function")
with self._child_level():
self._write_line(f"name: {stmt.name}")
self._write_sequence(
"posonlyargs",
stmt.posonlyargs,
print_func=self._print_argument,
)
self._write_sequence(
"args",
stmt.args,
print_func=self._print_argument,
)
self._write_sequence(
"kwonlyargs",
stmt.kwonlyargs,
print_func=self._print_argument,
)
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_argument(self, arg: p.Function.Argument) -> None:
self._write_line("FunctionArgument")
def _print_param_spec(self, spec: p.ParamSpec) -> None:
self._write_line("ParamSpec")
with self._child_level():
self._write_line(f"name: {arg.name}")
self._write_optional_child("type", arg.type, last=True)
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")

22
midas/ast/python.py
View File

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

82
midas/checker/dispatcher.py
View File

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

12
midas/checker/evaluator.py
View File

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

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

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

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

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

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

50
midas/checker/midas.py
View File

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

13
midas/checker/preamble.py
View File

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

90
midas/checker/python.py
View File

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

104
midas/checker/registry.py
View File

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

2
midas/checker/resolver.py
View File

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

83
midas/checker/types.py
View File

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

66
midas/checker/unifier.py
View File

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

14
midas/checker/variance.py
View File

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

15
midas/cli/highlighter.py
View File

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

2
midas/cli/hl_python.css
View File

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

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

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),
)
def make_args(self, params: p.ParamSpec) -> ast.arguments:
return ast.arguments(
posonlyargs=[ast.arg(arg=param.name) for param in params.pos],
args=[ast.arg(arg=param.name) for param in params.mixed],
kwonlyargs=[ast.arg(arg=param.name) for param in params.kw],
defaults=[
self.convert(param.default)
for param in params.pos + params.mixed
if param.default is not None
],
kw_defaults=[
self.convert(param.default) if param.default is not None else None
for param in params.kw
],
)
def visit_function(self, stmt: p.Function) -> ast.stmt:
return ast.FunctionDef(
name=stmt.name,
args=ast.arguments(
posonlyargs=[ast.arg(arg=arg.name) for arg in stmt.posonlyargs],
vararg=None,
args=[ast.arg(arg=arg.name) for arg in stmt.args],
kwonlyargs=[ast.arg(arg=arg.name) for arg in stmt.kwonlyargs],
kwarg=None,
defaults=[
self.convert(arg.default)
for arg in stmt.posonlyargs + stmt.args
if arg.default is not None
],
kw_defaults=[
self.convert(arg.default) if arg.default is not None else None
for arg in stmt.kwonlyargs
],
),
args=self.make_args(stmt.params),
body=self._visit_body(stmt.body),
decorator_list=[],
)

32
midas/generator/stubs.py
View File

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

44
midas/parser/midas.py
View File

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

90
midas/parser/python.py
View File

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

14
tests/serializer/midas.py
View File

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

35
tests/serializer/python.py
View File

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