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

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

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

@@ -142,26 +142,26 @@ class MidasPrinter(
return f"fn {spec} -> {type.returns.accept(self)}" return f"fn {spec} -> {type.returns.accept(self)}"
def _visit_param_spec(self, spec: m.ParamSpec) -> str: def _visit_param_spec(self, spec: m.ParamSpec) -> str:
pos_args: list[str] = [self._print_arg(arg) for arg in spec.pos] pos: list[str] = [self._print_param(param) for param in spec.pos]
mixed_args: list[str] = [self._print_arg(arg) for arg in spec.mixed] mixed: list[str] = [self._print_param(param) for param in spec.mixed]
kw_args: list[str] = [self._print_arg(arg) for arg in spec.kw] kw: list[str] = [self._print_param(param) for param in spec.kw]
args: list[str] = pos_args params: list[str] = pos
if len(pos_args) != 0: if len(pos) != 0:
args.append("/") params.append("/")
args += mixed_args params += mixed
if len(kw_args) != 0: if len(kw) != 0:
args.append("*") params.append("*")
args += kw_args params += kw
return f"({', '.join(args)})" return f"({', '.join(params)})"
def _print_arg(self, arg: m.FunctionType.Argument) -> str: def _print_param(self, param: m.FunctionType.Parameter) -> str:
res: str = "" res: str = ""
if arg.name is not None: if param.name is not None:
res += arg.name.lexeme res += param.name.lexeme
res += ": " res += ": "
res += arg.type.accept(self) res += param.type.accept(self)
if not arg.required: if not param.required:
res += "?" res += "?"
return res return res

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

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

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

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

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)

82
midas/checker/dispatcher.py
View File

@@ -30,9 +30,9 @@ TypedExpr = tuple[E, Type]
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class MappedArgument(Generic[E]): class MappedArgument(Generic[E]):
expr: E arg_expr: E
type: Type arg_type: Type
argument: Function.Argument parameter: Function.Parameter
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
@@ -219,11 +219,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
@@ -347,28 +347,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 +387,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 +416,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
@@ -474,11 +476,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

12
midas/checker/evaluator.py
View File

@@ -158,15 +158,17 @@ 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 positional: list[Function.Parameter] = (
keywords: dict[str, Function.Argument] = { function.params.pos + function.params.mixed
arg.name: arg for arg in function.args + function.kw_args )
keywords: dict[str, Function.Parameter] = {
param.name: param for param in function.params.mixed + function.params.kw
} }
for i, arg in enumerate(args): 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)

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.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
@@ -38,31 +45,31 @@ 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] = [] real_params: list[Function.Parameter] = []
for i, arg in enumerate(args): for i, param in enumerate(params):
match arg: match param:
case str() as name: 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

76
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,
@@ -82,14 +83,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 +166,14 @@ 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:
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 +181,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 +226,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 +244,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 +256,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 +280,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 +305,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),
) )

82
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,
@@ -150,14 +151,16 @@ class FrameMethodRegistry(MethodRegistry[Call]):
# 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 +230,11 @@ 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:
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 +242,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 +256,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
), ),
*kwargs, *kwargs,
], ],
),
returns=TopType(), returns=TopType(),
) )
overload = OverloadedFunction( overload = OverloadedFunction(
@@ -300,7 +307,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 +325,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 +337,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 +361,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 +386,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),
) )

50
midas/checker/midas.py
View File

@@ -21,6 +21,7 @@ from midas.checker.types import (
ExtensionType, ExtensionType,
Function, Function,
GenericType, GenericType,
ParamSpec,
Predicate, Predicate,
Type, Type,
TypeVar, TypeVar,
@@ -32,13 +33,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): class ReturnException(Exception):
pass pass
@@ -47,7 +41,7 @@ class ReturnException(Exception):
class MappedArgument: class MappedArgument:
expr: m.Expr expr: m.Expr
type: Type type: Type
argument: Function.Argument argument: Function.Parameter
@dataclass(frozen=True, kw_only=True) @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) 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 +201,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 = {}
@@ -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: 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:

13
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,
@@ -132,9 +133,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 +145,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:

90
midas/checker/python.py
View File

@@ -31,6 +31,7 @@ from midas.checker.types import (
FrameGroupBy, FrameGroupBy,
Function, Function,
GenericType, GenericType,
ParamSpec,
TopType, TopType,
TupleType, TupleType,
Type, Type,
@@ -59,7 +60,7 @@ class UndefinedMethodException(Exception):
class MappedArgument: class MappedArgument:
expr: p.Expr expr: p.Expr
type: Type type: Type
argument: Function.Argument argument: Function.Parameter
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
@@ -289,61 +290,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 +379,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:

104
midas/checker/registry.py
View File

@@ -234,92 +234,100 @@ 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

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

83
midas/checker/types.py
View File

@@ -45,28 +45,14 @@ class UnitType:
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class Function: class Function:
pos_args: list[Argument] = field(default_factory=list) params: ParamSpec
args: list[Argument] = field(default_factory=list)
kw_args: list[Argument] = field(default_factory=list)
returns: Type 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,6 +63,28 @@ class Function:
return f"{self.name}: {self.type}{opt}" 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) @dataclass(frozen=True, kw_only=True)
class OverloadedFunction: class OverloadedFunction:
overloads: list[Type] overloads: list[Type]
@@ -204,12 +212,19 @@ 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): def sub_parameter(param: Function.Parameter):
return Function.Argument( return Function.Parameter(
pos=arg.pos, pos=param.pos,
name=arg.name, name=param.name,
type=substitute_typevars(arg.type, substitutions), type=substitute_typevars(param.type, substitutions),
required=arg.required, required=param.required,
)
def sub_param_spec(spec: ParamSpec):
return ParamSpec(
pos=list(map(sub_parameter, spec.pos)),
mixed=list(map(sub_parameter, spec.mixed)),
kw=list(map(sub_parameter, spec.kw)),
) )
def sub_column(col: DataFrameType.Column): def sub_column(col: DataFrameType.Column):
@@ -235,15 +250,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),
) )
@@ -351,14 +362,14 @@ def unfold_type(type: Type) -> Type:
def to_annotation(type: Type) -> str: def to_annotation(type: Type) -> str:
def _args_annotation(func: Function) -> str: def _params_annotation(spec: ParamSpec) -> str:
if len(func.kw_args) != 0: if len(spec.kw) != 0:
return "..." 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 +387,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():

38
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,
@@ -29,8 +30,9 @@ class Unifier:
keywords: dict[str, Type], keywords: dict[str, Type],
) -> Optional[Type]: ) -> Optional[Type]:
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 +40,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 +49,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)
@@ -125,7 +127,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] = {}
@@ -161,19 +163,23 @@ 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 pos1: list[Function.Parameter] = func1.params.pos
mixed1: list[Function.Argument] = func1.args mixed1: list[Function.Parameter] = func1.params.mixed
kw1: list[Function.Argument] = func1.kw_args kw1: list[Function.Parameter] = func1.params.kw
pos2: list[Function.Argument] = func2.pos_args pos2: list[Function.Parameter] = func2.params.pos
mixed2: list[Function.Argument] = func2.args mixed2: list[Function.Parameter] = func2.params.mixed
kw2: list[Function.Argument] = func2.kw_args kw2: list[Function.Parameter] = func2.params.kw
mapped: list[tuple[Function.Argument, Function.Argument]] = [] mapped: list[tuple[Function.Parameter, Function.Parameter]] = []
by_pos2: dict[int, Function.Argument] = {arg.pos: arg for arg in pos2 + mixed2} by_pos2: dict[int, Function.Parameter] = {
by_name2: dict[str, Function.Argument] = {arg.name: arg for arg in mixed2 + kw2} param.pos: param for param in pos2 + mixed2
}
by_name2: dict[str, Function.Parameter] = {
param.name: param for param in mixed2 + kw2
}
for arg1 in pos1: for arg1 in pos1:
if (arg2 := by_pos2.get(arg1.pos)) is not None: 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: 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"])
@@ -109,10 +109,10 @@ class VarianceInferrer:
Variance.COVARIANT: 1, Variance.COVARIANT: 1,
Variance.CONTRAVARIANT: -1, Variance.CONTRAVARIANT: -1,
} }
for arg, param in zip(args, params): for param, param in zip(args, params):
param_polarity: Polarity = polarities[param.variance] param_polarity: Polarity = polarities[param.variance]
self.walk( self.walk(
arg, param,
cast(Polarity, polarity * param_polarity), cast(Polarity, polarity * param_polarity),
base_name, base_name,
path + [f"applied:'{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: 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=[],

44
midas/parser/midas.py
View File

@@ -587,7 +587,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 +599,7 @@ class MidasParser(Parser):
) )
def function(self) -> FunctionType: def function(self) -> FunctionType:
params: ParamSpec = self.function_args() 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 +610,38 @@ class MidasParser(Parser):
returns=result, returns=result,
) )
def function_args(self) -> ParamSpec: def function_params(self) -> ParamSpec:
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 +649,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)

88
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,
@@ -212,27 +213,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 +225,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,

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",