diff --git a/gen/midas.py b/gen/midas.py index fed8626..fb90199 100644 --- a/gen/midas.py +++ b/gen/midas.py @@ -29,9 +29,9 @@ class MemberKind(Enum): @dataclass(frozen=True, kw_only=True) class ParamSpec: l_paren: Token - pos: list[FunctionType.Argument] - mixed: list[FunctionType.Argument] - kw: list[FunctionType.Argument] + pos: list[FunctionType.Parameter] + mixed: list[FunctionType.Parameter] + kw: list[FunctionType.Parameter] ###< @@ -150,7 +150,7 @@ class FunctionType: returns: Type @dataclass(frozen=True, kw_only=True) - class Argument: + class Parameter: location: Optional[Location] = None name: Optional[Token] type: Type diff --git a/gen/python.py b/gen/python.py index 73f64c7..60a92f7 100644 --- a/gen/python.py +++ b/gen/python.py @@ -12,6 +12,21 @@ from midas.ast.location import Location ###< +###> Preamble +@dataclass(frozen=True, kw_only=True) +class ParamSpec: + pos: list[Function.Parameter] + mixed: list[Function.Parameter] + kw: list[Function.Parameter] + + @property + def all(self) -> list[Function.Parameter]: + return self.pos + self.mixed + self.kw + + +###< + + ###> MidasType | Type annotations | node class BaseType: base: str @@ -42,25 +57,17 @@ class ExpressionStmt: class Function: name: str - posonlyargs: list[Argument] - args: list[Argument] - sink: Optional[Argument] - kwonlyargs: list[Argument] - kw_sink: Optional[Argument] + params: ParamSpec returns: Optional[MidasType] body: list[Stmt] @dataclass(frozen=True, kw_only=True) - class Argument: + class Parameter: location: Optional[Location] = None name: str type: Optional[MidasType] default: Optional[Expr] - @property - def all_args(self) -> list[Argument]: - return self.posonlyargs + self.args + self.kwonlyargs - class TypeAssign: name: str diff --git a/midas/ast/midas.py b/midas/ast/midas.py index 128b1c0..eae434a 100644 --- a/midas/ast/midas.py +++ b/midas/ast/midas.py @@ -30,9 +30,9 @@ class MemberKind(Enum): @dataclass(frozen=True, kw_only=True) class ParamSpec: l_paren: Token - pos: list[FunctionType.Argument] - mixed: list[FunctionType.Argument] - kw: list[FunctionType.Argument] + pos: list[FunctionType.Parameter] + mixed: list[FunctionType.Parameter] + kw: list[FunctionType.Parameter] ############## @@ -318,7 +318,7 @@ class FunctionType(Type): returns: Type @dataclass(frozen=True, kw_only=True) - class Argument: + class Parameter: location: Optional[Location] = None name: Optional[Token] type: Type diff --git a/midas/ast/printer/midas.py b/midas/ast/printer/midas.py index 8339480..cbb62d7 100644 --- a/midas/ast/printer/midas.py +++ b/midas/ast/printer/midas.py @@ -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 diff --git a/midas/ast/printer/midas_ast.py b/midas/ast/printer/midas_ast.py index 75b68ac..69a0b16 100644 --- a/midas/ast/printer/midas_ast.py +++ b/midas/ast/printer/midas_ast.py @@ -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") diff --git a/midas/ast/printer/python_ast.py b/midas/ast/printer/python_ast.py index 25074de..b51c8f9 100644 --- a/midas/ast/printer/python_ast.py +++ b/midas/ast/printer/python_ast.py @@ -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") diff --git a/midas/ast/python.py b/midas/ast/python.py index 6202fcc..65e896c 100644 --- a/midas/ast/python.py +++ b/midas/ast/python.py @@ -14,6 +14,16 @@ from midas.ast.location import Location T = TypeVar("T") +@dataclass(frozen=True, kw_only=True) +class ParamSpec: + pos: list[Function.Parameter] + mixed: list[Function.Parameter] + kw: list[Function.Parameter] + + @property + def all(self) -> list[Function.Parameter]: + return self.pos + self.mixed + self.kw + #################### # Type annotations # @@ -128,25 +138,17 @@ class ExpressionStmt(Stmt): @dataclass(frozen=True) class Function(Stmt): name: str - posonlyargs: list[Argument] - args: list[Argument] - sink: Optional[Argument] - kwonlyargs: list[Argument] - kw_sink: Optional[Argument] + params: ParamSpec returns: Optional[MidasType] body: list[Stmt] @dataclass(frozen=True, kw_only=True) - class Argument: + class Parameter: location: Optional[Location] = None name: str type: Optional[MidasType] default: Optional[Expr] - @property - def all_args(self) -> list[Argument]: - return self.posonlyargs + self.args + self.kwonlyargs - def accept(self, visitor: Stmt.Visitor[T]) -> T: return visitor.visit_function(self) diff --git a/midas/checker/dispatcher.py b/midas/checker/dispatcher.py index 0a421d3..cef3a7d 100644 --- a/midas/checker/dispatcher.py +++ b/midas/checker/dispatcher.py @@ -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 diff --git a/midas/checker/evaluator.py b/midas/checker/evaluator.py index c21f1c2..3f3fa3e 100644 --- a/midas/checker/evaluator.py +++ b/midas/checker/evaluator.py @@ -158,15 +158,17 @@ class Evaluator(m.Expr.Visitor[Any]): return res def _map_args(self, function: Function, args: list[Any], kwargs: dict[str, Any]): - positional: list[Function.Argument] = function.pos_args + function.args - 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) diff --git a/midas/checker/frames/column_groupby_methods.py b/midas/checker/frames/column_groupby_methods.py index 7470f46..e633be7 100644 --- a/midas/checker/frames/column_groupby_methods.py +++ b/midas/checker/frames/column_groupby_methods.py @@ -7,7 +7,14 @@ import midas.ast.python as p from midas.ast.location import Location from midas.checker.dispatcher import CallResult from midas.checker.frames.utils import MethodRegistry, method -from midas.checker.types import ColumnGroupBy, ColumnType, Function, TopType, Type +from midas.checker.types import ( + ColumnGroupBy, + ColumnType, + Function, + ParamSpec, + TopType, + Type, +) if TYPE_CHECKING: from midas.checker.python import TypedExpr @@ -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 diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 95bce9e..9396d0b 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -13,6 +13,7 @@ from midas.checker.types import ( ColumnType, Function, GenericType, + ParamSpec, TopType, Type, TypeVar, @@ -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), ) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 1be481f..7866784 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -14,6 +14,7 @@ from midas.checker.types import ( FrameGroupBy, Function, OverloadedFunction, + ParamSpec, TopType, Type, UnknownType, @@ -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), ) diff --git a/midas/checker/midas.py b/midas/checker/midas.py index 32716fc..432897e 100644 --- a/midas/checker/midas.py +++ b/midas/checker/midas.py @@ -21,6 +21,7 @@ from midas.checker.types import ( ExtensionType, Function, GenericType, + ParamSpec, Predicate, Type, TypeVar, @@ -32,13 +33,6 @@ from midas.lexer.token import Token from midas.parser.midas import MidasParser -@dataclass(frozen=True, kw_only=True) -class TypedParamSpec: - pos: list[Function.Argument] - mixed: list[Function.Argument] - kw: list[Function.Argument] - - class ReturnException(Exception): pass @@ -47,7 +41,7 @@ class ReturnException(Exception): class MappedArgument: expr: m.Expr type: Type - argument: Function.Argument + argument: Function.Parameter @dataclass(frozen=True, kw_only=True) @@ -196,9 +190,7 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type self._predicate_params[param.name.lexeme] = param.type.accept(self) type: Type = self.type_of(stmt.body) - params: list[TypedParamSpec] = [ - 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: diff --git a/midas/checker/preamble.py b/midas/checker/preamble.py index a60dabb..8fcced5 100644 --- a/midas/checker/preamble.py +++ b/midas/checker/preamble.py @@ -7,6 +7,7 @@ from midas.checker.types import ( Function, GenericType, OverloadedFunction, + ParamSpec, TopType, Type, TypeVar, @@ -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: diff --git a/midas/checker/python.py b/midas/checker/python.py index 9103cf4..56d54ca 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -31,6 +31,7 @@ from midas.checker.types import ( FrameGroupBy, Function, GenericType, + ParamSpec, TopType, TupleType, Type, @@ -59,7 +60,7 @@ class UndefinedMethodException(Exception): class MappedArgument: expr: p.Expr type: Type - argument: Function.Argument + argument: Function.Parameter @dataclass(frozen=True, kw_only=True) @@ -289,61 +290,64 @@ class PythonTyper( def visit_function(self, stmt: p.Function) -> None: env: Environment = Environment(self.env) - pos_args: list[Function.Argument] = [] - 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: diff --git a/midas/checker/registry.py b/midas/checker/registry.py index 92b54a6..431cb47 100644 --- a/midas/checker/registry.py +++ b/midas/checker/registry.py @@ -234,92 +234,100 @@ class TypesRegistry: if not self.is_subtype(func1.returns, func2.returns): return False - pos1: list[Function.Argument] = func1.pos_args - 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 diff --git a/midas/checker/resolver.py b/midas/checker/resolver.py index 43768bd..24f7d69 100644 --- a/midas/checker/resolver.py +++ b/midas/checker/resolver.py @@ -93,7 +93,7 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]): function (p.Function): the function to resolve """ self.begin_scope() - for param in function.all_args: + for param in function.params.all: self.declare(param.name) self.define(param.name) self.resolve(*function.body) diff --git a/midas/checker/types.py b/midas/checker/types.py index 51658d7..1c6aff3 100644 --- a/midas/checker/types.py +++ b/midas/checker/types.py @@ -45,28 +45,14 @@ class UnitType: @dataclass(frozen=True, kw_only=True) class Function: - pos_args: list[Argument] = field(default_factory=list) - 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(): diff --git a/midas/checker/unifier.py b/midas/checker/unifier.py index 138a503..9394ef5 100644 --- a/midas/checker/unifier.py +++ b/midas/checker/unifier.py @@ -8,6 +8,7 @@ from midas.checker.types import ( DataFrameType, Function, GenericType, + ParamSpec, TopType, Type, TypeVar, @@ -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: diff --git a/midas/checker/variance.py b/midas/checker/variance.py index e7f5ac0..af205d7 100644 --- a/midas/checker/variance.py +++ b/midas/checker/variance.py @@ -77,14 +77,14 @@ class VarianceInferrer: match type: # Arguments are negative positions -> flip polarity # Return is positive position -> keep polarity - case Function(pos_args=pos_args, args=mixed_args, kw_args=kw_args): - 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}'"], diff --git a/midas/cli/highlighter.py b/midas/cli/highlighter.py index 8d142f1..15f3251 100644 --- a/midas/cli/highlighter.py +++ b/midas/cli/highlighter.py @@ -157,15 +157,18 @@ class PythonHighlighter( def visit_function(self, stmt: p.Function) -> None: self.wrap(stmt, "function") - for arg in stmt.posonlyargs + stmt.args + stmt.kwonlyargs: - self._highlight_function_argument(arg) + self._highlight_param_spec(stmt.params) for body_stmt in stmt.body: body_stmt.accept(self) - def _highlight_function_argument(self, arg: p.Function.Argument) -> None: - self.wrap(arg, "argument") - if arg.type is not None: - arg.type.accept(self) + def _highlight_param_spec(self, spec: p.ParamSpec) -> None: + for param in spec.all: + self._highlight_function_param(param) + + def _highlight_function_param(self, param: p.Function.Parameter) -> None: + self.wrap(param, "parameter") + if param.type is not None: + param.type.accept(self) def visit_type_assign(self, stmt: p.TypeAssign) -> None: stmt.type.accept(self) diff --git a/midas/cli/hl_python.css b/midas/cli/hl_python.css index 899c461..b0b29d0 100644 --- a/midas/cli/hl_python.css +++ b/midas/cli/hl_python.css @@ -23,7 +23,7 @@ span { --col: 215, 103, 224; } - &.argument { + &.parameter { --col: 103, 192, 224; } } \ No newline at end of file diff --git a/midas/generator/constraints.py b/midas/generator/constraints.py index e840b42..c936884 100644 --- a/midas/generator/constraints.py +++ b/midas/generator/constraints.py @@ -5,6 +5,7 @@ import midas.ast.midas as m from midas.checker.registry import TypesRegistry from midas.checker.types import ( Function, + ParamSpec, Predicate, Type, to_annotation, @@ -54,16 +55,16 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]): return expr.accept(self) case _: func = Function( - pos_args=[], - args=[ - Function.Argument( - pos=0, - name="_", - type=self.types.get_type("Any"), - required=True, - ) - ], - kw_args=[], + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="_", + type=self.types.get_type("Any"), + required=True, + ) + ], + ), returns=self.types.get_type("bool"), ) alias: str = self.make_alias(None) @@ -94,28 +95,28 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]): ) return self.make_func(name, [ast.Return(value=body)], predicate.type) - def make_args(self, func: Function) -> ast.arguments: + def make_args(self, params: ParamSpec) -> ast.arguments: return ast.arguments( posonlyargs=[ ast.arg( - arg=arg.name, - annotation=ast.Constant(value=to_annotation(arg.type)), + arg=param.name, + annotation=ast.Constant(value=to_annotation(param.type)), ) - for arg in func.pos_args + for param in params.pos ], args=[ ast.arg( - arg=arg.name, - annotation=ast.Constant(value=to_annotation(arg.type)), + arg=param.name, + annotation=ast.Constant(value=to_annotation(param.type)), ) - for arg in func.args + for param in params.mixed ], kwonlyargs=[ ast.arg( - arg=arg.name, - annotation=ast.Constant(value=to_annotation(arg.type)), + arg=param.name, + annotation=ast.Constant(value=to_annotation(param.type)), ) - for arg in func.kw_args + for param in params.kw ], defaults=[], kw_defaults=[], @@ -125,11 +126,11 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]): self, name: str, inner_body: list[ast.stmt], type: Type, level: int = 0 ) -> ast.stmt: match type: - case Function(returns=Function()): + case Function(params=params, returns=Function()): inner_name: str = f"inner{level}" return ast.FunctionDef( name=name, - args=self.make_args(type), + args=self.make_args(params), body=[ self.make_func(inner_name, inner_body, type.returns, level + 1), ast.Return(value=ast.Name(id=inner_name)), @@ -138,10 +139,10 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]): decorator_list=[], ) - case Function(): + case Function(params=params): return ast.FunctionDef( name=name, - args=self.make_args(type), + args=self.make_args(params), body=inner_body, returns=ast.Constant(value=to_annotation(type.returns)), decorator_list=[], diff --git a/midas/generator/generator.py b/midas/generator/generator.py index 352a544..07799b8 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -250,25 +250,26 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): value=self.convert(stmt.expr), ) + def make_args(self, params: p.ParamSpec) -> ast.arguments: + return ast.arguments( + posonlyargs=[ast.arg(arg=param.name) for param in params.pos], + args=[ast.arg(arg=param.name) for param in params.mixed], + kwonlyargs=[ast.arg(arg=param.name) for param in params.kw], + defaults=[ + self.convert(param.default) + for param in params.pos + params.mixed + if param.default is not None + ], + kw_defaults=[ + self.convert(param.default) if param.default is not None else None + for param in params.kw + ], + ) + def visit_function(self, stmt: p.Function) -> ast.stmt: return ast.FunctionDef( name=stmt.name, - args=ast.arguments( - posonlyargs=[ast.arg(arg=arg.name) for arg in stmt.posonlyargs], - vararg=None, - args=[ast.arg(arg=arg.name) for arg in stmt.args], - kwonlyargs=[ast.arg(arg=arg.name) for arg in stmt.kwonlyargs], - kwarg=None, - defaults=[ - self.convert(arg.default) - for arg in stmt.posonlyargs + stmt.args - if arg.default is not None - ], - kw_defaults=[ - self.convert(arg.default) if arg.default is not None else None - for arg in stmt.kwonlyargs - ], - ), + args=self.make_args(stmt.params), body=self._visit_body(stmt.body), decorator_list=[], ) diff --git a/midas/generator/stubs.py b/midas/generator/stubs.py index 07ab0db..7ddcb06 100644 --- a/midas/generator/stubs.py +++ b/midas/generator/stubs.py @@ -17,6 +17,7 @@ from midas.checker.types import ( Function, GenericType, OverloadedFunction, + ParamSpec, TopType, TupleType, Type, @@ -328,7 +329,7 @@ class StubsGenerator: return [ ast.FunctionDef( name=name, - args=self.dump_args(method, with_self=True), + args=self.dump_params(method.params, with_self=True), returns=self.dump_type(method.returns), body=[ast.Expr(value=Empty)], decorator_list=[ast.Name(id="overload")] if overloaded else [], @@ -348,24 +349,33 @@ class StubsGenerator: ) ] - def dump_args(self, func: Function, with_self: bool = False) -> ast.arguments: + def dump_params(self, params: ParamSpec, with_self: bool = False) -> ast.arguments: pos: list[ast.arg] = [ - ast.arg(arg=f"_{arg.pos}", annotation=self.dump_type(arg.type)) - for arg in func.pos_args + ast.arg( + arg=f"_{param.pos}", + annotation=self.dump_type(param.type), + ) + for param in params.pos ] mixed: list[ast.arg] = [ - ast.arg(arg=arg.name, annotation=self.dump_type(arg.type)) - for arg in func.args + ast.arg( + arg=param.name, + annotation=self.dump_type(param.type), + ) + for param in params.mixed ] kw: list[ast.arg] = [ - ast.arg(arg=arg.name, annotation=self.dump_type(arg.type)) - for arg in func.kw_args + ast.arg( + arg=param.name, + annotation=self.dump_type(param.type), + ) + for param in params.kw ] defaults: list[ast.expr] = [ - Empty for arg in func.pos_args + func.args if not arg.required + Empty for param in params.pos + params.mixed if not param.required ] kw_defaults: list[Optional[ast.expr]] = [ - None if arg.required else Empty for arg in func.kw_args + None if param.required else Empty for param in params.kw ] if with_self: arg = ast.arg(arg="self", annotation=None) @@ -391,7 +401,7 @@ class StubsGenerator: body=[ ast.FunctionDef( name="__call__", - args=self.dump_args(func, with_self=True), + args=self.dump_params(func.params, with_self=True), returns=self.dump_type(func.returns), body=[ast.Expr(value=Empty)], decorator_list=[], diff --git a/midas/parser/midas.py b/midas/parser/midas.py index 534f95e..a78af5a 100644 --- a/midas/parser/midas.py +++ b/midas/parser/midas.py @@ -587,7 +587,7 @@ class MidasParser(Parser): params: list[ParamSpec] = [] while self.check(TokenType.LEFT_PAREN): - params.append(self.function_args()) + params.append(self.function_params()) self.consume(TokenType.EQUAL, "Expected '=' after predicate subject") body: Expr = self.constraint() @@ -599,7 +599,7 @@ class MidasParser(Parser): ) def function(self) -> FunctionType: - params: ParamSpec = self.function_args() + params: ParamSpec = self.function_params() self.consume(TokenType.ARROW, "Expected '->' before result type") result: Type = self.type_expr() @@ -610,36 +610,38 @@ class MidasParser(Parser): returns=result, ) - def function_args(self) -> ParamSpec: + def function_params(self) -> ParamSpec: l_paren: Token = self.consume( TokenType.LEFT_PAREN, "Expected '(' before function parameters" ) - pos_args: list[FunctionType.Argument] = [] - 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) diff --git a/midas/parser/python.py b/midas/parser/python.py index a6c9a78..f2e4554 100644 --- a/midas/parser/python.py +++ b/midas/parser/python.py @@ -23,6 +23,7 @@ from midas.ast.python import ( LiteralExpr, LogicalExpr, MidasType, + ParamSpec, RawExpr, RawStmt, ReturnStmt, @@ -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, diff --git a/tests/serializer/midas.py b/tests/serializer/midas.py index 081059c..13fcbaa 100644 --- a/tests/serializer/midas.py +++ b/tests/serializer/midas.py @@ -188,16 +188,16 @@ class MidasAstJsonSerializer( def _serialize_param_spec(self, spec: ParamSpec) -> dict: return { "_type": "ParamSpec", - "pos": [self._serialize_func_arg(arg) for arg in spec.pos], - "mixed": [self._serialize_func_arg(arg) for arg in spec.mixed], - "kw": [self._serialize_func_arg(arg) for arg in spec.kw], + "pos": [self._serialize_func_param(arg) for arg in spec.pos], + "mixed": [self._serialize_func_param(arg) for arg in spec.mixed], + "kw": [self._serialize_func_param(arg) for arg in spec.kw], } - def _serialize_func_arg(self, arg: FunctionType.Argument) -> dict: + def _serialize_func_param(self, param: FunctionType.Parameter) -> dict: return { - "name": arg.name.lexeme if arg.name is not None else None, - "type": arg.type.accept(self), - "required": arg.required, + "name": param.name.lexeme if param.name is not None else None, + "type": param.type.accept(self), + "required": param.required, } def visit_extension_type(self, type: ExtensionType) -> dict: diff --git a/tests/serializer/python.py b/tests/serializer/python.py index 4e531cb..afb5521 100644 --- a/tests/serializer/python.py +++ b/tests/serializer/python.py @@ -22,6 +22,7 @@ from midas.ast.python import ( LiteralExpr, LogicalExpr, MidasType, + ParamSpec, Pass, RawExpr, RawStmt, @@ -128,32 +129,30 @@ class PythonAstJsonSerializer( "expr": stmt.expr.accept(self), } - def _serialize_argument(self, arg: Function.Argument) -> dict: - return { - "name": arg.name, - "type": self._serialize_optional(arg.type), - "default": self._serialize_optional(arg.default), - } - def visit_function(self, stmt: Function) -> dict: return { "_type": "Function", "name": stmt.name, - "posonlyargs": [self._serialize_argument(arg) for arg in stmt.posonlyargs], - "args": [self._serialize_argument(arg) for arg in stmt.args], - "sink": ( - self._serialize_argument(stmt.sink) if stmt.sink is not None else None - ), - "kwonlyargs": [self._serialize_argument(arg) for arg in stmt.kwonlyargs], - "kw_sink": ( - self._serialize_argument(stmt.kw_sink) - if stmt.kw_sink is not None - else None - ), + "params": self._serialize_param_spec(stmt.params), "returns": self._serialize_optional(stmt.returns), "body": self._serialize_list(stmt.body), } + def _serialize_param_spec(self, spec: ParamSpec) -> dict: + return { + "_type": "ParamSpec", + "pos": [self._serialize_func_param(arg) for arg in spec.pos], + "mixed": [self._serialize_func_param(arg) for arg in spec.mixed], + "kw": [self._serialize_func_param(arg) for arg in spec.kw], + } + + def _serialize_func_param(self, param: Function.Parameter) -> dict: + return { + "name": param.name, + "type": self._serialize_optional(param.type), + "default": self._serialize_optional(param.default), + } + def visit_type_assign(self, stmt: TypeAssign) -> dict: return { "_type": "TypeAssign",