diff --git a/gen/midas.py b/gen/midas.py index 4c4ca32..fed8626 100644 --- a/gen/midas.py +++ b/gen/midas.py @@ -157,4 +157,14 @@ class FunctionType: required: bool +class FrameType: + columns: list[Column] + + @dataclass(frozen=True, kw_only=True) + class Column: + location: Optional[Location] = None + name: Token + type: Type + + ###< diff --git a/gen/python.py b/gen/python.py index df83e6f..73f64c7 100644 --- a/gen/python.py +++ b/gen/python.py @@ -15,7 +15,7 @@ from midas.ast.location import Location ###> MidasType | Type annotations | node class BaseType: base: str - param: Optional[MidasType] + args: tuple[MidasType, ...] class ConstraintType: @@ -174,6 +174,10 @@ class SliceExpr: step: Optional[Expr] +class TupleExpr: + items: tuple[Expr, ...] + + class RawExpr: expr: ast.expr diff --git a/midas/ast/midas.py b/midas/ast/midas.py index a6e2ed9..128b1c0 100644 --- a/midas/ast/midas.py +++ b/midas/ast/midas.py @@ -265,6 +265,9 @@ class Type(ABC): @abstractmethod def visit_function_type(self, type: FunctionType) -> T: ... + @abstractmethod + def visit_frame_type(self, type: FrameType) -> T: ... + @dataclass(frozen=True) class NamedType(Type): @@ -323,3 +326,17 @@ class FunctionType(Type): def accept(self, visitor: Type.Visitor[T]) -> T: return visitor.visit_function_type(self) + + +@dataclass(frozen=True) +class FrameType(Type): + columns: list[Column] + + @dataclass(frozen=True, kw_only=True) + class Column: + location: Optional[Location] = None + name: Token + type: Type + + def accept(self, visitor: Type.Visitor[T]) -> T: + return visitor.visit_frame_type(self) diff --git a/midas/ast/printer.py b/midas/ast/printer.py index 17f0dd0..95c4109 100644 --- a/midas/ast/printer.py +++ b/midas/ast/printer.py @@ -358,6 +358,25 @@ class MidasAstPrinter( arg.type.accept(self) self._write_line(f"required: {arg.required}", last=True) + def visit_frame_type(self, type: m.FrameType) -> None: + self._write_line("FrameType") + with self._child_level(single=True): + self._write_line("columns") + with self._child_level(): + for i, column in enumerate(type.columns): + self._idx = i + if i == len(type.columns) - 1: + self._mark_last() + self._print_frame_column(column) + + def _print_frame_column(self, column: m.FrameType.Column) -> None: + self._write_line("Column") + with self._child_level(): + self._write_line(f'name: "{column.name.lexeme}"') + self._write_line("type") + with self._child_level(single=True): + column.type.accept(self) + class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str]): def __init__(self, indent: int = 4): @@ -513,6 +532,23 @@ class MidasPrinter(m.Expr.Visitor[str], m.Stmt.Visitor[str], m.Type.Visitor[str] res += "?" return res + def visit_frame_type(self, type: m.FrameType) -> str: + res: str = self.indented("Frame[") + if len(type.columns) != 0: + res += "\n" + self.level += 1 + columns: list[str] = [] + for column in type.columns: + columns.append(self.indented(self._print_frame_column(column))) + res += ",\n".join(columns) + self.level -= 1 + res += "\n" + res += "]" + return res + + def _print_frame_column(self, column: m.FrameType.Column) -> str: + return f"{column.name.lexeme}: {column.type.accept(self)}" + class PythonAstPrinter( AstPrinter, @@ -524,7 +560,13 @@ class PythonAstPrinter( self._write_line("BaseType") with self._child_level(): self._write_line(f"base: {node.base}") - self._write_optional_child("param", node.param, last=True) + self._write_line("args:", last=True) + with self._child_level(): + for i, arg in enumerate(node.args): + self._idx = i + if i == len(node.args) - 1: + self._mark_last() + arg.accept(self) def visit_constraint_type(self, node: p.ConstraintType) -> None: self._write_line("ConstraintType") @@ -837,6 +879,17 @@ class PythonAstPrinter( self._write_optional_child("upper", expr.upper) self._write_optional_child("step", expr.step, last=True) + def visit_tuple_expr(self, expr: p.TupleExpr) -> None: + self._write_line("TupleExpr") + with self._child_level(): + self._write_line("items", last=True) + with self._child_level(): + for i, item in enumerate(expr.items): + self._idx = i + if i == len(expr.items) - 1: + self._mark_last() + item.accept(self) + def visit_raw_expr(self, expr: p.RawExpr) -> None: self._write_line("RawExpr") with self._child_level(single=True): diff --git a/midas/ast/python.py b/midas/ast/python.py index 20d7279..6202fcc 100644 --- a/midas/ast/python.py +++ b/midas/ast/python.py @@ -44,7 +44,7 @@ class MidasType(ABC): @dataclass(frozen=True) class BaseType(MidasType): base: str - param: Optional[MidasType] + args: tuple[MidasType, ...] def accept(self, visitor: MidasType.Visitor[T]) -> T: return visitor.visit_base_type(self) @@ -268,6 +268,9 @@ class Expr(ABC): @abstractmethod def visit_slice_expr(self, expr: SliceExpr) -> T: ... + @abstractmethod + def visit_tuple_expr(self, expr: TupleExpr) -> T: ... + @abstractmethod def visit_raw_expr(self, expr: RawExpr) -> T: ... @@ -402,6 +405,14 @@ class SliceExpr(Expr): return visitor.visit_slice_expr(self) +@dataclass(frozen=True) +class TupleExpr(Expr): + items: tuple[Expr, ...] + + def accept(self, visitor: Expr.Visitor[T]) -> T: + return visitor.visit_tuple_expr(self) + + @dataclass(frozen=True) class RawExpr(Expr): expr: ast.expr diff --git a/midas/checker/builtins.midas b/midas/checker/builtins.midas index 3110cbf..e53cab6 100644 --- a/midas/checker/builtins.midas +++ b/midas/checker/builtins.midas @@ -178,4 +178,100 @@ extend dict[K, V] { // def __ior__: fn(value: SupportsKeysAndGetItem[K, V], /) -> dict[K, V] // def __ior__: fn(value: Iterable[tuple[K, V]], /) -> dict[K, V] -} \ No newline at end of file +} + +extend str { + def capitalize: fn() -> str + def casefold: fn() -> str + def center: fn(width: int, fillchar: str?, /) -> str + def count: fn(sub: str, start: None?, end: None?, /) -> int + def count: fn(sub: str, start: int, end: None?, /) -> int + def count: fn(sub: str, start: None, end: int, /) -> int + def count: fn(sub: str, start: int, end: int, /) -> int + def encode: fn(encoding: str?, errors: str?) -> bytes + def endswith: fn(suffix: str, start: None?, end: None?, /) -> bool + def endswith: fn(suffix: str, start: int, end: None?, /) -> bool + def endswith: fn(suffix: str, start: None, end: int, /) -> bool + def endswith: fn(suffix: str, start: int, end: int, /) -> bool + def expandtabs: fn(tabsize: int?) -> str + def find: fn(sub: str, start: None?, end: None?, /) -> int + def find: fn(sub: str, start: int, end: None?, /) -> int + def find: fn(sub: str, start: None, end: int, /) -> int + def find: fn(sub: str, start: int, end: int, /) -> int + // def format: fn(*args: object, **kwargs: object) -> str + // def format_map: fn(mapping: _FormatMapMapping, /) -> str + def index: fn(sub: str, start: None?, end: None?, /) -> int + def index: fn(sub: str, start: int, end: None?, /) -> int + def index: fn(sub: str, start: None, end: int, /) -> int + def index: fn(sub: str, start: int, end: int, /) -> int + def isalnum: fn() -> bool + def isalpha: fn() -> bool + def isascii: fn() -> bool + def isdecimal: fn() -> bool + def isdigit: fn() -> bool + def isidentifier: fn() -> bool + def islower: fn() -> bool + def isnumeric: fn() -> bool + def isprintable: fn() -> bool + def isspace: fn() -> bool + def istitle: fn() -> bool + def isupper: fn() -> bool + def join: fn(iterable: list[str], /) -> str // TODO: use Iterable + def ljust: fn(width: int, fillchar: str?, /) -> str + def lower: fn() -> str + def lstrip: fn(chars: None?, /) -> str + def lstrip: fn(chars: str, /) -> str + def partition: fn(sep: str, /) -> tuple[str, str, str] + + def replace: fn(old: str, new: str, count: int?, /) -> str + + def removeprefix: fn(prefix: str, /) -> str + def removesuffix: fn(suffix: str, /) -> str + def rfind: fn(sub: str, start: None?, end: None?, /) -> int + def rfind: fn(sub: str, start: int, end: None?, /) -> int + def rfind: fn(sub: str, start: None, end: int, /) -> int + def rfind: fn(sub: str, start: int, end: int, /) -> int + def rindex: fn(sub: str, start: None?, end: None?, /) -> int + def rindex: fn(sub: str, start: int, end: None?, /) -> int + def rindex: fn(sub: str, start: None, end: int, /) -> int + def rindex: fn(sub: str, start: int, end: int, /) -> int + def rjust: fn(width: int, fillchar: str?, /) -> str + def rpartition: fn(sep: str, /) -> tuple[str, str, str] + def rsplit: fn(sep: None?, maxsplit: int?) -> list[str] + def rsplit: fn(sep: str, maxsplit: int?) -> list[str] + def rstrip: fn(chars: None?, /) -> str + def rstrip: fn(chars: str, /) -> str + def split: fn(sep: None?, maxsplit: int?) -> list[str] + def split: fn(sep: str, maxsplit: int?) -> list[str] + def splitlines: fn(keepends: bool?) -> list[str] + def startswith: fn(prefix: str, start: None?, end: None?, /) -> bool + def startswith: fn(prefix: str, start: int, end: None?, /) -> bool + def startswith: fn(prefix: str, start: None, end: int, /) -> bool + def startswith: fn(prefix: str, start: int, end: int, /) -> bool + def strip: fn(chars: None?, /) -> str + def strip: fn(chars: str, /) -> str + def swapcase: fn() -> str + def title: fn() -> str + // def translate: fn(table: _TranslateTable, /) -> str + def upper: fn() -> str + def zfill: fn(width: int, /) -> str + def __add__: fn(value: str, /) -> str + // Incompatible with Sequence.__contains__ + def __contains__: fn(key: str, /) -> bool + def __eq__: fn(value: object, /) -> bool + def __ge__: fn(value: str, /) -> bool + def __getitem__: fn(key: slice, /) -> str + def __getitem__: fn(key: int, /) -> str + def __gt__: fn(value: str, /) -> bool + def __hash__: fn() -> int + // def __iter__: fn() -> Iterator[str] + def __le__: fn(value: str, /) -> bool + def __len__: fn() -> int + def __lt__: fn(value: str, /) -> bool + def __mod__: fn(value: Any, /) -> str + def __mul__: fn(value: int, /) -> str + def __ne__: fn(value: object, /) -> bool + def __rmul__: fn(value: int, /) -> str + def __getnewargs__: fn() -> tuple[str] + def __format__: fn(format_spec: str, /) -> str +} diff --git a/midas/checker/builtins.py b/midas/checker/builtins.py index 78ba18f..ce97d9c 100644 --- a/midas/checker/builtins.py +++ b/midas/checker/builtins.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: BUILTIN_SUBTYPES: dict[str, set[str]] = { - "object": {"float", "list", "dict", "str"}, + "object": {"float", "list", "dict", "str", "bytes", "tuple"}, "float": {"int"}, "int": {"bool"}, } @@ -26,12 +26,15 @@ def define_builtins(reg: TypesRegistry): any = reg.define_type("Any", TopType()) unit = reg.define_type("None", UnitType()) object = reg.define_type("object", BaseType(name="object")) + bytes = reg.define_type("bytes", BaseType(name="bytes")) bool = reg.define_type("bool", BaseType(name="bool")) int = reg.define_type("int", BaseType(name="int")) float = reg.define_type("float", BaseType(name="float")) str = reg.define_type("str", BaseType(name="str")) slice = reg.define_type("slice", BaseType(name="slice")) + tuple = reg.define_type("tuple", BaseType(name="tuple")) + list = reg.define_type( "list", GenericType( diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py new file mode 100644 index 0000000..8a7dfbc --- /dev/null +++ b/midas/checker/frame_methods.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Optional + +from midas.ast.location import Location +from midas.checker.registry import TypesRegistry +from midas.checker.reporter import FileReporter +from midas.checker.types import ( + ColumnType, + DataFrameType, + Function, + OverloadedFunction, + TopType, + Type, + UnknownType, + unfold_type, +) + +if TYPE_CHECKING: + from midas.checker.python import PythonTyper, TypedExpr + + +@staticmethod +def frame_method(*names: str): + def wrapper(func): + names_: tuple[str, ...] = names + if len(names_) == 0: + names_ = (func.__name__,) + setattr(func, "__method_names__", names_) + return func + + return wrapper + + +@dataclass(frozen=True, kw_only=True) +class Call: + location: Location + frame: DataFrameType + positional: list[TypedExpr] + keywords: dict[str, TypedExpr] + + +class _MethodRegistryMeta(type): + _methods: dict[str, Callable[..., Type]] = {} + + def __new__( + cls, + name: str, + bases: tuple[type, ...], + namespace: dict[str, Any], + ): + new_class = super().__new__(cls, name, bases, namespace) + new_class._methods = {} + for attr in namespace.values(): + if callable(attr) and hasattr(attr, "__method_names__"): + for name in attr.__method_names__: # type: ignore + new_class._methods[name] = attr # type: ignore + return new_class + + +class MethodRegistry(metaclass=_MethodRegistryMeta): + def __init__(self, typer: PythonTyper) -> None: + self.typer: PythonTyper = typer + + @property + def reporter(self) -> FileReporter: + return self.typer.reporter + + @property + def types(self) -> TypesRegistry: + return self.typer.types + + def call( + self, + method: str, + call: Call, + ) -> Type: + func: Optional[Callable[..., Type]] = self._methods.get(method) + if func is None: + self.reporter.warning(call.location, f"Unknown method {method}") + return UnknownType() + return func(self, call) + + @frame_method("add", "__add__") + def add( + self, + call: Call, + ) -> Type: + # TODO: support add with scalar, sequence, Series, dict + # TODO: check operation exists on inner column types + + new_columns: list[DataFrameType.Column] = [] + + by_name: dict[str, DataFrameType.Column] = {} + frame2: Optional[DataFrameType] = None + if len(call.positional) != 0: + other: Type = call.positional[0][1] + unfolded_other: Type = unfold_type(other) + if isinstance(unfolded_other, DataFrameType): + frame2 = unfolded_other + by_name = { + col.name: col for col in frame2.columns if col.name is not None + } + + in_frame1: set[str] = set() + for column in call.frame.columns: + if column.name is not None: + in_frame1.add(column.name) + + col_type1: Type = column.type + col_type: Type = ColumnType(type=UnknownType()) + if column.name in by_name: + column2 = by_name[column.name] + col_type2: Type = column2.type + if self.types.are_equivalent(col_type2, col_type1): + col_type = col_type1 + + new_column = DataFrameType.Column( + index=column.index, + name=column.name, + type=col_type, + ) + new_columns.append(new_column) + + if frame2 is not None: + for column in frame2.columns: + if column.name in in_frame1: + continue + new_columns.append( + DataFrameType.Column( + index=len(new_columns), + name=column.name, + type=ColumnType(type=UnknownType()), + ) + ) + + signature = Function( + args=[ + Function.Argument( + pos=0, + name="other", + type=DataFrameType(columns=[]), + required=True, + ), + ], + returns=DataFrameType(columns=new_columns), + ) + + return ( + self.typer._get_call_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + or UnknownType() + ) + + @frame_method() + def mean(self, call: Call) -> Type: + with_axis = Function( + kw_args=[ + Function.Argument( + pos=0, + name="axis", + type=self.types.get_type("int"), + required=False, + ) + ], + returns=ColumnType(type=TopType()), + ) + without_axis = Function( + kw_args=[ + Function.Argument( + pos=0, + name="axis", + type=self.types.get_type("None"), + required=True, + ) + ], + returns=TopType(), + ) + overload = OverloadedFunction( + overloads=[ + with_axis, + without_axis, + ] + ) + return ( + self.typer._get_call_result( + location=call.location, + callee=overload, + positional=call.positional, + keywords=call.keywords, + ) + or UnknownType() + ) diff --git a/midas/checker/frames.py b/midas/checker/frames.py new file mode 100644 index 0000000..da9eb0e --- /dev/null +++ b/midas/checker/frames.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, TypeGuard, cast + +import midas.ast.python as p +from midas.ast.location import Location +from midas.checker.frame_methods import Call, MethodRegistry +from midas.checker.reporter import FileReporter +from midas.checker.types import ColumnType, DataFrameType, TupleType, Type, UnknownType + +if TYPE_CHECKING: + from midas.checker.python import PythonTyper, TypedExpr + + +def is_list_of_literals(exprs: list[p.Expr]) -> TypeGuard[list[p.LiteralExpr]]: + return all(isinstance(expr, p.LiteralExpr) for expr in exprs) + + +class FrameManager: + def __init__(self, typer: PythonTyper) -> None: + self.typer: PythonTyper = typer + self.method_resolver: MethodRegistry = MethodRegistry(self.typer) + + def assign( + self, + reporter: FileReporter, + location: Location, + frame: DataFrameType, + index: p.Expr, + value_type: Type, + ) -> Type: + match index: + case p.LiteralExpr(value=str() as name): + return self.assign_column(reporter, location, frame, name, value_type) + + case p.ListExpr(items=indices) if is_list_of_literals(indices) and all( + isinstance(idx, str) for idx in indices + ): + raise NotImplementedError + + case _: + reporter.error(location, f"Invalid index type {index} on {frame}") + return UnknownType() + + def assign_column( + self, + reporter: FileReporter, + location: Location, + frame: DataFrameType, + name: str, + type: Type, + ) -> Type: + if not isinstance(type, ColumnType): + reporter.error( + location, + f"Cannot assign {type} to dataframe column. Must be a ColumnType", + ) + return frame + return self._set_column(frame, name, type) + + def get( + self, + reporter: FileReporter, + location: Location, + frame: DataFrameType, + index: p.Expr, + ) -> Type: + match index: + case p.LiteralExpr(value=str() as name): + column: Optional[ColumnType] = FrameManager._get_column(frame, name) + if column is None: + reporter.error(location, f"Unknown column '{name}' on {frame}") + return UnknownType() + return column + + case p.ListExpr(items=indices) if is_list_of_literals(indices) and all( + isinstance(index.value, str) for index in indices + ): + names: list[str] = [cast(str, index.value) for index in indices] + columns: list[ColumnType] = [] + for name in names: + column: Optional[ColumnType] = FrameManager._get_column(frame, name) + if column is None: + reporter.error(location, f"Unknown column '{name}' on {frame}") + return UnknownType() + columns.append(column) + return TupleType(items=tuple(columns)) + + case _: + reporter.error(location, f"Invalid index type {index} on {frame}") + return UnknownType() + + @classmethod + def _set_column( + cls, frame: DataFrameType, name: str, column: ColumnType + ) -> DataFrameType: + new_columns: list[DataFrameType.Column] = [] + index: int = len(frame.columns) + replace: bool = False + for i, col in enumerate(frame.columns): + if col.name == name: + index = i + replace = True + # TODO: check column type here to prevent changing it + new_columns.append(col) + + new_col: DataFrameType.Column = DataFrameType.Column( + index=index, + name=name, + type=column, + ) + if replace: + new_columns[index] = new_col + else: + new_columns.append(new_col) + + return DataFrameType(columns=new_columns) + + @classmethod + def _set_columns( + cls, frame: DataFrameType, names: list[str], columns: list[ColumnType] + ) -> DataFrameType: + for name, col in zip(names, columns): + frame = cls._set_column(frame, name, col) + return frame + + @classmethod + def _get_column(cls, frame: DataFrameType, name: str) -> Optional[ColumnType]: + for col in frame.columns: + if col.name == name: + return col.type + return None + + @classmethod + def _get_columns( + cls, frame: DataFrameType, names: list[str] + ) -> list[Optional[ColumnType]]: + return [cls._get_column(frame, name) for name in names] + + def call( + self, + method: str, + location: Location, + frame: DataFrameType, + positional: list[TypedExpr], + keywords: dict[str, TypedExpr], + ) -> Type: + call: Call = Call( + location=location, + frame=frame, + positional=positional, + keywords=keywords, + ) + return self.method_resolver.call(method, call) diff --git a/midas/checker/midas.py b/midas/checker/midas.py index c446e6c..01aa09f 100644 --- a/midas/checker/midas.py +++ b/midas/checker/midas.py @@ -13,8 +13,10 @@ from midas.checker.registry import TypesRegistry from midas.checker.reporter import FileReporter, Reporter from midas.checker.types import ( AppliedType, + ColumnType, ComplexType, ConstraintType, + DataFrameType, DerivedType, ExtensionType, Function, @@ -408,6 +410,18 @@ class MidasTyper(m.Stmt.Visitor[None], m.Expr.Visitor[Type], m.Type.Visitor[Type kw=[process_arg(arg, i + n_pos + n_mixed) for i, arg in enumerate(spec.kw)], ) + def visit_frame_type(self, type: m.FrameType) -> Type: + def process_column(i: int, col: m.FrameType.Column) -> DataFrameType.Column: + return DataFrameType.Column( + index=i, + name=col.name.lexeme, + type=ColumnType(type=col.type.accept(self)), + ) + + return DataFrameType( + columns=[process_column(i, col) for i, col in enumerate(type.columns)] + ) + def _resolve_type_params(self, params: list[m.TypeParam]): vars: list[TypeVar] = [] for param in params: diff --git a/midas/checker/preamble.py b/midas/checker/preamble.py index feea8b7..4323762 100644 --- a/midas/checker/preamble.py +++ b/midas/checker/preamble.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import Callable, Optional +from typing import Any, Callable, Optional from midas.checker.environment import Environment from midas.checker.registry import TypesRegistry @@ -17,7 +17,7 @@ class Preamble(Environment): def __init__(self, types: TypesRegistry) -> None: super().__init__() self._types: TypesRegistry = types - self._python_funcs: dict[str, Callable] = {} + self._python_funcs: dict[str, Callable[..., Any]] = {} self._def_type_constructor("object", object) self._def_type_constructor("float", float) @@ -34,7 +34,7 @@ class Preamble(Environment): # TODO: use sink self._def_function( name="print", - pos=[Param("object", TopType())], + pos=[Param("object", TopType(), required=False)], returns=UnitType(), py_function=print, ) @@ -64,11 +64,18 @@ class Preamble(Environment): pos=[Param("prompt", TopType(), required=False)], returns=self._types.get_type("str"), ) + self._def_function( + name="len", + pos=[Param("object", TopType())], + returns=self._types.get_type("int"), + ) def _list_of(self, item_type: Type) -> Type: return self._types.apply_generic(self._types.get_type("list"), [item_type]) - def _def_type_constructor(self, name: str, py_function: Optional[Callable] = None): + def _def_type_constructor( + self, name: str, py_function: Optional[Callable[..., Any]] = None + ): # TODO: more specific arg types self._def_function( name=name, @@ -121,7 +128,7 @@ class Preamble(Environment): kw: list[Param] = [], returns: Type = UnitType(), type_vars: list[TypeVar] = [], - py_function: Optional[Callable] = None, + py_function: Optional[Callable[..., Any]] = None, ): function: Type = self._make_function( name=name, @@ -135,5 +142,5 @@ class Preamble(Environment): if py_function is not None: self._python_funcs[name] = py_function - def get_py_func(self, name: str) -> Optional[Callable]: + def get_py_func(self, name: str) -> Optional[Callable[..., Any]]: return self._python_funcs.get(name) diff --git a/midas/checker/python.py b/midas/checker/python.py index 6a01764..c149d8d 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -8,6 +8,7 @@ from midas.ast.location import Location from midas.ast.printer import MidasPrinter from midas.checker.environment import Environment from midas.checker.evaluator import Evaluator +from midas.checker.frames import FrameManager from midas.checker.operators import ( PY_COMPARATOR_METHODS, PY_OPERATOR_METHODS, @@ -20,11 +21,14 @@ from midas.checker.resolver import Resolver from midas.checker.types import ( AppliedType, BaseType, + ColumnType, ConstraintType, + DataFrameType, DerivedType, Function, GenericType, OverloadedFunction, + TupleType, Type, TypeVar, UnitType, @@ -43,6 +47,10 @@ class ReturnException(Exception): pass +class UndefinedMethodException(Exception): + pass + + @dataclass(frozen=True, kw_only=True) class MappedArgument: expr: p.Expr @@ -71,6 +79,7 @@ class PythonTyper( self.logger: logging.Logger = logging.getLogger("PythonTyper") self.reporter: FileReporter = reporter.for_file(None) self.types: TypesRegistry = types + self.frame_mgr: FrameManager = FrameManager(self) self.global_env: Environment = Preamble(self.types) self.env: Environment = self.global_env self.locals: dict[p.Expr, int] = {} @@ -190,6 +199,36 @@ class PythonTyper( return self.env.get_at(distance, name) return self.global_env.get(name) + def call_method( + self, + location: Location, + obj: Type, + method_name: str, + positional: list[TypedExpr], + keywords: dict[str, TypedExpr], + ) -> Optional[Type]: + unfolded: Type = unfold_type(obj) + match unfolded: + case DataFrameType(): + return self.frame_mgr.call( + method=method_name, + location=location, + frame=unfolded, + positional=positional, + keywords=keywords, + ) + + method: Optional[Type] = self.types.lookup_member(obj, method_name) + if method is None: + raise UndefinedMethodException + + return self._get_call_result( + location, + method, + positional, + keywords, + ) + def is_subtype(self, type1: Type, type2: Type) -> bool: return self.types.is_subtype(type1, type2) @@ -319,9 +358,15 @@ class PythonTyper( case p.VariableExpr(): self._assign_var(location, target, value_type) + # Allow any kind of object because we disallow creating new attributes case p.GetExpr(object=object, name=name): self._assign_attr(location, object, name, value_type) + # Only support variable expressions because modifying + # the underlying value would require reference types + case p.SubscriptExpr(object=p.VariableExpr() as var, index=index): + self._assign_sub(location, var, index, value_type) + case _: if not isinstance(target, p.VariableExpr): self.logger.warning(f"Unsupported assignment to {target}") @@ -360,6 +405,30 @@ class PythonTyper( f"Cannot assign {value_type} to member '{object_type}.{name}' of type {member}", ) + def _assign_sub( + self, + location: Location, + var: p.VariableExpr, + index: p.Expr, + value_type: Type, + ): + var_type: Type = self.type_of(var) + unfolded_type: Type = unfold_type(var_type) + # TODO: what happens if type is an alias of a dataframe type + match unfolded_type: + case DataFrameType() as frame: + new_type: Type = self.frame_mgr.assign( + self.reporter, location, frame, index, value_type + ) + self.env.assign(var.name, new_type) + case UnknownType(): + return + case _: + self.reporter.error( + location, + f"Cannot assign {value_type} to index {index} of {var_type}", + ) + def visit_return_stmt(self, stmt: p.ReturnStmt) -> None: type: Type = self.type_of(stmt.value) if stmt.value is not None else UnitType() self.env.return_types.append(type) @@ -373,8 +442,10 @@ class PythonTyper( # print(m) # <- m is still defined test_type: Type = self.type_of(stmt.test) - # TODO Allow subtypes or any type - if test_type != self.types.get_type("bool"): + if ( + not self.types.is_subtype(test_type, self.types.get_type("bool")) + and test_type != UnknownType() + ): self.reporter.error( stmt.test.location, f"If test must be a boolean, got {test_type}" ) @@ -390,13 +461,16 @@ class PythonTyper( pass def visit_for_stmt(self, stmt: p.ForStmt) -> None: - item_type: Optional[Type] = self._get_iterator_type(stmt.iterator) - if item_type is None: - iterator_type: Type = self.compute_type(stmt.iterator) - self.reporter.error( - stmt.iterator.location, f"{iterator_type} is not iterable" - ) - item_type = UnknownType() + item_type: Type = UnknownType() + iterator_type: Type = self.type_of(stmt.iterator) + if iterator_type != UnknownType(): + maybe_item_type = self._get_iterator_type(stmt.iterator, iterator_type) + if maybe_item_type is None: + self.reporter.error( + stmt.iterator.location, f"{iterator_type} is not iterable" + ) + else: + item_type = maybe_item_type self._assign(stmt.location, stmt.target, item_type) self.judge(stmt.target, item_type) @@ -436,20 +510,16 @@ class PythonTyper( left: Type = self.type_of(left_expr) right: Type = self.type_of(right_expr) - operation: Optional[Type] = self.types.lookup_member(left, method) - if operation is None: + result: Optional[Type] + try: + result = self.call_method(location, left, method, [(right_expr, right)], {}) + except UndefinedMethodException: self.reporter.error( location, f"Undefined operation {method} between {left} and {right}", ) return UnknownType() - result: Optional[Type] = self._get_call_result( - location, - operation, - [(right_expr, right)], - {}, - ) return result or UnknownType() def visit_unary_expr(self, expr: p.UnaryExpr) -> Type: @@ -462,20 +532,17 @@ class PythonTyper( return UnknownType() operand: Type = self.type_of(expr.right) - operation: Optional[Type] = self.types.lookup_member(operand, method) - if operation is None: + + result: Optional[Type] + try: + result = self.call_method(expr.location, operand, method, [], {}) + except UndefinedMethodException: self.reporter.error( expr.location, f"Undefined operation {method} for {operand}", ) return UnknownType() - result: Optional[Type] = self._get_call_result( - expr.location, - operation, - [], - {}, - ) return result or UnknownType() def visit_call_expr(self, expr: p.CallExpr) -> Type: @@ -483,13 +550,27 @@ class PythonTyper( case p.VariableExpr(name="TypeVar"): return self.define_typevar(expr) or UnknownType() - callee: Type = self.type_of(expr.callee) positional: list[TypedExpr] = [ (arg, self.type_of(arg)) for arg in expr.arguments ] keywords: dict[str, TypedExpr] = { name: (arg, self.type_of(arg)) for name, arg in expr.keywords.items() } + + match expr.callee: + case p.GetExpr(object=obj, name=method): + obj_type: Type = self.type_of(obj) + unfolded: Type = unfold_type(obj_type) + if isinstance(unfolded, DataFrameType): + return self.frame_mgr.call( + method, + expr.location, + unfolded, + positional, + keywords, + ) + + callee: Type = self.type_of(expr.callee) return ( self._get_call_result( location=expr.location, @@ -504,7 +585,7 @@ class PythonTyper( object: Type = self.type_of(expr.object) member: Optional[Type] = self.types.lookup_member(object, expr.name) if member is None: - self.reporter.error( + self.reporter.warning( expr.location, f"Unknown member '{expr.name}' of {object}" ) return UnknownType() @@ -521,6 +602,8 @@ class PythonTyper( return self.types.get_type("float") case str(): return self.types.get_type("str") + case None: + return self.types.get_type("None") case _: self.reporter.warning(expr.location, f"Unknown literal {expr}") return UnknownType() @@ -563,7 +646,10 @@ class PythonTyper( test_type: Type = self.type_of(expr.test) # TODO Allow subtypes or any type - if test_type != self.types.get_type("bool"): + if ( + not self.is_subtype(test_type, self.types.get_type("bool")) + and test_type != UnknownType() + ): self.reporter.error( expr.test.location, f"If test must be a boolean, got {test_type}" ) @@ -592,9 +678,9 @@ class PythonTyper( if len(item_types) == 1: item_type: Type = item_types[0] return self.types.apply_generic(list_type, [item_type]) - self.reporter.error( + self.reporter.warning( expr.location, - f"Heterogeneous list items: {item_types}", + f"Heterogeneous list items: [{', '.join(map(str, item_types))}]", ) return self.types.apply_generic(list_type, [UnknownType()]) @@ -624,22 +710,29 @@ class PythonTyper( if len(key_types) == 1: key_type = key_types[0] else: - self.reporter.error( + self.reporter.warning( expr.location, - f"Heterogeneous dict keys: {key_types}", + f"Heterogeneous dict keys: [{', '.join(map(str, key_types))}]", ) if len(value_types) == 1: value_type = value_types[0] else: - self.reporter.error( + self.reporter.warning( expr.location, - f"Heterogeneous dict values: {value_types}", + f"Heterogeneous dict values: [{', '.join(map(str, value_types))}]", ) return self.types.apply_generic(dict_type, [key_type, value_type]) def visit_subscript_expr(self, expr: p.SubscriptExpr) -> Type: object: Type = self.type_of(expr.object) + unfolded: Type = unfold_type(object) + match unfolded: + case TupleType(): + return self._visit_tuple_subscript(unfolded, expr) + case DataFrameType(): + return self._visit_frame_subscript(unfolded, expr) + operation: Optional[Type] = self.types.lookup_member(object, "__getitem__") if operation is None: self.reporter.error( @@ -657,6 +750,11 @@ class PythonTyper( def visit_slice_expr(self, expr: p.SliceExpr) -> Type: return self.types.get_type("slice") + def visit_tuple_expr(self, expr: p.TupleExpr) -> Type: + return TupleType( + items=tuple(self.type_of(item) for item in expr.items), + ) + def visit_raw_expr(self, expr: p.RawExpr) -> Type: return UnknownType() @@ -668,22 +766,35 @@ class PythonTyper( self.reporter.warning(node.location, f"Unknown type '{node.base}'") return UnknownType() - if node.param is not None: - param: Type = self.resolve_type_expr(node.param) - return self.types.apply_generic(base, [param]) + if len(node.args) != 0: + args: list[Type] = [self.resolve_type_expr(arg) for arg in node.args] + return self.types.apply_generic(base, args) return base def visit_constraint_type(self, node: p.ConstraintType) -> Type: self.reporter.warning(node.location, "ConstraintType not yet supported") return UnknownType() - def visit_frame_column(self, node: p.FrameColumn) -> Type: - self.reporter.warning(node.location, "FrameColumn not yet supported") - return UnknownType() + def visit_frame_column(self, node: p.FrameColumn) -> ColumnType: + return ColumnType( + type=( + self.resolve_type_expr(node.type) + if node.type is not None + else UnknownType() + ) + ) def visit_frame_type(self, node: p.FrameType) -> Type: - self.reporter.warning(node.location, "FrameType not yet supported") - return UnknownType() + return DataFrameType( + columns=[ + DataFrameType.Column( + index=i, + name=column.name, + type=self.visit_frame_column(column), + ) + for i, column in enumerate(node.columns) + ] + ) def _get_call_result( self, @@ -1055,9 +1166,8 @@ class PythonTyper( return False return True - def _get_iterator_type(self, expr: p.Expr) -> Optional[Type]: + def _get_iterator_type(self, expr: p.Expr, type: Type) -> Optional[Type]: # TODO: lookup __iter__ - type: Type = self.type_of(expr) getitem: Optional[Type] = self.types.lookup_member(type, "__getitem__") if getitem is None: return None @@ -1123,7 +1233,7 @@ class PythonTyper( node: ast.Expression = ast.parse(value, mode="eval") return parser._parse_type(node.body) case p.VariableExpr(name=name): - return p.BaseType(location=location, base=name, param=None) + return p.BaseType(location=location, base=name, args=()) case _: raise NotImplementedError @@ -1211,8 +1321,34 @@ class PythonTyper( return False return True + case DataFrameType() | ColumnType(): + self.reporter.error( + expr.location, f"Cannot cast {lit_value!r} to {target_type}" + ) + return False + case _: self.reporter.info( expr.location, f"Cannot evaluate cast to {target_type} statically" ) return False + + def _visit_tuple_subscript(self, tup: TupleType, expr: p.SubscriptExpr) -> Type: + match expr.index: + case p.LiteralExpr(value=int() as index): + if index < 0 or index >= len(tup.items): + self.reporter.error( + expr.location, f"Index {index} out of range for tuple {tup}" + ) + return UnknownType() + return tup.items[index] + case _: + self.reporter.error( + expr.location, f"Invalid index type {expr.index} on {tup}" + ) + return UnknownType() + + def _visit_frame_subscript( + self, frame: DataFrameType, expr: p.SubscriptExpr + ) -> Type: + return self.frame_mgr.get(self.reporter, expr.location, frame, expr.index) diff --git a/midas/checker/registry.py b/midas/checker/registry.py index 7520981..02c96c3 100644 --- a/midas/checker/registry.py +++ b/midas/checker/registry.py @@ -7,8 +7,10 @@ from midas.checker.builtins import BUILTIN_SUBTYPES from midas.checker.types import ( AppliedType, BaseType, + ColumnType, ComplexType, ConstraintType, + DataFrameType, DerivedType, ExtensionType, Function, @@ -16,6 +18,7 @@ from midas.checker.types import ( OverloadedFunction, Predicate, TopType, + TupleType, Type, TypeVar, UnknownType, @@ -157,6 +160,24 @@ class TypesRegistry: return False return True + case (DataFrameType(columns=columns1), DataFrameType(columns=columns2)): + # TODO: check order? + by_name1: dict[str, DataFrameType.Column] = { + col.name: col for col in columns1 if col.name is not None + } + for col2 in columns2: + if col2.name not in by_name1: + return False + if not self.is_subtype(by_name1[col2.name].type, col2.type): + return False + return True + + case (ColumnType(type=inner1), ColumnType(type=inner2)): + # TODO: invariant, replace ColumnType with simple GenericType + if not self.are_equivalent(inner1, inner2): + return False + return True + case (Function(), Function()): return self.is_func_subtype(type1, type2) @@ -187,6 +208,9 @@ class TypesRegistry: return False + def are_equivalent(self, type1: Type, type2: Type) -> bool: + return self.is_subtype(type1, type2) and self.is_subtype(type2, type1) + # TODO: verify the logic in here def is_func_subtype(self, func1: Function, func2: Function) -> bool: """Check whether a function is a subtype of another @@ -323,6 +347,9 @@ class TypesRegistry: body=substitute_typevars(body, substitutions), ) + case BaseType(name="tuple"): + return TupleType(items=tuple(args)) + case _: raise ValueError(f"{type} is not a generic type") diff --git a/midas/checker/resolver.py b/midas/checker/resolver.py index 3bf73d7..43768bd 100644 --- a/midas/checker/resolver.py +++ b/midas/checker/resolver.py @@ -128,6 +128,10 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]): case p.GetExpr(): target.accept(self) + + case p.SubscriptExpr(): + target.accept(self) + case _: raise Exception(f"Unsupported assignment to {target}") @@ -232,5 +236,9 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]): if expr.step is not None: self.resolve(expr.step) + def visit_tuple_expr(self, expr: p.TupleExpr) -> None: + for item in expr.items: + self.resolve(item) + def visit_raw_expr(self, expr: p.RawExpr) -> None: pass diff --git a/midas/checker/types.py b/midas/checker/types.py index bf1703b..a4279af 100644 --- a/midas/checker/types.py +++ b/midas/checker/types.py @@ -2,7 +2,7 @@ from __future__ import annotations from dataclasses import dataclass, field from enum import StrEnum -from typing import Optional, assert_never +from typing import Optional, assert_never, cast import midas.ast.midas as m from midas.ast.printer import MidasPrinter @@ -156,6 +156,37 @@ class ConstraintType: return f"{self.type} where {printer.print(self.constraint)}" +@dataclass(frozen=True, kw_only=True) +class TupleType: + items: tuple[Type, ...] + + def __str__(self) -> str: + return f"({', '.join(map(str, self.items))})" + + +@dataclass(frozen=True, kw_only=True) +class ColumnType: + type: Type + + def __str__(self) -> str: + return f"Column[{self.type}]" + + +@dataclass(frozen=True, kw_only=True) +class DataFrameType: + columns: list[Column] + + def __str__(self) -> str: + schema: list[str] = [f"{col.name}: {col.type}" for col in self.columns] + return f"Frame[{', '.join(schema)}]" + + @dataclass(frozen=True, kw_only=True) + class Column: + index: int + name: Optional[str] + type: ColumnType + + def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: def sub_argument(arg: Function.Argument): return Function.Argument( @@ -165,6 +196,13 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: required=arg.required, ) + def sub_column(col: DataFrameType.Column): + return DataFrameType.Column( + index=col.index, + name=col.name, + type=cast(ColumnType, substitute_typevars(col.type, substitutions)), + ) + match type: case TopType(): return type @@ -252,10 +290,26 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: body=substitute_typevars(body, substitutions), ) + case TupleType(items=items): + return TupleType( + items=tuple(substitute_typevars(item, substitutions) for item in items), + ) + + case ColumnType(type=items_type): + return ColumnType( + type=substitute_typevars(items_type, substitutions), + ) + + case DataFrameType(columns=columns): + return DataFrameType( + columns=list(map(sub_column, columns)), + ) + case UnknownType() | UnitType(): return type case TopType() | GenericType(): + raise NotImplementedError(f"Unsupported type {type}") # Ensure exhaustiveness @@ -319,6 +373,15 @@ def to_annotation(type: Type) -> str: case ConstraintType(): return str(type) + case TupleType(items=items): + return f"Tuple[{', '.join(map(to_annotation, items))}]" + + case ColumnType(): + return "pd.Series" + + case DataFrameType(): + return "pd.DataFrame" + case _: assert_never(type) @@ -344,4 +407,7 @@ Type = ( | GenericType | AppliedType | ConstraintType + | TupleType + | ColumnType + | DataFrameType ) diff --git a/midas/cli/commands/compile.py b/midas/cli/commands/compile.py index 497a0f9..a59eea1 100644 --- a/midas/cli/commands/compile.py +++ b/midas/cli/commands/compile.py @@ -5,7 +5,7 @@ import sys from pathlib import Path -from typing import TextIO +from typing import Optional, TextIO import click @@ -19,18 +19,23 @@ from midas.utils import TypedAST @click.command(help="Compile source") @click.argument("file", type=click.File("r")) @click.option("-t", "--types", type=click.File("r"), multiple=True) +@click.option("-s", "--stubs", type=str, multiple=True) @click.option("--ignore-errors", is_flag=True) def compile( file: TextIO, types: tuple[TextIO], + stubs: tuple[str], ignore_errors: bool, ): source: str = file.read() source_path: Path = Path(file.name).resolve() checker = TypeChecker() - for types_file in types: - checker.import_midas(Path(types_file.name).resolve()) + type_files: list[tuple[Path, Optional[str]]] = [] + for i, types_file in enumerate(types): + in_path: Path = Path(types_file.name).resolve() + checker.import_midas(in_path) + type_files.append((in_path, stubs[i] if i < len(stubs) else None)) typed_ast: TypedAST = checker.type_check_source(source, str(source_path)) diagnostics: list[Diagnostic] = checker.diagnostics.copy() @@ -43,4 +48,4 @@ def compile( sys.exit(1) generator = Generator(workdir=source_path.parent, types=checker.types) - generator.generate(typed_ast, source_path) + generator.generate(typed_ast, source_path, type_files=type_files) diff --git a/midas/cli/commands/stubs.py b/midas/cli/commands/stubs.py index a5267c9..9eb028c 100644 --- a/midas/cli/commands/stubs.py +++ b/midas/cli/commands/stubs.py @@ -1,7 +1,7 @@ import ast import time from pathlib import Path -from typing import TextIO +from typing import Optional, TextIO import black import click @@ -38,15 +38,17 @@ class Handler(FileSystemEventHandler): @click.command(help="Generate stubs from Midas definitions") @click.argument("file", type=click.File("r")) -@click.option("-o", "--output", type=click.File("w"), default="-") +@click.option("-o", "--output", type=click.File("w")) @click.option("-w", "--watch", is_flag=True) def stubs( file: TextIO, - output: TextIO, + output: Optional[TextIO], watch: bool, ): source_path: Path = Path(file.name).resolve() - out_path: Path = Path(output.name).resolve() + out_path: Path = source_path.with_suffix(".pyi") + if output is not None: + out_path = Path(output.name).resolve() generate_stubs(source_path, out_path) if watch: diff --git a/midas/cli/highlighter.py b/midas/cli/highlighter.py index ce6c2b2..8d142f1 100644 --- a/midas/cli/highlighter.py +++ b/midas/cli/highlighter.py @@ -134,9 +134,9 @@ class PythonHighlighter( def visit_base_type(self, node: p.BaseType) -> None: self.wrap(node, "base-type") - if node.param is not None: - self.wrap(node.param, "param") - node.param.accept(self) + for arg in node.args: + self.wrap(arg, "arg") + arg.accept(self) def visit_constraint_type(self, node: p.ConstraintType) -> None: self.wrap(node, "constraint-type") @@ -247,6 +247,10 @@ class PythonHighlighter( if expr.step is not None: expr.step.accept(self) + def visit_tuple_expr(self, expr: p.TupleExpr) -> None: + for item in expr.items: + item.accept(self) + def visit_raw_expr(self, expr: p.RawExpr) -> None: ... def visit_raw_stmt(self, stmt: p.RawStmt) -> None: ... @@ -350,6 +354,14 @@ class MidasHighlighter( for param in spec.pos + spec.mixed + spec.kw: param.type.accept(self) + def visit_frame_type(self, type: m.FrameType) -> None: + self.wrap(type, "frame") + for column in type.columns: + self._visit_frame_column(column) + + def _visit_frame_column(self, column: m.FrameType.Column) -> None: + self.wrap(column, "column") + class DiagnosticsHighlighter(Highlighter): EXTRA_CSS_PATH: Optional[Path] = Path(__file__).parent / "hl_diagnostic.css" diff --git a/midas/cli/hl_python.css b/midas/cli/hl_python.css index e6dc43b..899c461 100644 --- a/midas/cli/hl_python.css +++ b/midas/cli/hl_python.css @@ -3,7 +3,7 @@ span { --col: 108, 233, 108; } - &.param { + &.arg { --col: 103, 192, 224; } diff --git a/midas/cli/utils.py b/midas/cli/utils.py index 38ceaee..9dc3802 100644 --- a/midas/cli/utils.py +++ b/midas/cli/utils.py @@ -68,7 +68,7 @@ class DiagnosticPrinter: loc: Location = diagnostic.location if loc.lineno != loc.end_lineno: - print(diagnostic) + self.print_multiline(lines, diagnostic, indent) return start_offset: int = loc.col_offset @@ -95,3 +95,27 @@ class DiagnosticPrinter: print(indent_str + before + subject + after) print(indent_str + cursor) print() + + def print_multiline( + self, all_lines: list[str], diagnostic: Diagnostic, indent: int = 4 + ): + loc: Location = diagnostic.location + lines: list[str] = all_lines[loc.lineno - 1 : loc.end_lineno] + + start_offset: int = loc.col_offset + end_offset: int = loc.end_col_offset or (start_offset + 1) + + indent_str: str = " " * indent + color: int = self.COLORS.get(diagnostic.type, Ansi.WHITE) + res: str = indent_str + lines[0][:start_offset] + res += Ansi.FG(color) + lines[0][start_offset:] + for line in lines[1:-1]: + res += "\n" + indent_str + line + res += "\n" + indent_str + lines[-1][:end_offset] + res += Ansi.RESET + lines[-1][end_offset:] + + print(diagnostic.location_str + ":") + print(res) + print() + print(Ansi.FG(color) + diagnostic.message + Ansi.RESET) + print() diff --git a/midas/generator/generator.py b/midas/generator/generator.py index 3f87ef2..c2f46ee 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -1,4 +1,5 @@ import ast +import logging import shutil from dataclasses import dataclass, field from pathlib import Path @@ -8,38 +9,47 @@ import midas.ast.midas as m import midas.ast.python as p from midas.ast.location import Location from midas.ast.printer import MidasPrinter +from midas.checker.checker import TypeChecker from midas.checker.registry import TypesRegistry from midas.checker.types import ( AppliedType, BaseType, + ColumnType, ComplexType, ConstraintType, + DataFrameType, DerivedType, ExtensionType, Function, GenericType, OverloadedFunction, TopType, + TupleType, Type, TypeVar, UnitType, UnknownType, ) from midas.generator.constraints import ConstraintGenerator +from midas.generator.stubs import StubsGenerator from midas.utils import TypedAST @dataclass class Scope: - pre_assertions: list[ast.stmt] = field(default_factory=list) - aliases: list[str] = field(default_factory=list) + pre_assertions: list[ast.stmt] = field(default_factory=list[ast.stmt]) + aliases: list[str] = field(default_factory=list[str]) class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): + IS_DATAFRAME_FUNC = "__midas_is_dataframe__" + IS_COLUMN_FUNC = "__midas_is_column__" + def __init__(self, workdir: Path, types: TypesRegistry) -> None: self.workdir: Path = workdir.resolve() self.build_dir: Path = self.workdir / "build" / "midas" self.rel_src_path: Path = Path() + self.logger: logging.Logger = logging.getLogger("Generator") self._typed_ast: TypedAST = TypedAST( stmts=[], @@ -53,20 +63,37 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): self._constraint_generator: ConstraintGenerator = ConstraintGenerator(types) self._constraints: list[tuple[m.Expr, ast.expr]] = [] - def generate_ast(self, typed_ast: TypedAST, src_path: Path) -> ast.AST: - self.rel_src_path = src_path.resolve().relative_to(self.workdir) + self.define_is_dataframe: bool = False + self.define_is_column: bool = False + + def set_src_path(self, path: Path): + self.rel_src_path = path.resolve().relative_to(self.workdir) + + def generate_ast(self, typed_ast: TypedAST) -> ast.AST: self._typed_ast = typed_ast body: list[ast.stmt] = self._visit_body(typed_ast.stmts) predicates: list[ast.stmt] = self._constraint_generator.get_definitions() - module = ast.Module(body=predicates + body, type_ignores=[]) + + body = predicates + body + + if self.define_is_dataframe: + body = [self._is_dataframe_definition()] + body + + if self.define_is_column: + body = [self._is_column_definition()] + body + + module = ast.Module(body=body, type_ignores=[]) module = ast.fix_missing_locations(module) return module def generate( - self, typed_ast: TypedAST, src_path: Path, out_path: Optional[Path] = None + self, + typed_ast: TypedAST, + src_path: Path, + out_path: Optional[Path] = None, + type_files: Optional[list[tuple[Path, Optional[str]]]] = None, ) -> Path: - module: ast.AST = self.generate_ast(typed_ast, src_path) - compiled: str = ast.unparse(module) + self.set_src_path(src_path) if out_path is None: if self.build_dir.exists(): shutil.rmtree(self.build_dir) @@ -78,10 +105,30 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): raise ValueError( f"Directory traversal, {self.rel_src_path} points outside of parent directory" ) - out_path.parent.mkdir(parents=True, exist_ok=True) + out_dir: Path = out_path.parent + out_dir.parent.mkdir(parents=True, exist_ok=True) + + if type_files is not None: + for in_path, out_name in type_files: + if out_name is None: + out_name = in_path.stem + self.generate_stubs(in_path, out_dir / f"{out_name}.py") + + module: ast.AST = self.generate_ast(typed_ast) + compiled: str = ast.unparse(module) + out_path.write_text(compiled) return out_path + def generate_stubs(self, in_path: Path, out_path: Path): + checker = TypeChecker() + checker.import_midas(in_path) + generator = StubsGenerator(checker.types) + module: ast.Module = generator.generate_stubs() + module = ast.fix_missing_locations(module) + output: str = ast.unparse(module) + out_path.write_text(output) + def visit_binary_expr(self, expr: p.BinaryExpr) -> ast.expr: return ast.BinOp( left=expr.left.accept(self), @@ -139,7 +186,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): alias: ast.expr = self._make_alias(expr2) type: Type = self._get_expr_type(expr) - self._make_cast_asserts(expr.location, alias, type) + asserts: list[ast.stmt] = self._make_cast_asserts(expr.location, alias, type) + for assert_ in asserts: + self._add_assert(assert_) return alias @@ -174,6 +223,11 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): step=expr.step.accept(self) if expr.step is not None else None, ) + def visit_tuple_expr(self, expr: p.TupleExpr) -> ast.expr: + return ast.Tuple( + elts=[item.accept(self) for item in expr.items], + ) + def visit_raw_expr(self, expr: p.RawExpr) -> ast.expr: return expr.expr @@ -274,63 +328,156 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): ) return alias - def _add_assert(self, expr: ast.expr, message: str | ast.expr): + def _build_assert(self, expr: ast.expr, message: str | ast.expr) -> ast.stmt: if isinstance(message, str): message = ast.Constant(value=message) - self._scopes[-1].pre_assertions.append( - ast.Assert( - test=expr, - msg=message, - ) + return ast.Assert( + test=expr, + msg=message, ) + def _add_assert(self, assertion: ast.stmt): + self._scopes[-1].pre_assertions.append(assertion) + def _get_expr_type(self, query: p.Expr) -> Type: for expr, type in self._typed_ast.judgements: if expr == query: return type raise RuntimeError(f"Cannot get type judgement for {query}") - def _make_cast_asserts(self, src_location: Location, expr: ast.expr, type: Type): + def _make_cast_asserts( + self, src_location: Location, expr: ast.expr, type: Type + ) -> list[ast.stmt]: match type: case UnknownType(): - pass + return [] case BaseType(name=name): - self._add_assert( - ast.Call( - func=ast.Name(id="isinstance"), - args=[expr, ast.Name(id=name)], - keywords=[], - ), - self._make_cast_assert_message(src_location, expr, type), - ) + return [ + self._build_assert( + ast.Call( + func=ast.Name(id="isinstance"), + args=[expr, ast.Name(id=name)], + keywords=[], + ), + self._make_cast_assert_message(src_location, expr, type), + ) + ] case DerivedType(type=base): - self._make_cast_asserts(src_location, expr, base) + return self._make_cast_asserts(src_location, expr, base) case UnitType(): - self._add_assert( - ast.Compare( - left=expr, - ops=[ast.Is()], - comparators=[ - ast.Constant(value=None), - ], + return [ + self._build_assert( + ast.Compare( + left=expr, + ops=[ast.Is()], + comparators=[ + ast.Constant(value=None), + ], + ), + self._make_cast_assert_message(src_location, expr, type), ), - self._make_cast_assert_message(src_location, expr, type), - ) + ] case AppliedType(body=body): - self._make_cast_asserts(src_location, expr, body) + return self._make_cast_asserts(src_location, expr, body) case ConstraintType(type=base, constraint=constraint): - self._make_cast_asserts(src_location, expr, base) - self._make_constraint_assert(src_location, expr, constraint) + asserts: list[ast.stmt] = self._make_cast_asserts( + src_location, expr, base + ) + asserts.append( + self._make_constraint_assert(src_location, expr, constraint) + ) + return asserts case TypeVar(bound=bound): # TODO: check with type from arguments / use call-site context - if bound is not None: - self._make_cast_asserts(src_location, expr, bound) + if bound is None: + return [] + return self._make_cast_asserts(src_location, expr, bound) + + case TupleType(items=items): + asserts: list[ast.stmt] = [ + self._build_assert( + ast.Call( + func=ast.Name(id="isinstance"), + args=[expr, ast.Name(id="tuple")], + keywords=[], + ), + self._make_cast_assert_message(src_location, expr, type), + ), + ] + assert isinstance(expr, ast.Tuple) + for item, item_type in zip(expr.elts, items): + asserts.extend( + self._make_cast_asserts(src_location, item, item_type) + ) + return asserts + + case DataFrameType(columns=columns): + self.define_is_dataframe = True + asserts: list[ast.stmt] = [ + self._build_assert( + ast.Call( + func=ast.Name(id=self.IS_DATAFRAME_FUNC), + args=[expr], + keywords=[], + ), + self._make_cast_assert_message( + src_location, expr, type, ": Not a dataframe" + ), + ), + ] + for column in columns: + asserts.append( + self._build_assert( + ast.Compare( + left=ast.Constant(value=column.name), + ops=[ast.In()], + comparators=[expr], + ), + self._make_cast_assert_message( + src_location, + expr, + type, + f": Missing column {column.name}", + ), + ) + ) + asserts.extend( + self._make_cast_asserts( + src_location, + ast.Subscript( + value=expr, slice=ast.Constant(value=column.name) + ), + column.type, + ) + ) + return asserts + + case ColumnType(): + self.define_is_column = True + asserts: list[ast.stmt] = [ + self._build_assert( + ast.Call( + func=ast.Name(id=self.IS_COLUMN_FUNC), + args=[expr], + keywords=[], + ), + self._make_cast_assert_message( + src_location, expr, type, ": Not a column" + ), + ), + ] + inner_assert: Optional[ast.stmt] = self._make_column_inner_assert( + src_location, expr, type + ) + if inner_assert is not None: + asserts.append(inner_assert) + return asserts case ( TopType() @@ -340,14 +487,19 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): | ExtensionType() | GenericType() ): - raise NotImplementedError(f"Can't make assertion for type {type}") + self.logger.warning(f"Can't make assertion for type {type}") + return [] # Ensure exhaustiveness case _: assert_never(type) def _make_cast_assert_message( - self, location: Location, expr: ast.expr, type: Type + self, + location: Location, + expr: ast.expr, + type: Type, + extra: Optional[str] = None, ) -> ast.expr: loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}" # f"file.py:L1:1: CastError: Cannot cast {type(expr).__name__} to Type" @@ -365,15 +517,15 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): ), conversion=-1, ), - ast.Constant(f" to {type}"), + ast.Constant(f" to {type}{extra or ''}"), ] ) def _make_constraint_assert( self, src_location: Location, expr: ast.expr, constraint: m.Expr - ): + ) -> ast.stmt: test_func: ast.expr = self._get_constraint(constraint) - self._add_assert( + return self._build_assert( ast.Call( func=test_func, args=[expr], @@ -401,3 +553,90 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): constraint: ast.expr = self._constraint_generator.generate(expr) self._constraints.append((expr, constraint)) return constraint + + def _is_dataframe_definition(self) -> ast.stmt: + """ + def IS_DATAFRAME_FUNC(obj) -> bool: + import pandas as pd + return isinstance(obj, pd.DataFrame) + """ + + return ast.FunctionDef( + name=self.IS_DATAFRAME_FUNC, + args=ast.arguments( + posonlyargs=[ast.arg(arg="obj")], + args=[], + kwonlyargs=[], + defaults=[], + kw_defaults=[], + ), + body=[ + ast.Import(names=[ast.alias(name="pandas", asname="pd")]), + ast.Return( + value=ast.Call( + func=ast.Name(id="isinstance"), + args=[ + ast.Name(id="obj"), + ast.Attribute( + value=ast.Name(id="pd"), + attr="DataFrame", + ), + ], + keywords=[], + ) + ), + ], + decorator_list=[], + returns=ast.Name(id="bool"), + ) + + def _is_column_definition(self) -> ast.stmt: + """ + def IS_COLUMN_FUNC(obj) -> bool: + import pandas as pd + return isinstance(obj, pd.Series) + """ + + return ast.FunctionDef( + name=self.IS_COLUMN_FUNC, + args=ast.arguments( + posonlyargs=[ast.arg(arg="obj")], + args=[], + kwonlyargs=[], + defaults=[], + kw_defaults=[], + ), + body=[ + ast.Import(names=[ast.alias(name="pandas", asname="pd")]), + ast.Return( + value=ast.Call( + func=ast.Name(id="isinstance"), + args=[ + ast.Name(id="obj"), + ast.Attribute( + value=ast.Name(id="pd"), + attr="Series", + ), + ], + keywords=[], + ) + ), + ], + decorator_list=[], + returns=ast.Name(id="bool"), + ) + + def _make_column_inner_assert( + self, src_location: Location, column: ast.expr, type: ColumnType + ) -> Optional[ast.stmt]: + # TODO: improve message, maybe chain contexts + col: ast.expr = ast.Name(id="col") + body: list[ast.stmt] = self._make_cast_asserts(src_location, col, type.type) + if len(body) == 0: + return None + return ast.For( + target=col, + iter=column, + body=body, + orelse=[], + ) diff --git a/midas/generator/stubs.py b/midas/generator/stubs.py index 2ebcbd2..d0a9b23 100644 --- a/midas/generator/stubs.py +++ b/midas/generator/stubs.py @@ -6,14 +6,17 @@ from midas.checker.registry import Member, TypesRegistry from midas.checker.types import ( AppliedType, BaseType, + ColumnType, ComplexType, ConstraintType, + DataFrameType, DerivedType, ExtensionType, Function, GenericType, OverloadedFunction, TopType, + TupleType, Type, TypeVar, UnitType, @@ -30,6 +33,7 @@ class StubsGenerator: self.types: TypesRegistry = types self.stubs: list[ast.stmt] = [] self.typing_imports: set[str] = set() + self.import_pandas: bool = False self.protocol_idx: int = 0 self.stub_idx: int = 0 self.type_var_idx: int = 0 @@ -38,6 +42,7 @@ class StubsGenerator: def generate_stubs(self) -> ast.Module: self.stubs = [] self.typing_imports = set() + self.import_pandas = False for name, type in self.types._types.items(): # Skip builtin types, not just based on name so the user can override # TODO: check if added members on builtin type @@ -53,7 +58,7 @@ class StubsGenerator: continue self.generate_stub(name, type) - imports = [ + imports: list[ast.stmt] = [ ast.ImportFrom( module="__future__", names=[ast.alias(name="annotations")], @@ -70,6 +75,17 @@ class StubsGenerator: level=0, ) ) + if self.import_pandas: + imports.append( + ast.Import( + names=[ + ast.alias( + name="pandas", + asname="pd", + ) + ], + ) + ) return ast.Module(body=imports + self.stubs, type_ignores=[]) def generate_stub(self, name: str, type: Type): @@ -231,6 +247,31 @@ class StubsGenerator: case ConstraintType(): return self.dump_type(type.type) + case TupleType(items=items): + return ast.Subscript( + value=ast.Name(id="tuple"), + slice=ast.Tuple( + elts=[self.dump_type(item) for item in items], + ), + ) + + case ColumnType(type=inner): + self.import_pandas = True + return ast.Subscript( + value=ast.Attribute( + value=ast.Name(id="pd"), + attr="Series", + ), + slice=self.dump_type(inner), + ) + + case DataFrameType(): + self.import_pandas = True + return ast.Attribute( + value=ast.Name(id="pd"), + attr="DataFrame", + ) + case _: assert_never(type) diff --git a/midas/parser/midas.py b/midas/parser/midas.py index c5fced1..3bcee32 100644 --- a/midas/parser/midas.py +++ b/midas/parser/midas.py @@ -10,6 +10,7 @@ from midas.ast.midas import ( Expr, ExtendStmt, ExtensionType, + FrameType, FunctionType, GenericType, GetExpr, @@ -226,8 +227,10 @@ class MidasParser(Parser): return self.generic_type() def generic_type(self) -> Type: - type: Type = self.named_type() + type: NamedType = self.named_type() if self.check(TokenType.LEFT_BRACKET): + if type.name.lexeme == "Frame": + return self.frame_type() args: list[Type] = self.type_args() return GenericType( location=Location.span(type.location, self.previous().get_location()), @@ -246,7 +249,7 @@ class MidasParser(Parser): self.consume(TokenType.RIGHT_BRACKET, "Missing ']' after generic arguments") return args - def named_type(self) -> Type: + def named_type(self) -> NamedType: name: Token = self.consume_identifier("Expected type name") return NamedType( location=name.get_location(), @@ -281,6 +284,32 @@ class MidasParser(Parser): members=members, ) + def frame_type(self) -> FrameType: + keyword: Token = self.previous() + self.consume(TokenType.LEFT_BRACKET, "Expected '[' to start frame schema") + + columns: list[FrameType.Column] = [] + while not self.check(TokenType.RIGHT_BRACKET) and not self.is_at_end(): + name: Token = self.advance() + self.consume(TokenType.COLON, "Expected ':' between column name and type") + type: Type = self.type_expr() + columns.append( + FrameType.Column( + location=name.location_to(self.previous()), + name=name, + type=type, + ) + ) + if not self.match(TokenType.COMMA): + break + + self.consume(TokenType.RIGHT_BRACKET, "Unclosed frame schema") + + return FrameType( + location=keyword.location_to(self.previous()), + columns=columns, + ) + def constraint(self) -> Expr: """Parse a constraint @@ -370,7 +399,7 @@ class MidasParser(Parser): pos_args: list[Expr] = [] kw_args: dict[str, Expr] = {} keywords: bool = False - while not self.match(TokenType.RIGHT_PAREN): + while not self.check(TokenType.RIGHT_PAREN): if self.check_identifier() and self.check_next(TokenType.EQUAL): keywords = True keyword: Token = self.advance() diff --git a/midas/parser/python.py b/midas/parser/python.py index 7839f52..a6c9a78 100644 --- a/midas/parser/python.py +++ b/midas/parser/python.py @@ -30,6 +30,7 @@ from midas.ast.python import ( Stmt, SubscriptExpr, TernaryExpr, + TupleExpr, TypeAssign, UnaryExpr, VariableExpr, @@ -300,26 +301,28 @@ class PythonParser: case ast.Subscript(value=ast.Name(id="Frame"), slice=schema): return self._parse_frame_type(schema) - case ast.Subscript(value=ast.Name(id=name), slice=param): + case ast.Subscript(value=ast.Name(id=name), slice=arg): + args: tuple[MidasType, ...] = ( + tuple(self._parse_type(a) for a in arg.elts) + if isinstance(arg, ast.Tuple) + else (self._parse_type(arg),) + ) return BaseType( location=loc, base=name, - param=self._parse_type(param), + args=args, ) case ast.Name(id=name): return BaseType( location=loc, base=name, - param=None, + args=(), ) case ast.BinOp(left=left_expr, op=ast.Add(), right=right_expr): left = self._parse_type(left_expr) match left: - case None: - raise InvalidSyntaxError() - # If chained constraints, separate base type and rebuild constraint case ConstraintType(type=left_type, constraint=left_constraint): constraint = ast.BinOp( @@ -345,7 +348,7 @@ class PythonParser: return BaseType( location=loc, base="None", - param=None, + args=(), ) case _: @@ -477,6 +480,12 @@ class PythonParser: step=self.parse_expr(step) if step is not None else None, ) + case ast.Tuple(elts=items): + return TupleExpr( + location=location, + items=tuple(self.parse_expr(item) for item in items), + ) + case _: print(f"Unsupported expression: {ast.unparse(node)}") return RawExpr(location=location, expr=node) diff --git a/midas/typing.py b/midas/typing.py index 9c5b407..fc121ce 100644 --- a/midas/typing.py +++ b/midas/typing.py @@ -1,3 +1,4 @@ +from typing import Generic, TypeVar from typing import cast as typing_cast cast = typing_cast @@ -32,3 +33,20 @@ This operation is unsound, use at your own risk! _**Internal Python documentation**_ """ + + +T = TypeVar("T") + + +class Frame(Generic[T]): + """A `Frame` is the abstract type implemented by `DataFrame` + + A frame contains any number of named columns (see :class:`Column`) + """ + + +class Column(Generic[T]): + """A `Column` is the abstract type implemented by `Series` + + A column contains a any number of values of the same type + """ diff --git a/tests/cases/checker/01_simple_types.py.ref.json b/tests/cases/checker/01_simple_types.py.ref.json index ac24fcd..8767cf3 100644 --- a/tests/cases/checker/01_simple_types.py.ref.json +++ b/tests/cases/checker/01_simple_types.py.ref.json @@ -4,7 +4,35 @@ "type": "Warning", "location": { "start": [ - 6, + 8, + 12 + ], + "end": [ + 8, + 43 + ] + }, + "message": "ConstraintType not yet supported" + }, + { + "type": "Warning", + "location": { + "start": [ + 10, + 10 + ], + "end": [ + 10, + 18 + ] + }, + "message": "Unknown type 'datetime'" + }, + { + "type": "Warning", + "location": { + "start": [ + 13, 4 ], "end": [ @@ -12,7 +40,7 @@ 5 ] }, - "message": "FrameType not yet supported" + "message": "Unknown type '_'" } ], "judgments": [] diff --git a/tests/cases/checker/03_functions.py.ref.json b/tests/cases/checker/03_functions.py.ref.json index fa06642..820a7ea 100644 --- a/tests/cases/checker/03_functions.py.ref.json +++ b/tests/cases/checker/03_functions.py.ref.json @@ -328,6 +328,19 @@ }, "type": {} }, + { + "location": { + "from": "L6:9", + "to": "L6:10" + }, + "expr": { + "_type": "LiteralExpr", + "value": 1 + }, + "type": { + "name": "int" + } + }, { "location": { "from": "L6:5", @@ -373,19 +386,6 @@ } } }, - { - "location": { - "from": "L6:9", - "to": "L6:10" - }, - "expr": { - "_type": "LiteralExpr", - "value": 1 - }, - "type": { - "name": "int" - } - }, { "location": { "from": "L6:5", @@ -407,6 +407,32 @@ }, "type": {} }, + { + "location": { + "from": "L7:9", + "to": "L7:10" + }, + "expr": { + "_type": "LiteralExpr", + "value": 1 + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L7:12", + "to": "L7:15" + }, + "expr": { + "_type": "LiteralExpr", + "value": 2.0 + }, + "type": { + "name": "float" + } + }, { "location": { "from": "L7:5", @@ -452,32 +478,6 @@ } } }, - { - "location": { - "from": "L7:9", - "to": "L7:10" - }, - "expr": { - "_type": "LiteralExpr", - "value": 1 - }, - "type": { - "name": "int" - } - }, - { - "location": { - "from": "L7:12", - "to": "L7:15" - }, - "expr": { - "_type": "LiteralExpr", - "value": 2.0 - }, - "type": { - "name": "float" - } - }, { "location": { "from": "L7:5", @@ -503,6 +503,32 @@ }, "type": {} }, + { + "location": { + "from": "L8:9", + "to": "L8:10" + }, + "expr": { + "_type": "LiteralExpr", + "value": 1 + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L8:14", + "to": "L8:17" + }, + "expr": { + "_type": "LiteralExpr", + "value": 2.0 + }, + "type": { + "name": "float" + } + }, { "location": { "from": "L8:5", @@ -548,32 +574,6 @@ } } }, - { - "location": { - "from": "L8:9", - "to": "L8:10" - }, - "expr": { - "_type": "LiteralExpr", - "value": 1 - }, - "type": { - "name": "int" - } - }, - { - "location": { - "from": "L8:14", - "to": "L8:17" - }, - "expr": { - "_type": "LiteralExpr", - "value": 2.0 - }, - "type": { - "name": "float" - } - }, { "location": { "from": "L8:5", @@ -600,6 +600,45 @@ }, "type": {} }, + { + "location": { + "from": "L9:9", + "to": "L9:10" + }, + "expr": { + "_type": "LiteralExpr", + "value": 1 + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L9:12", + "to": "L9:15" + }, + "expr": { + "_type": "LiteralExpr", + "value": 2.0 + }, + "type": { + "name": "float" + } + }, + { + "location": { + "from": "L9:17", + "to": "L9:23" + }, + "expr": { + "_type": "LiteralExpr", + "value": "test" + }, + "type": { + "name": "str" + } + }, { "location": { "from": "L9:5", @@ -645,45 +684,6 @@ } } }, - { - "location": { - "from": "L9:9", - "to": "L9:10" - }, - "expr": { - "_type": "LiteralExpr", - "value": 1 - }, - "type": { - "name": "int" - } - }, - { - "location": { - "from": "L9:12", - "to": "L9:15" - }, - "expr": { - "_type": "LiteralExpr", - "value": 2.0 - }, - "type": { - "name": "float" - } - }, - { - "location": { - "from": "L9:17", - "to": "L9:23" - }, - "expr": { - "_type": "LiteralExpr", - "value": "test" - }, - "type": { - "name": "str" - } - }, { "location": { "from": "L9:5", @@ -713,6 +713,45 @@ }, "type": {} }, + { + "location": { + "from": "L10:9", + "to": "L10:10" + }, + "expr": { + "_type": "LiteralExpr", + "value": 1 + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L10:12", + "to": "L10:15" + }, + "expr": { + "_type": "LiteralExpr", + "value": 2.0 + }, + "type": { + "name": "float" + } + }, + { + "location": { + "from": "L10:19", + "to": "L10:22" + }, + "expr": { + "_type": "LiteralExpr", + "value": 3.0 + }, + "type": { + "name": "float" + } + }, { "location": { "from": "L10:5", @@ -758,45 +797,6 @@ } } }, - { - "location": { - "from": "L10:9", - "to": "L10:10" - }, - "expr": { - "_type": "LiteralExpr", - "value": 1 - }, - "type": { - "name": "int" - } - }, - { - "location": { - "from": "L10:12", - "to": "L10:15" - }, - "expr": { - "_type": "LiteralExpr", - "value": 2.0 - }, - "type": { - "name": "float" - } - }, - { - "location": { - "from": "L10:19", - "to": "L10:22" - }, - "expr": { - "_type": "LiteralExpr", - "value": 3.0 - }, - "type": { - "name": "float" - } - }, { "location": { "from": "L10:5", @@ -827,6 +827,19 @@ }, "type": {} }, + { + "location": { + "from": "L11:11", + "to": "L11:12" + }, + "expr": { + "_type": "LiteralExpr", + "value": 1 + }, + "type": { + "name": "int" + } + }, { "location": { "from": "L11:5", @@ -872,19 +885,6 @@ } } }, - { - "location": { - "from": "L11:11", - "to": "L11:12" - }, - "expr": { - "_type": "LiteralExpr", - "value": 1 - }, - "type": { - "name": "int" - } - }, { "location": { "from": "L11:5", @@ -906,6 +906,19 @@ }, "type": {} }, + { + "location": { + "from": "L12:11", + "to": "L12:17" + }, + "expr": { + "_type": "LiteralExpr", + "value": "test" + }, + "type": { + "name": "str" + } + }, { "location": { "from": "L12:5", @@ -951,19 +964,6 @@ } } }, - { - "location": { - "from": "L12:11", - "to": "L12:17" - }, - "expr": { - "_type": "LiteralExpr", - "value": "test" - }, - "type": { - "name": "str" - } - }, { "location": { "from": "L12:5", @@ -985,6 +985,45 @@ }, "type": {} }, + { + "location": { + "from": "L14:10", + "to": "L14:11" + }, + "expr": { + "_type": "LiteralExpr", + "value": 1 + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L14:13", + "to": "L14:16" + }, + "expr": { + "_type": "LiteralExpr", + "value": 2.0 + }, + "type": { + "name": "float" + } + }, + { + "location": { + "from": "L14:20", + "to": "L14:26" + }, + "expr": { + "_type": "LiteralExpr", + "value": "test" + }, + "type": { + "name": "str" + } + }, { "location": { "from": "L14:6", @@ -1030,45 +1069,6 @@ } } }, - { - "location": { - "from": "L14:10", - "to": "L14:11" - }, - "expr": { - "_type": "LiteralExpr", - "value": 1 - }, - "type": { - "name": "int" - } - }, - { - "location": { - "from": "L14:13", - "to": "L14:16" - }, - "expr": { - "_type": "LiteralExpr", - "value": 2.0 - }, - "type": { - "name": "float" - } - }, - { - "location": { - "from": "L14:20", - "to": "L14:26" - }, - "expr": { - "_type": "LiteralExpr", - "value": "test" - }, - "type": { - "name": "str" - } - }, { "location": { "from": "L14:6", @@ -1101,6 +1101,45 @@ "name": "bool" } }, + { + "location": { + "from": "L15:10", + "to": "L15:11" + }, + "expr": { + "_type": "LiteralExpr", + "value": 1 + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L15:15", + "to": "L15:18" + }, + "expr": { + "_type": "LiteralExpr", + "value": 2.0 + }, + "type": { + "name": "float" + } + }, + { + "location": { + "from": "L15:22", + "to": "L15:28" + }, + "expr": { + "_type": "LiteralExpr", + "value": "test" + }, + "type": { + "name": "str" + } + }, { "location": { "from": "L15:6", @@ -1146,45 +1185,6 @@ } } }, - { - "location": { - "from": "L15:10", - "to": "L15:11" - }, - "expr": { - "_type": "LiteralExpr", - "value": 1 - }, - "type": { - "name": "int" - } - }, - { - "location": { - "from": "L15:15", - "to": "L15:18" - }, - "expr": { - "_type": "LiteralExpr", - "value": 2.0 - }, - "type": { - "name": "float" - } - }, - { - "location": { - "from": "L15:22", - "to": "L15:28" - }, - "expr": { - "_type": "LiteralExpr", - "value": "test" - }, - "type": { - "name": "str" - } - }, { "location": { "from": "L15:6", @@ -1217,6 +1217,45 @@ "name": "bool" } }, + { + "location": { + "from": "L16:10", + "to": "L16:11" + }, + "expr": { + "_type": "LiteralExpr", + "value": 1 + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L16:15", + "to": "L16:21" + }, + "expr": { + "_type": "LiteralExpr", + "value": "test" + }, + "type": { + "name": "str" + } + }, + { + "location": { + "from": "L16:25", + "to": "L16:28" + }, + "expr": { + "_type": "LiteralExpr", + "value": 2.0 + }, + "type": { + "name": "float" + } + }, { "location": { "from": "L16:6", @@ -1262,45 +1301,6 @@ } } }, - { - "location": { - "from": "L16:10", - "to": "L16:11" - }, - "expr": { - "_type": "LiteralExpr", - "value": 1 - }, - "type": { - "name": "int" - } - }, - { - "location": { - "from": "L16:15", - "to": "L16:21" - }, - "expr": { - "_type": "LiteralExpr", - "value": "test" - }, - "type": { - "name": "str" - } - }, - { - "location": { - "from": "L16:25", - "to": "L16:28" - }, - "expr": { - "_type": "LiteralExpr", - "value": 2.0 - }, - "type": { - "name": "float" - } - }, { "location": { "from": "L16:6", @@ -1333,6 +1333,45 @@ "name": "bool" } }, + { + "location": { + "from": "L18:10", + "to": "L18:13" + }, + "expr": { + "_type": "LiteralExpr", + "value": "a" + }, + "type": { + "name": "str" + } + }, + { + "location": { + "from": "L18:15", + "to": "L18:16" + }, + "expr": { + "_type": "LiteralExpr", + "value": 3 + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L18:20", + "to": "L18:25" + }, + "expr": { + "_type": "LiteralExpr", + "value": false + }, + "type": { + "name": "bool" + } + }, { "location": { "from": "L18:6", @@ -1378,45 +1417,6 @@ } } }, - { - "location": { - "from": "L18:10", - "to": "L18:13" - }, - "expr": { - "_type": "LiteralExpr", - "value": "a" - }, - "type": { - "name": "str" - } - }, - { - "location": { - "from": "L18:15", - "to": "L18:16" - }, - "expr": { - "_type": "LiteralExpr", - "value": 3 - }, - "type": { - "name": "int" - } - }, - { - "location": { - "from": "L18:20", - "to": "L18:25" - }, - "expr": { - "_type": "LiteralExpr", - "value": false - }, - "type": { - "name": "bool" - } - }, { "location": { "from": "L18:6", diff --git a/tests/cases/checker/04_custom_types.py.ref.json b/tests/cases/checker/04_custom_types.py.ref.json index 01177d9..5698ca2 100644 --- a/tests/cases/checker/04_custom_types.py.ref.json +++ b/tests/cases/checker/04_custom_types.py.ref.json @@ -24,7 +24,7 @@ "type": { "_type": "BaseType", "base": "Meter", - "param": null + "args": [] }, "expr": { "_type": "LiteralExpr", @@ -62,7 +62,7 @@ "type": { "_type": "BaseType", "base": "Second", - "param": null + "args": [] }, "expr": { "_type": "LiteralExpr", diff --git a/tests/cases/checker/06_subtyping.py.ref.json b/tests/cases/checker/06_subtyping.py.ref.json index 3435f45..03194f1 100644 --- a/tests/cases/checker/06_subtyping.py.ref.json +++ b/tests/cases/checker/06_subtyping.py.ref.json @@ -100,6 +100,32 @@ "name": "float" } }, + { + "location": { + "from": "L11:13", + "to": "L11:15" + }, + "expr": { + "_type": "VariableExpr", + "name": "v1" + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L11:17", + "to": "L11:19" + }, + "expr": { + "_type": "VariableExpr", + "name": "v2" + }, + "type": { + "name": "float" + } + }, { "location": { "from": "L11:5", @@ -135,32 +161,6 @@ } } }, - { - "location": { - "from": "L11:13", - "to": "L11:15" - }, - "expr": { - "_type": "VariableExpr", - "name": "v1" - }, - "type": { - "name": "int" - } - }, - { - "location": { - "from": "L11:17", - "to": "L11:19" - }, - "expr": { - "_type": "VariableExpr", - "name": "v2" - }, - "type": { - "name": "float" - } - }, { "location": { "from": "L11:5", diff --git a/tests/cases/checker/07_variance.py.ref.json b/tests/cases/checker/07_variance.py.ref.json index a43b957..abfebf3 100644 --- a/tests/cases/checker/07_variance.py.ref.json +++ b/tests/cases/checker/07_variance.py.ref.json @@ -72,29 +72,6 @@ } ], "judgments": [ - { - "location": { - "from": "L26:0", - "to": "L26:5" - }, - "expr": { - "_type": "VariableExpr", - "name": "print" - }, - "type": { - "pos_args": [ - { - "pos": 0, - "name": "object", - "type": {}, - "required": true - } - ], - "args": [], - "kw_args": [], - "returns": {} - } - }, { "location": { "from": "L27:4", @@ -325,6 +302,29 @@ } } }, + { + "location": { + "from": "L26:0", + "to": "L26:5" + }, + "expr": { + "_type": "VariableExpr", + "name": "print" + }, + "type": { + "pos_args": [ + { + "pos": 0, + "name": "object", + "type": {}, + "required": false + } + ], + "args": [], + "kw_args": [], + "returns": {} + } + }, { "location": { "from": "L26:0", diff --git a/tests/cases/checker/08_unification.py.ref.json b/tests/cases/checker/08_unification.py.ref.json index bfaa7bd..c384d8a 100644 --- a/tests/cases/checker/08_unification.py.ref.json +++ b/tests/cases/checker/08_unification.py.ref.json @@ -63,31 +63,6 @@ "name": "float" } }, - { - "location": { - "from": "L6:11", - "to": "L6:15" - }, - "expr": { - "_type": "VariableExpr", - "name": "bool" - }, - "type": { - "pos_args": [ - { - "pos": 0, - "name": "object", - "type": {}, - "required": false - } - ], - "args": [], - "kw_args": [], - "returns": { - "name": "bool" - } - } - }, { "location": { "from": "L6:16", @@ -135,6 +110,31 @@ "name": "int" } }, + { + "location": { + "from": "L6:11", + "to": "L6:15" + }, + "expr": { + "_type": "VariableExpr", + "name": "bool" + }, + "type": { + "pos_args": [ + { + "pos": 0, + "name": "object", + "type": {}, + "required": false + } + ], + "args": [], + "kw_args": [], + "returns": { + "name": "bool" + } + } + }, { "location": { "from": "L6:11", @@ -367,6 +367,54 @@ } } }, + { + "location": { + "from": "L12:21", + "to": "L12:27" + }, + "expr": { + "_type": "VariableExpr", + "name": "double" + }, + "type": { + "pos_args": [], + "args": [ + { + "pos": 0, + "name": "value", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw_args": [], + "returns": { + "name": "float" + } + } + }, + { + "location": { + "from": "L12:29", + "to": "L12:35" + }, + "expr": { + "_type": "VariableExpr", + "name": "floats" + }, + "type": { + "name": "list", + "args": [ + { + "name": "float" + } + ], + "body": { + "name": "list" + } + } + }, { "location": { "from": "L12:17", @@ -455,54 +503,6 @@ } } }, - { - "location": { - "from": "L12:21", - "to": "L12:27" - }, - "expr": { - "_type": "VariableExpr", - "name": "double" - }, - "type": { - "pos_args": [], - "args": [ - { - "pos": 0, - "name": "value", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [], - "returns": { - "name": "float" - } - } - }, - { - "location": { - "from": "L12:29", - "to": "L12:35" - }, - "expr": { - "_type": "VariableExpr", - "name": "floats" - }, - "type": { - "name": "list", - "args": [ - { - "name": "float" - } - ], - "body": { - "name": "list" - } - } - }, { "location": { "from": "L12:17", @@ -538,6 +538,54 @@ } } }, + { + "location": { + "from": "L13:19", + "to": "L13:25" + }, + "expr": { + "_type": "VariableExpr", + "name": "double" + }, + "type": { + "pos_args": [], + "args": [ + { + "pos": 0, + "name": "value", + "type": { + "name": "float" + }, + "required": true + } + ], + "kw_args": [], + "returns": { + "name": "float" + } + } + }, + { + "location": { + "from": "L13:27", + "to": "L13:31" + }, + "expr": { + "_type": "VariableExpr", + "name": "ints" + }, + "type": { + "name": "list", + "args": [ + { + "name": "int" + } + ], + "body": { + "name": "list" + } + } + }, { "location": { "from": "L13:15", @@ -626,54 +674,6 @@ } } }, - { - "location": { - "from": "L13:19", - "to": "L13:25" - }, - "expr": { - "_type": "VariableExpr", - "name": "double" - }, - "type": { - "pos_args": [], - "args": [ - { - "pos": 0, - "name": "value", - "type": { - "name": "float" - }, - "required": true - } - ], - "kw_args": [], - "returns": { - "name": "float" - } - } - }, - { - "location": { - "from": "L13:27", - "to": "L13:31" - }, - "expr": { - "_type": "VariableExpr", - "name": "ints" - }, - "type": { - "name": "list", - "args": [ - { - "name": "int" - } - ], - "body": { - "name": "list" - } - } - }, { "location": { "from": "L13:15", @@ -699,6 +699,54 @@ }, "type": {} }, + { + "location": { + "from": "L14:15", + "to": "L14:21" + }, + "expr": { + "_type": "VariableExpr", + "name": "is_odd" + }, + "type": { + "pos_args": [], + "args": [ + { + "pos": 0, + "name": "value", + "type": { + "name": "int" + }, + "required": true + } + ], + "kw_args": [], + "returns": { + "name": "bool" + } + } + }, + { + "location": { + "from": "L14:23", + "to": "L14:27" + }, + "expr": { + "_type": "VariableExpr", + "name": "ints" + }, + "type": { + "name": "list", + "args": [ + { + "name": "int" + } + ], + "body": { + "name": "list" + } + } + }, { "location": { "from": "L14:11", @@ -787,54 +835,6 @@ } } }, - { - "location": { - "from": "L14:15", - "to": "L14:21" - }, - "expr": { - "_type": "VariableExpr", - "name": "is_odd" - }, - "type": { - "pos_args": [], - "args": [ - { - "pos": 0, - "name": "value", - "type": { - "name": "int" - }, - "required": true - } - ], - "kw_args": [], - "returns": { - "name": "bool" - } - } - }, - { - "location": { - "from": "L14:23", - "to": "L14:27" - }, - "expr": { - "_type": "VariableExpr", - "name": "ints" - }, - "type": { - "name": "list", - "args": [ - { - "name": "int" - } - ], - "body": { - "name": "list" - } - } - }, { "location": { "from": "L14:11", diff --git a/tests/cases/python-parser/01_simple_types.py.ref.json b/tests/cases/python-parser/01_simple_types.py.ref.json index 452b9c0..50730f3 100644 --- a/tests/cases/python-parser/01_simple_types.py.ref.json +++ b/tests/cases/python-parser/01_simple_types.py.ref.json @@ -16,7 +16,7 @@ "type": { "_type": "BaseType", "base": "bool", - "param": null + "args": [] } }, { @@ -25,7 +25,7 @@ "type": { "_type": "BaseType", "base": "int", - "param": null + "args": [] } }, { @@ -36,7 +36,7 @@ "type": { "_type": "BaseType", "base": "float", - "param": null + "args": [] }, "constraint": "(_ > 0) + (_ < 250)" } @@ -47,7 +47,7 @@ "type": { "_type": "BaseType", "base": "str", - "param": null + "args": [] } }, { @@ -56,7 +56,7 @@ "type": { "_type": "BaseType", "base": "datetime", - "param": null + "args": [] } }, { @@ -65,7 +65,7 @@ "type": { "_type": "BaseType", "base": "float", - "param": null + "args": [] } }, { @@ -79,7 +79,7 @@ "type": { "_type": "BaseType", "base": "_", - "param": null + "args": [] } } ] diff --git a/tests/cases/python-parser/02_custom_types.py.ref.json b/tests/cases/python-parser/02_custom_types.py.ref.json index 9d77ebd..fdb89d3 100644 --- a/tests/cases/python-parser/02_custom_types.py.ref.json +++ b/tests/cases/python-parser/02_custom_types.py.ref.json @@ -16,7 +16,7 @@ "type": { "_type": "BaseType", "base": "GeoLocation", - "param": null + "args": [] } } ] @@ -28,11 +28,13 @@ "type": { "_type": "BaseType", "base": "Column", - "param": { - "_type": "BaseType", - "base": "GeoLocation", - "param": null - } + "args": [ + { + "_type": "BaseType", + "base": "GeoLocation", + "args": [] + } + ] } }, { @@ -65,11 +67,13 @@ "type": { "_type": "BaseType", "base": "Column", - "param": { - "_type": "BaseType", - "base": "GeoLocation", - "param": null - } + "args": [ + { + "_type": "BaseType", + "base": "GeoLocation", + "args": [] + } + ] } }, { @@ -117,7 +121,7 @@ "type": { "_type": "BaseType", "base": "Latitude", - "param": null + "args": [] } }, { @@ -146,7 +150,7 @@ "type": { "_type": "BaseType", "base": "Latitude", - "param": null + "args": [] } }, { @@ -175,11 +179,13 @@ "type": { "_type": "BaseType", "base": "Difference", - "param": { - "_type": "BaseType", - "base": "Latitude", - "param": null - } + "args": [ + { + "_type": "BaseType", + "base": "Latitude", + "args": [] + } + ] } }, { @@ -217,7 +223,7 @@ "type": { "_type": "BaseType", "base": "int", - "param": null + "args": [] }, "constraint": "_ >= 0" } @@ -230,7 +236,7 @@ "type": { "_type": "BaseType", "base": "float", - "param": null + "args": [] }, "constraint": "_ >= 0" } @@ -252,7 +258,7 @@ "type": { "_type": "BaseType", "base": "int", - "param": null + "args": [] }, "constraint": "Positive" } @@ -265,7 +271,7 @@ "type": { "_type": "BaseType", "base": "float", - "param": null + "args": [] }, "constraint": "Positive" } diff --git a/tests/cases/python-parser/03_functions.py.ref.json b/tests/cases/python-parser/03_functions.py.ref.json index a8f261f..2288b82 100644 --- a/tests/cases/python-parser/03_functions.py.ref.json +++ b/tests/cases/python-parser/03_functions.py.ref.json @@ -14,15 +14,17 @@ "type": { "_type": "BaseType", "base": "Column", - "param": { - "_type": "ConstraintType", - "type": { - "_type": "BaseType", - "base": "float", - "param": null - }, - "constraint": "0 <= _ <= 1" - } + "args": [ + { + "_type": "ConstraintType", + "type": { + "_type": "BaseType", + "base": "float", + "args": [] + }, + "constraint": "0 <= _ <= 1" + } + ] }, "default": null }, @@ -31,15 +33,17 @@ "type": { "_type": "BaseType", "base": "Column", - "param": { - "_type": "ConstraintType", - "type": { - "_type": "BaseType", - "base": "float", - "param": null - }, - "constraint": "0 <= _ <= 1" - } + "args": [ + { + "_type": "ConstraintType", + "type": { + "_type": "BaseType", + "base": "float", + "args": [] + }, + "constraint": "0 <= _ <= 1" + } + ] }, "default": null } @@ -50,15 +54,17 @@ "returns": { "_type": "BaseType", "base": "Column", - "param": { - "_type": "ConstraintType", - "type": { - "_type": "BaseType", - "base": "float", - "param": null - }, - "constraint": "0 <= _ <= 2" - } + "args": [ + { + "_type": "ConstraintType", + "type": { + "_type": "BaseType", + "base": "float", + "args": [] + }, + "constraint": "0 <= _ <= 2" + } + ] }, "body": [ { @@ -67,15 +73,17 @@ "type": { "_type": "BaseType", "base": "Column", - "param": { - "_type": "ConstraintType", - "type": { - "_type": "BaseType", - "base": "float", - "param": null - }, - "constraint": "0 <= _ <= 2" - } + "args": [ + { + "_type": "ConstraintType", + "type": { + "_type": "BaseType", + "base": "float", + "args": [] + }, + "constraint": "0 <= _ <= 2" + } + ] } }, { @@ -117,7 +125,7 @@ "type": { "_type": "BaseType", "base": "int", - "param": null + "args": [] }, "default": null } @@ -128,7 +136,7 @@ "type": { "_type": "BaseType", "base": "float", - "param": null + "args": [] }, "default": null } @@ -140,7 +148,7 @@ "type": { "_type": "BaseType", "base": "str", - "param": null + "args": [] }, "default": null } diff --git a/tests/generator.py b/tests/generator.py index 067790f..f04992b 100644 --- a/tests/generator.py +++ b/tests/generator.py @@ -46,7 +46,8 @@ class GeneratorTester(Tester): if not any(d.type == DiagnosticType.ERROR for d in checker.diagnostics): generator = Generator(workdir=path.parent, types=checker.types) - result.compiled_ast = generator.generate_ast(typed_ast, path) + generator.set_src_path(path) + result.compiled_ast = generator.generate_ast(typed_ast) return result diff --git a/tests/serializer/midas.py b/tests/serializer/midas.py index 0ad86f3..7a3f8e8 100644 --- a/tests/serializer/midas.py +++ b/tests/serializer/midas.py @@ -8,6 +8,7 @@ from midas.ast.midas import ( Expr, ExtendStmt, ExtensionType, + FrameType, FunctionType, GenericType, GetExpr, @@ -197,3 +198,15 @@ class MidasAstJsonSerializer( "base": type.base.accept(self), "extension": type.extension.accept(self), } + + def visit_frame_type(self, type: FrameType) -> dict: + return { + "_type": "FrameType", + "columns": [self._serialize_column(col) for col in type.columns], + } + + def _serialize_column(self, column: FrameType.Column): + return { + "name": column.name.lexeme, + "type": column.type.accept(self), + } diff --git a/tests/serializer/python.py b/tests/serializer/python.py index 3904739..4e531cb 100644 --- a/tests/serializer/python.py +++ b/tests/serializer/python.py @@ -30,6 +30,7 @@ from midas.ast.python import ( Stmt, SubscriptExpr, TernaryExpr, + TupleExpr, TypeAssign, UnaryExpr, VariableExpr, @@ -98,7 +99,7 @@ class PythonAstJsonSerializer( return { "_type": "BaseType", "base": node.base, - "param": self._serialize_optional(node.param), + "args": self._serialize_list(node.args), } def visit_constraint_type(self, node: ConstraintType) -> dict: @@ -302,6 +303,12 @@ class PythonAstJsonSerializer( "step": self._serialize_optional(expr.step), } + def visit_tuple_expr(self, expr: TupleExpr) -> dict: + return { + "_type": "TupleExpr", + "items": [item.accept(self) for item in expr.items], + } + def visit_raw_expr(self, expr: RawExpr) -> dict: return { "_type": "RawExpr",