From 556765fd3516daa9137280c04468b9ae469f61fa Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 22 Jun 2026 10:29:06 +0200 Subject: [PATCH 01/39] feat(types): add DataFrameType and ColumnType --- midas/checker/python.py | 25 ++++++++++++++++++----- midas/checker/types.py | 44 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 63 insertions(+), 6 deletions(-) diff --git a/midas/checker/python.py b/midas/checker/python.py index 6a01764..38bef94 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -20,7 +20,9 @@ from midas.checker.resolver import Resolver from midas.checker.types import ( AppliedType, BaseType, + ColumnType, ConstraintType, + DataFrameType, DerivedType, Function, GenericType, @@ -677,13 +679,26 @@ class PythonTyper( 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, diff --git a/midas/checker/types.py b/midas/checker/types.py index bf1703b..8bbac53 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,22 @@ class ConstraintType: return f"{self.type} where {printer.print(self.constraint)}" +@dataclass(frozen=True, kw_only=True) +class ColumnType: + type: Type + + +@dataclass(frozen=True, kw_only=True) +class DataFrameType: + columns: list[Column] + + @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 +181,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 +275,21 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: body=substitute_typevars(body, substitutions), ) + 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 +353,12 @@ def to_annotation(type: Type) -> str: case ConstraintType(): return str(type) + case ColumnType(): + return "pd.Series" + + case DataFrameType(): + return "pd.DataFrame" + case _: assert_never(type) @@ -344,4 +384,6 @@ Type = ( | GenericType | AppliedType | ConstraintType + | ColumnType + | DataFrameType ) -- 2.49.1 From 5357ca8e58b0dd0f5660af35b353749325a5ea6f Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 23 Jun 2026 10:58:45 +0200 Subject: [PATCH 02/39] fix(types): add str methods to dataframe types --- midas/checker/types.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/midas/checker/types.py b/midas/checker/types.py index 8bbac53..67c7462 100644 --- a/midas/checker/types.py +++ b/midas/checker/types.py @@ -160,11 +160,18 @@ class ConstraintType: 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 -- 2.49.1 From 0d20993f02374d8b5e169ec23c000772ef29544e Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 23 Jun 2026 12:26:39 +0200 Subject: [PATCH 03/39] feat(types): add TupleType --- midas/checker/python.py | 21 +++++++++++++++++++++ midas/checker/types.py | 17 +++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/midas/checker/python.py b/midas/checker/python.py index 38bef94..743c8d2 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -27,6 +27,7 @@ from midas.checker.types import ( Function, GenericType, OverloadedFunction, + TupleType, Type, TypeVar, UnitType, @@ -642,6 +643,11 @@ class PythonTyper( 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) + operation: Optional[Type] = self.types.lookup_member(object, "__getitem__") if operation is None: self.reporter.error( @@ -1231,3 +1237,18 @@ class PythonTyper( 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() diff --git a/midas/checker/types.py b/midas/checker/types.py index 67c7462..a4279af 100644 --- a/midas/checker/types.py +++ b/midas/checker/types.py @@ -156,6 +156,14 @@ 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 @@ -282,6 +290,11 @@ 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), @@ -360,6 +373,9 @@ 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" @@ -391,6 +407,7 @@ Type = ( | GenericType | AppliedType | ConstraintType + | TupleType | ColumnType | DataFrameType ) -- 2.49.1 From a31d295eb15f0842920555793752647ce45314a4 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 23 Jun 2026 12:27:31 +0200 Subject: [PATCH 04/39] feat(checker): type check subscript on dataframes --- midas/checker/frames.py | 51 +++++++++++++++++++++++++++++++++++++++++ midas/checker/python.py | 41 ++++++++++++++++++++++++++++++++- 2 files changed, 91 insertions(+), 1 deletion(-) create mode 100644 midas/checker/frames.py diff --git a/midas/checker/frames.py b/midas/checker/frames.py new file mode 100644 index 0000000..069143a --- /dev/null +++ b/midas/checker/frames.py @@ -0,0 +1,51 @@ +from typing import Optional + +from midas.checker.types import ColumnType, DataFrameType + + +class FrameManager: + @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 + 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] diff --git a/midas/checker/python.py b/midas/checker/python.py index 743c8d2..2e75698 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -1,13 +1,14 @@ import ast import logging from dataclasses import dataclass -from typing import Any, Optional +from typing import Any, Optional, cast import midas.ast.python as p 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, @@ -647,6 +648,8 @@ class PythonTyper( 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: @@ -1252,3 +1255,39 @@ class PythonTyper( expr.location, f"Invalid index type {expr.index} on {tup}" ) return UnknownType() + + def _visit_frame_subscript( + self, frame: DataFrameType, expr: p.SubscriptExpr + ) -> Type: + match expr.index: + case p.LiteralExpr(value=str() as name): + column: Optional[ColumnType] = FrameManager.get_column(frame, name) + if column is None: + self.reporter.error( + expr.location, f"Unknown column '{name}' on {frame}" + ) + return UnknownType() + return column + + case p.ListExpr(items=indices) if all( + isinstance(index, p.LiteralExpr) and isinstance(index.value, str) + for index in indices + ): + indices = cast(list[p.LiteralExpr], 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: + self.reporter.error( + expr.location, f"Unknown column '{name}' on {frame}" + ) + return UnknownType() + columns.append(column) + return TupleType(items=tuple(columns)) + + case _: + self.reporter.error( + expr.location, f"Invalid index type {expr.index} on {frame}" + ) + return UnknownType() -- 2.49.1 From c91b20679198020b71cf36416e7619e392a3b039 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 23 Jun 2026 14:02:13 +0200 Subject: [PATCH 05/39] feat(checker): handle setting dataframe column --- midas/checker/frames.py | 98 +++++++++++++++++++++++++++++++++++---- midas/checker/python.py | 63 ++++++++++++------------- midas/checker/resolver.py | 4 ++ 3 files changed, 124 insertions(+), 41 deletions(-) diff --git a/midas/checker/frames.py b/midas/checker/frames.py index 069143a..0dbd95b 100644 --- a/midas/checker/frames.py +++ b/midas/checker/frames.py @@ -1,11 +1,92 @@ -from typing import Optional +from typing import Optional, TypeGuard, cast -from midas.checker.types import ColumnType, DataFrameType +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, TupleType, Type, UnknownType + +import midas.ast.python as p + + +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, types: TypesRegistry) -> None: + self.types: TypesRegistry = types + + 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( + def _set_column( cls, frame: DataFrameType, name: str, column: ColumnType ) -> DataFrameType: new_columns: list[DataFrameType.Column] = [] @@ -15,6 +96,7 @@ class FrameManager: 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( @@ -30,22 +112,22 @@ class FrameManager: return DataFrameType(columns=new_columns) @classmethod - def set_columns( + 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) + frame = cls._set_column(frame, name, col) return frame @classmethod - def get_column(cls, frame: DataFrameType, name: str) -> Optional[ColumnType]: + 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( + def _get_columns( cls, frame: DataFrameType, names: list[str] ) -> list[Optional[ColumnType]]: - return [cls.get_column(frame, name) for name in names] + return [cls._get_column(frame, name) for name in names] diff --git a/midas/checker/python.py b/midas/checker/python.py index 2e75698..39888cf 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -1,7 +1,7 @@ import ast import logging from dataclasses import dataclass -from typing import Any, Optional, cast +from typing import Any, Optional import midas.ast.python as p from midas.ast.location import Location @@ -75,6 +75,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.types) self.global_env: Environment = Preamble(self.types) self.env: Environment = self.global_env self.locals: dict[p.Expr, int] = {} @@ -323,9 +324,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}") @@ -364,6 +371,27 @@ 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) + # TODO: what happens if type is an alias of a dataframe type + match var_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 _: + 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) @@ -1259,35 +1287,4 @@ class PythonTyper( def _visit_frame_subscript( self, frame: DataFrameType, expr: p.SubscriptExpr ) -> Type: - match expr.index: - case p.LiteralExpr(value=str() as name): - column: Optional[ColumnType] = FrameManager.get_column(frame, name) - if column is None: - self.reporter.error( - expr.location, f"Unknown column '{name}' on {frame}" - ) - return UnknownType() - return column - - case p.ListExpr(items=indices) if all( - isinstance(index, p.LiteralExpr) and isinstance(index.value, str) - for index in indices - ): - indices = cast(list[p.LiteralExpr], 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: - self.reporter.error( - expr.location, f"Unknown column '{name}' on {frame}" - ) - return UnknownType() - columns.append(column) - return TupleType(items=tuple(columns)) - - case _: - self.reporter.error( - expr.location, f"Invalid index type {expr.index} on {frame}" - ) - return UnknownType() + return self.frame_mgr.get(self.reporter, expr.location, frame, expr.index) diff --git a/midas/checker/resolver.py b/midas/checker/resolver.py index 3bf73d7..753f4db 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}") -- 2.49.1 From 87d5e286d2cf88d8d41a4effb66ea5b867b44fcc Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 23 Jun 2026 14:45:19 +0200 Subject: [PATCH 06/39] feat(gen): add support for tuples and dataframes --- midas/generator/generator.py | 22 +++++++++++++++++- midas/generator/stubs.py | 43 +++++++++++++++++++++++++++++++++++- 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/midas/generator/generator.py b/midas/generator/generator.py index 3f87ef2..a7641b5 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 @@ -12,14 +13,17 @@ 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, @@ -40,6 +44,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): 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=[], @@ -332,6 +337,19 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): if bound is not None: self._make_cast_asserts(src_location, expr, bound) + case TupleType(items=items): + self._add_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): + self._make_cast_asserts(src_location, item, item_type) + case ( TopType() | Function() @@ -339,8 +357,10 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): | ComplexType() | ExtensionType() | GenericType() + | ColumnType() + | DataFrameType() ): - raise NotImplementedError(f"Can't make assertion for type {type}") + self.logger.warning(f"Can't make assertion for type {type}") # Ensure exhaustiveness case _: 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) -- 2.49.1 From ede7272c097d75e770cf3e24df88669943f059a4 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 23 Jun 2026 15:56:07 +0200 Subject: [PATCH 07/39] feat(parser): add frame type to midas syntax --- gen/midas.py | 10 ++++++++++ midas/ast/midas.py | 17 +++++++++++++++++ midas/ast/printer.py | 36 ++++++++++++++++++++++++++++++++++++ midas/parser/midas.py | 33 +++++++++++++++++++++++++++++++-- 4 files changed, 94 insertions(+), 2 deletions(-) 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/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..d713a74 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, diff --git a/midas/parser/midas.py b/midas/parser/midas.py index c5fced1..77ad40b 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 -- 2.49.1 From 7074b074bca9896481d27192741196df184f73aa Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 23 Jun 2026 15:58:27 +0200 Subject: [PATCH 08/39] feat(cli): add frame type to highlighter --- midas/cli/highlighter.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/midas/cli/highlighter.py b/midas/cli/highlighter.py index ce6c2b2..c131143 100644 --- a/midas/cli/highlighter.py +++ b/midas/cli/highlighter.py @@ -350,6 +350,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" -- 2.49.1 From cdf1725c2661217a14c470f95ea49605cbe9a484 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 23 Jun 2026 16:12:27 +0200 Subject: [PATCH 09/39] feat(checker): process frame type definitions --- midas/checker/midas.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) 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: -- 2.49.1 From 7219109e5d2329d85aa324c000ade3ec6e8ae19f Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 25 Jun 2026 16:32:15 +0200 Subject: [PATCH 10/39] feat(cli): print context for multiline diagnostics --- midas/cli/utils.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) 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() -- 2.49.1 From b33fadf768f006324c7c87ac76550bf1f0c8e49c Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 25 Jun 2026 21:09:14 +0200 Subject: [PATCH 11/39] feat(checker): add structural subtyping rule for dataframes --- midas/checker/registry.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/midas/checker/registry.py b/midas/checker/registry.py index 7520981..0ef2d76 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, @@ -157,6 +159,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 +207,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 -- 2.49.1 From 08dd7408ece5f2e7b6729b7d4601d84392766105 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 25 Jun 2026 21:34:00 +0200 Subject: [PATCH 12/39] feat(checker): defined add method of dataframes --- midas/checker/frame_methods.py | 124 +++++++++++++++++++++++++++++++++ 1 file changed, 124 insertions(+) create mode 100644 midas/checker/frame_methods.py diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py new file mode 100644 index 0000000..60c6eb1 --- /dev/null +++ b/midas/checker/frame_methods.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Protocol + +from midas.ast.location import Location +from midas.checker.reporter import FileReporter +from midas.checker.types import ( + ColumnType, + DataFrameType, + Function, + Type, + UnknownType, + unfold_type, +) + +if TYPE_CHECKING: + from midas.checker.frames import FrameManager + from midas.checker.python import PythonTyper, TypedExpr + + +class FrameMethod(Protocol): + @property + def __name__(self) -> str: ... + def __call__( + self, + typer: PythonTyper, + manager: FrameManager, + reporter: FileReporter, + location: Location, + frame: DataFrameType, + positional: list[TypedExpr], + keywords: dict[str, TypedExpr], + ) -> Type: ... + + +FRAME_METHODS: dict[str, FrameMethod] = {} + + +def frame_method(*names: str): + def wrapper(func: FrameMethod): + names_: tuple[str, ...] = names + if len(names_) == 0: + names_ = (func.__name__,) + for name in names_: + FRAME_METHODS[name] = func + return func + + return wrapper + + +@frame_method("add", "__add__") +def add( + typer: PythonTyper, + manager: FrameManager, + reporter: FileReporter, + location: Location, + frame: DataFrameType, + positional: list[TypedExpr], + keywords: dict[str, TypedExpr], +) -> Type: + new_columns: list[DataFrameType.Column] = [] + + by_name: dict[str, DataFrameType.Column] = {} + frame2: Optional[DataFrameType] = None + if len(positional) != 0: + other: Type = 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 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 manager.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 ( + typer._get_call_result( + location=location, + callee=signature, + positional=positional, + keywords=keywords, + ) + or UnknownType() + ) -- 2.49.1 From 946b2e0d2e4760b86cd79dfd1e19ce741b0fc6c3 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 25 Jun 2026 21:34:59 +0200 Subject: [PATCH 13/39] feat(checker): lookup dataframe methods --- midas/checker/frames.py | 22 ++++++++++++++++++++-- midas/checker/python.py | 24 +++++++++++++++++++++++- 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/midas/checker/frames.py b/midas/checker/frames.py index 0dbd95b..50ae080 100644 --- a/midas/checker/frames.py +++ b/midas/checker/frames.py @@ -1,11 +1,16 @@ -from typing import Optional, TypeGuard, cast +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 FRAME_METHODS, FrameMethod from midas.checker.registry import TypesRegistry from midas.checker.reporter import FileReporter from midas.checker.types import ColumnType, DataFrameType, TupleType, Type, UnknownType -import midas.ast.python as p +if TYPE_CHECKING: + from midas.checker.python import PythonTyper, TypedExpr def is_list_of_literals(exprs: list[p.Expr]) -> TypeGuard[list[p.LiteralExpr]]: @@ -131,3 +136,16 @@ class FrameManager: cls, frame: DataFrameType, names: list[str] ) -> list[Optional[ColumnType]]: return [cls._get_column(frame, name) for name in names] + + def call_method( + self, + typer: PythonTyper, + reporter: FileReporter, + location: Location, + frame: DataFrameType, + method: str, + positional: list[TypedExpr], + keywords: dict[str, TypedExpr], + ) -> Type: + function: FrameMethod = FRAME_METHODS[method] + return function(typer, self, reporter, location, frame, positional, keywords) diff --git a/midas/checker/python.py b/midas/checker/python.py index 39888cf..3a2c96d 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.frame_methods import FRAME_METHODS from midas.checker.frames import FrameManager from midas.checker.operators import ( PY_COMPARATOR_METHODS, @@ -515,13 +516,31 @@ 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) and self._is_frame_method( + method + ): + return self.frame_mgr.call_method( + self, + self.reporter, + expr.location, + unfolded, + method, + positional, + keywords, + ) + + callee: Type = self.type_of(expr.callee) return ( self._get_call_result( location=expr.location, @@ -1288,3 +1307,6 @@ class PythonTyper( self, frame: DataFrameType, expr: p.SubscriptExpr ) -> Type: return self.frame_mgr.get(self.reporter, expr.location, frame, expr.index) + + def _is_frame_method(self, method: str) -> bool: + return method in FRAME_METHODS -- 2.49.1 From f84a19159ffba57556f0981e2af2f9842be7279e Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 25 Jun 2026 21:35:19 +0200 Subject: [PATCH 14/39] fix(checker): improve heterogeneous error message --- midas/checker/python.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/midas/checker/python.py b/midas/checker/python.py index 3a2c96d..3a3c547 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -645,7 +645,7 @@ class PythonTyper( return self.types.apply_generic(list_type, [item_type]) self.reporter.error( 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()]) @@ -677,7 +677,7 @@ class PythonTyper( else: self.reporter.error( expr.location, - f"Heterogeneous dict keys: {key_types}", + f"Heterogeneous dict keys: [{', '.join(map(str, key_types))}]", ) if len(value_types) == 1: @@ -685,7 +685,7 @@ class PythonTyper( else: self.reporter.error( 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]) -- 2.49.1 From 4efb01285cb2fc151972170c932bc1ff155c582c Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 25 Jun 2026 21:35:47 +0200 Subject: [PATCH 15/39] feat: add dummy classes for typing frames and columns --- midas/typing.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) 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 + """ -- 2.49.1 From 00e2ca8fe32e7ff4fae36f7b11187c64407b30d3 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 25 Jun 2026 22:14:25 +0200 Subject: [PATCH 16/39] refactor: add MethodResolver class --- midas/checker/frame_methods.py | 215 ++++++++++++++++++--------------- midas/checker/frames.py | 23 ++-- midas/checker/python.py | 16 +-- 3 files changed, 137 insertions(+), 117 deletions(-) diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py index 60c6eb1..2bbc016 100644 --- a/midas/checker/frame_methods.py +++ b/midas/checker/frame_methods.py @@ -1,8 +1,10 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Optional, Protocol +from dataclasses import dataclass +from typing import TYPE_CHECKING, 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, @@ -14,111 +16,134 @@ from midas.checker.types import ( ) if TYPE_CHECKING: - from midas.checker.frames import FrameManager from midas.checker.python import PythonTyper, TypedExpr -class FrameMethod(Protocol): +@dataclass(frozen=True, kw_only=True) +class Call: + location: Location + frame: DataFrameType + positional: list[TypedExpr] + keywords: dict[str, TypedExpr] + + +def _is_method(obj: object, method: str) -> bool: + if not callable(obj): + return False + if not hasattr(obj, "__method_names__"): + return False + return method in obj.__method_names__ # type: ignore + + +class MethodResolver: + @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 + + def __init__(self, typer: PythonTyper) -> None: + self.typer: PythonTyper = typer + @property - def __name__(self) -> str: ... - def __call__( + def reporter(self) -> FileReporter: + return self.typer.reporter + + @property + def types(self) -> TypesRegistry: + return self.typer.types + + def _get_method_by_name(self, method: str) -> Optional[Callable]: + for name in dir(self): + attr = getattr(self, name) + if _is_method(attr, method): + return attr + return None + + def call( self, - typer: PythonTyper, - manager: FrameManager, - reporter: FileReporter, - location: Location, - frame: DataFrameType, - positional: list[TypedExpr], - keywords: dict[str, TypedExpr], - ) -> Type: ... + method: str, + call: Call, + ) -> Type: + func: Optional[Callable] = self._get_method_by_name(method) + if func is None: + self.reporter.error(call.location, f"Unknown method {method}") + return UnknownType() + return func(call) + @frame_method("add", "__add__") + def add( + self, + call: Call, + ) -> Type: + new_columns: list[DataFrameType.Column] = [] -FRAME_METHODS: dict[str, FrameMethod] = {} + 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) -def frame_method(*names: str): - def wrapper(func: FrameMethod): - names_: tuple[str, ...] = names - if len(names_) == 0: - names_ = (func.__name__,) - for name in names_: - FRAME_METHODS[name] = func - return func + 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 - return wrapper - - -@frame_method("add", "__add__") -def add( - typer: PythonTyper, - manager: FrameManager, - reporter: FileReporter, - location: Location, - frame: DataFrameType, - positional: list[TypedExpr], - keywords: dict[str, TypedExpr], -) -> Type: - new_columns: list[DataFrameType.Column] = [] - - by_name: dict[str, DataFrameType.Column] = {} - frame2: Optional[DataFrameType] = None - if len(positional) != 0: - other: Type = 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 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 manager.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()), - ) + new_column = DataFrameType.Column( + index=column.index, + name=column.name, + type=col_type, ) + new_columns.append(new_column) - signature = Function( - args=[ - Function.Argument( - pos=0, - name="other", - type=DataFrameType(columns=[]), - required=True, - ), - ], - returns=DataFrameType(columns=new_columns), - ) + 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()), + ) + ) - return ( - typer._get_call_result( - location=location, - callee=signature, - positional=positional, - keywords=keywords, + 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() ) - or UnknownType() - ) diff --git a/midas/checker/frames.py b/midas/checker/frames.py index 50ae080..dff5f71 100644 --- a/midas/checker/frames.py +++ b/midas/checker/frames.py @@ -4,8 +4,7 @@ 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 FRAME_METHODS, FrameMethod -from midas.checker.registry import TypesRegistry +from midas.checker.frame_methods import Call, MethodResolver from midas.checker.reporter import FileReporter from midas.checker.types import ColumnType, DataFrameType, TupleType, Type, UnknownType @@ -18,8 +17,9 @@ def is_list_of_literals(exprs: list[p.Expr]) -> TypeGuard[list[p.LiteralExpr]]: class FrameManager: - def __init__(self, types: TypesRegistry) -> None: - self.types: TypesRegistry = types + def __init__(self, typer: PythonTyper) -> None: + self.typer: PythonTyper = typer + self.method_resolver: MethodResolver = MethodResolver(self.typer) def assign( self, @@ -137,15 +137,18 @@ class FrameManager: ) -> list[Optional[ColumnType]]: return [cls._get_column(frame, name) for name in names] - def call_method( + def call( self, - typer: PythonTyper, - reporter: FileReporter, + method: str, location: Location, frame: DataFrameType, - method: str, positional: list[TypedExpr], keywords: dict[str, TypedExpr], ) -> Type: - function: FrameMethod = FRAME_METHODS[method] - return function(typer, self, reporter, location, frame, positional, keywords) + call: Call = Call( + location=location, + frame=frame, + positional=positional, + keywords=keywords, + ) + return self.method_resolver.call(method, call) diff --git a/midas/checker/python.py b/midas/checker/python.py index 3a3c547..c0a446f 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -8,7 +8,6 @@ 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.frame_methods import FRAME_METHODS from midas.checker.frames import FrameManager from midas.checker.operators import ( PY_COMPARATOR_METHODS, @@ -76,7 +75,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.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] = {} @@ -527,15 +526,11 @@ class PythonTyper( case p.GetExpr(object=obj, name=method): obj_type: Type = self.type_of(obj) unfolded: Type = unfold_type(obj_type) - if isinstance(unfolded, DataFrameType) and self._is_frame_method( - method - ): - return self.frame_mgr.call_method( - self, - self.reporter, + if isinstance(unfolded, DataFrameType): + return self.frame_mgr.call( + method, expr.location, unfolded, - method, positional, keywords, ) @@ -1307,6 +1302,3 @@ class PythonTyper( self, frame: DataFrameType, expr: p.SubscriptExpr ) -> Type: return self.frame_mgr.get(self.reporter, expr.location, frame, expr.index) - - def _is_frame_method(self, method: str) -> bool: - return method in FRAME_METHODS -- 2.49.1 From 3cdc9031d39883160d183bd5793621b5f003403f Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 25 Jun 2026 22:31:10 +0200 Subject: [PATCH 17/39] refactor: use metaclass to collect frame methods --- midas/checker/frame_methods.py | 60 ++++++++++++++++++---------------- midas/checker/frames.py | 4 +-- 2 files changed, 34 insertions(+), 30 deletions(-) diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py index 2bbc016..96c868c 100644 --- a/midas/checker/frame_methods.py +++ b/midas/checker/frame_methods.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, Optional +from typing import TYPE_CHECKING, Any, Callable, Optional from midas.ast.location import Location from midas.checker.registry import TypesRegistry @@ -19,6 +19,18 @@ 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 @@ -27,26 +39,25 @@ class Call: keywords: dict[str, TypedExpr] -def _is_method(obj: object, method: str) -> bool: - if not callable(obj): - return False - if not hasattr(obj, "__method_names__"): - return False - return method in obj.__method_names__ # type: ignore +class _MethodRegistryMeta(type): + _methods: dict[str, Callable] = {} + + 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 + return new_class -class MethodResolver: - @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 - +class MethodRegistry(metaclass=_MethodRegistryMeta): def __init__(self, typer: PythonTyper) -> None: self.typer: PythonTyper = typer @@ -58,23 +69,16 @@ class MethodResolver: def types(self) -> TypesRegistry: return self.typer.types - def _get_method_by_name(self, method: str) -> Optional[Callable]: - for name in dir(self): - attr = getattr(self, name) - if _is_method(attr, method): - return attr - return None - def call( self, method: str, call: Call, ) -> Type: - func: Optional[Callable] = self._get_method_by_name(method) + func: Optional[Callable] = self._methods.get(method) if func is None: self.reporter.error(call.location, f"Unknown method {method}") return UnknownType() - return func(call) + return func(self, call) @frame_method("add", "__add__") def add( diff --git a/midas/checker/frames.py b/midas/checker/frames.py index dff5f71..da9eb0e 100644 --- a/midas/checker/frames.py +++ b/midas/checker/frames.py @@ -4,7 +4,7 @@ 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, MethodResolver +from midas.checker.frame_methods import Call, MethodRegistry from midas.checker.reporter import FileReporter from midas.checker.types import ColumnType, DataFrameType, TupleType, Type, UnknownType @@ -19,7 +19,7 @@ def is_list_of_literals(exprs: list[p.Expr]) -> TypeGuard[list[p.LiteralExpr]]: class FrameManager: def __init__(self, typer: PythonTyper) -> None: self.typer: PythonTyper = typer - self.method_resolver: MethodResolver = MethodResolver(self.typer) + self.method_resolver: MethodRegistry = MethodRegistry(self.typer) def assign( self, -- 2.49.1 From dbf6f9e2dbaeffe68585ea5b377e7e85afee9540 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 26 Jun 2026 10:27:54 +0200 Subject: [PATCH 18/39] tests: update with reordered argument typing --- .../cases/checker/01_simple_types.py.ref.json | 32 +- tests/cases/checker/03_functions.py.ref.json | 650 +++++++++--------- tests/cases/checker/06_subtyping.py.ref.json | 52 +- tests/cases/checker/07_variance.py.ref.json | 46 +- .../cases/checker/08_unification.py.ref.json | 338 ++++----- tests/serializer/midas.py | 13 + 6 files changed, 586 insertions(+), 545 deletions(-) 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/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..a4387ee 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": true + } + ], + "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/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), + } -- 2.49.1 From d7bb6326de2676cadd358d95f704fe7b8e4b8b1a Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 26 Jun 2026 10:35:50 +0200 Subject: [PATCH 19/39] feat(checker): lookup dunders on dataframes --- midas/checker/python.py | 59 ++++++++++++++++++++++++++++++----------- 1 file changed, 43 insertions(+), 16 deletions(-) diff --git a/midas/checker/python.py b/midas/checker/python.py index c0a446f..a849a62 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -47,6 +47,10 @@ class ReturnException(Exception): pass +class UndefinedMethodException(Exception): + pass + + @dataclass(frozen=True, kw_only=True) class MappedArgument: expr: p.Expr @@ -195,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) @@ -468,20 +502,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: @@ -494,20 +524,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: -- 2.49.1 From 24bffdabd4a4b34ce90989eedd9065ea3642728e Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 26 Jun 2026 11:21:17 +0200 Subject: [PATCH 20/39] fix(checker): type check None literal --- midas/checker/python.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/midas/checker/python.py b/midas/checker/python.py index a849a62..9c21372 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -594,6 +594,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() -- 2.49.1 From 4ea400265c0fc657c1023bdb9cb4ec1b73904dda Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 26 Jun 2026 11:21:38 +0200 Subject: [PATCH 21/39] feat(checker): add mean method on frames --- midas/checker/frame_methods.py | 45 ++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py index 96c868c..0220ba0 100644 --- a/midas/checker/frame_methods.py +++ b/midas/checker/frame_methods.py @@ -10,6 +10,8 @@ from midas.checker.types import ( ColumnType, DataFrameType, Function, + OverloadedFunction, + TopType, Type, UnknownType, unfold_type, @@ -85,6 +87,9 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): 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] = {} @@ -151,3 +156,43 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): ) 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() + ) -- 2.49.1 From d79e1dee18d0bb93024db9a5d2dc3e12f619a698 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 26 Jun 2026 11:55:31 +0200 Subject: [PATCH 22/39] fix(checker): change heterogeneous errors to warnings --- midas/checker/python.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/midas/checker/python.py b/midas/checker/python.py index 9c21372..3917e36 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -667,7 +667,7 @@ 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: [{', '.join(map(str, item_types))}]", ) @@ -699,7 +699,7 @@ 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: [{', '.join(map(str, key_types))}]", ) @@ -707,7 +707,7 @@ class PythonTyper( if len(value_types) == 1: value_type = value_types[0] else: - self.reporter.error( + self.reporter.warning( expr.location, f"Heterogeneous dict values: [{', '.join(map(str, value_types))}]", ) -- 2.49.1 From db4019835780eeb5599102575a09abefdf9aa42c Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 26 Jun 2026 14:56:15 +0200 Subject: [PATCH 23/39] feat(gen): generate asserts for dataframes and columns --- midas/checker/python.py | 6 ++ midas/generator/generator.py | 144 +++++++++++++++++++++++++++++++++-- 2 files changed, 145 insertions(+), 5 deletions(-) diff --git a/midas/checker/python.py b/midas/checker/python.py index 3917e36..4d6825e 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -1306,6 +1306,12 @@ 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" diff --git a/midas/generator/generator.py b/midas/generator/generator.py index a7641b5..96c2c2f 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -40,6 +40,9 @@ class Scope: 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" @@ -58,12 +61,24 @@ 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]] = [] + self.define_is_dataframe: bool = False + self.define_is_column: bool = False + def generate_ast(self, typed_ast: TypedAST, src_path: Path) -> ast.AST: self.rel_src_path = src_path.resolve().relative_to(self.workdir) 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 @@ -350,6 +365,51 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): for item, item_type in zip(expr.elts, items): self._make_cast_asserts(src_location, item, item_type) + case DataFrameType(columns=columns): + self.define_is_dataframe = True + self._add_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: + self._add_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}" + ), + ) + self._make_cast_asserts( + src_location, + ast.Subscript( + value=expr, slice=ast.Constant(value=column.name) + ), + column.type, + ) + + case ColumnType(type=inner): + self.define_is_column = True + self._add_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" + ), + ) + # TODO: check value type + case ( TopType() | Function() @@ -357,8 +417,6 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): | ComplexType() | ExtensionType() | GenericType() - | ColumnType() - | DataFrameType() ): self.logger.warning(f"Can't make assertion for type {type}") @@ -367,7 +425,11 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): 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" @@ -385,7 +447,7 @@ 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 ''}"), ] ) @@ -421,3 +483,75 @@ 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"), + ) -- 2.49.1 From 75bd203d4aa4913ce91914fe6fa2b06d10a29d2c Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 11:01:53 +0200 Subject: [PATCH 24/39] fix(checker): allow calling unknown method on dataframes --- midas/checker/frame_methods.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py index 0220ba0..8a7dfbc 100644 --- a/midas/checker/frame_methods.py +++ b/midas/checker/frame_methods.py @@ -42,7 +42,7 @@ class Call: class _MethodRegistryMeta(type): - _methods: dict[str, Callable] = {} + _methods: dict[str, Callable[..., Type]] = {} def __new__( cls, @@ -55,7 +55,7 @@ class _MethodRegistryMeta(type): 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 + new_class._methods[name] = attr # type: ignore return new_class @@ -76,9 +76,9 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): method: str, call: Call, ) -> Type: - func: Optional[Callable] = self._methods.get(method) + func: Optional[Callable[..., Type]] = self._methods.get(method) if func is None: - self.reporter.error(call.location, f"Unknown method {method}") + self.reporter.warning(call.location, f"Unknown method {method}") return UnknownType() return func(self, call) -- 2.49.1 From 7e0319906a1771a2eb72d769f300829089bf24fa Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 11:04:25 +0200 Subject: [PATCH 25/39] feat(gen): assertions for column values --- midas/generator/generator.py | 202 +++++++++++++++++++++-------------- 1 file changed, 124 insertions(+), 78 deletions(-) diff --git a/midas/generator/generator.py b/midas/generator/generator.py index 96c2c2f..249f9ff 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -35,8 +35,8 @@ 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]): @@ -159,7 +159,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 @@ -294,121 +296,152 @@ 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): - self._add_assert( - ast.Call( - func=ast.Name(id="isinstance"), - args=[expr, ast.Name(id="tuple")], - keywords=[], + 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), ), - self._make_cast_assert_message(src_location, expr, type), - ) + ] assert isinstance(expr, ast.Tuple) for item, item_type in zip(expr.elts, items): - self._make_cast_asserts(src_location, item, item_type) + asserts.extend( + self._make_cast_asserts(src_location, item, item_type) + ) + return asserts case DataFrameType(columns=columns): self.define_is_dataframe = True - self._add_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: - self._add_assert( - ast.Compare( - left=ast.Constant(value=column.name), - ops=[ast.In()], - comparators=[expr], + 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, f": Missing column {column.name}" + 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}", + ), + ) ) - self._make_cast_asserts( - src_location, - ast.Subscript( - value=expr, slice=ast.Constant(value=column.name) - ), - column.type, + asserts.extend( + self._make_cast_asserts( + src_location, + ast.Subscript( + value=expr, slice=ast.Constant(value=column.name) + ), + column.type, + ) ) + return asserts - case ColumnType(type=inner): + case ColumnType(): self.define_is_column = True - self._add_assert( - ast.Call( - func=ast.Name(id=self.IS_COLUMN_FUNC), - args=[expr], - keywords=[], + 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" + ), ), - self._make_cast_assert_message( - src_location, expr, type, ": Not a column" - ), - ) - # TODO: check value type + ] + asserts.append(self._make_column_inner_assert(src_location, expr, type)) + return asserts case ( TopType() @@ -419,6 +452,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): | GenericType() ): self.logger.warning(f"Can't make assertion for type {type}") + return [] # Ensure exhaustiveness case _: @@ -453,9 +487,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): 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], @@ -555,3 +589,15 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): decorator_list=[], returns=ast.Name(id="bool"), ) + + def _make_column_inner_assert( + self, src_location: Location, column: ast.expr, type: ColumnType + ) -> ast.stmt: + # TODO: improve message, maybe chain contexts + col: ast.expr = ast.Name(id="col") + return ast.For( + target=col, + iter=column, + body=self._make_cast_asserts(src_location, col, type.type), + orelse=[], + ) -- 2.49.1 From 3392bc347dc6a25dd2dec7372101f615ddb540af Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 11:06:35 +0200 Subject: [PATCH 26/39] fix(checker): allow subtypes and unknown as if test --- midas/checker/python.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/midas/checker/python.py b/midas/checker/python.py index 4d6825e..8eff54a 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -439,8 +439,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}" ) @@ -638,7 +640,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}" ) -- 2.49.1 From 6f6245d28357e1c24179c666191c131a0b160d12 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 11:13:47 +0200 Subject: [PATCH 27/39] fix(checker): allow iterating on unknown --- midas/checker/python.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/midas/checker/python.py b/midas/checker/python.py index 8eff54a..33699cb 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -458,13 +458,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) @@ -1155,9 +1158,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 -- 2.49.1 From 64f4314f0d9fd72b6895ebbb6bb41532383a2f10 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 11:19:26 +0200 Subject: [PATCH 28/39] fix(gen): prevent empty loop for column asserts --- midas/generator/generator.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/midas/generator/generator.py b/midas/generator/generator.py index 249f9ff..2aaaea2 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -440,7 +440,11 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): ), ), ] - asserts.append(self._make_column_inner_assert(src_location, expr, type)) + 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 ( @@ -592,12 +596,15 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): def _make_column_inner_assert( self, src_location: Location, column: ast.expr, type: ColumnType - ) -> ast.stmt: + ) -> 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=self._make_cast_asserts(src_location, col, type.type), + body=body, orelse=[], ) -- 2.49.1 From 256536562fb86d3dfa42366484cb49bcfe830f16 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 13:59:03 +0200 Subject: [PATCH 29/39] fix(parser): parse empty calls --- midas/parser/midas.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/midas/parser/midas.py b/midas/parser/midas.py index 77ad40b..3bcee32 100644 --- a/midas/parser/midas.py +++ b/midas/parser/midas.py @@ -399,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() -- 2.49.1 From e324f414e6999e4ef4c08ca97c347ff424586542 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 14:00:37 +0200 Subject: [PATCH 30/39] feat(checker): type check tuple instantiation in Midas --- midas/checker/builtins.py | 4 +++- midas/checker/registry.py | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/midas/checker/builtins.py b/midas/checker/builtins.py index 78ba18f..2702381 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", "tuple"}, "float": {"int"}, "int": {"bool"}, } @@ -32,6 +32,8 @@ def define_builtins(reg: TypesRegistry): 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/registry.py b/midas/checker/registry.py index 0ef2d76..02c96c3 100644 --- a/midas/checker/registry.py +++ b/midas/checker/registry.py @@ -18,6 +18,7 @@ from midas.checker.types import ( OverloadedFunction, Predicate, TopType, + TupleType, Type, TypeVar, UnknownType, @@ -346,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") -- 2.49.1 From a5fe985eb207ad4e831d835460545f0accc04b46 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 14:01:33 +0200 Subject: [PATCH 31/39] feat(checker): add methods on str --- midas/checker/builtins.midas | 98 +++++++++++++++++++++++++++++++++++- midas/checker/builtins.py | 3 +- 2 files changed, 99 insertions(+), 2 deletions(-) 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 2702381..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", "tuple"}, + "object": {"float", "list", "dict", "str", "bytes", "tuple"}, "float": {"int"}, "int": {"bool"}, } @@ -26,6 +26,7 @@ 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")) -- 2.49.1 From cb75878ae9dd5d466a3b9425670f9bd8eec5d868 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 14:02:29 +0200 Subject: [PATCH 32/39] fix(checker): allow some assignments to unknown --- midas/checker/python.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/midas/checker/python.py b/midas/checker/python.py index 33699cb..4b714e6 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -413,13 +413,16 @@ class PythonTyper( 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 var_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, @@ -582,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() -- 2.49.1 From b36896cc7b3fab92a2e922dcdc2677d8b3288e07 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 14:02:52 +0200 Subject: [PATCH 33/39] feat(checker): add `len()` --- midas/checker/preamble.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) 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) -- 2.49.1 From 3f99563ac8173fef4fe41007f57604b791259564 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 14:24:38 +0200 Subject: [PATCH 34/39] feat: handle multi-parameter generic in Python --- gen/python.py | 2 +- midas/ast/printer.py | 8 +++++++- midas/ast/python.py | 2 +- midas/checker/python.py | 6 +++--- midas/cli/highlighter.py | 6 +++--- midas/cli/hl_python.css | 2 +- midas/parser/python.py | 16 +++++++++------- tests/serializer/python.py | 2 +- 8 files changed, 26 insertions(+), 18 deletions(-) diff --git a/gen/python.py b/gen/python.py index df83e6f..56a5f41 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: diff --git a/midas/ast/printer.py b/midas/ast/printer.py index d713a74..82cf4b3 100644 --- a/midas/ast/printer.py +++ b/midas/ast/printer.py @@ -560,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") diff --git a/midas/ast/python.py b/midas/ast/python.py index 20d7279..b67a043 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) diff --git a/midas/checker/python.py b/midas/checker/python.py index 4b714e6..26b6301 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -761,9 +761,9 @@ 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: diff --git a/midas/cli/highlighter.py b/midas/cli/highlighter.py index c131143..bced867 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") 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/parser/python.py b/midas/parser/python.py index 7839f52..9868229 100644 --- a/midas/parser/python.py +++ b/midas/parser/python.py @@ -300,26 +300,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 +347,7 @@ class PythonParser: return BaseType( location=loc, base="None", - param=None, + args=(), ) case _: diff --git a/tests/serializer/python.py b/tests/serializer/python.py index 3904739..57e47c9 100644 --- a/tests/serializer/python.py +++ b/tests/serializer/python.py @@ -98,7 +98,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: -- 2.49.1 From 9db56adf5676b8f30fb25a4de1458dc87f3d5427 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 14:35:31 +0200 Subject: [PATCH 35/39] feat: add Python tuple expression --- gen/python.py | 4 ++++ midas/ast/printer.py | 11 +++++++++++ midas/ast/python.py | 11 +++++++++++ midas/checker/python.py | 5 +++++ midas/checker/resolver.py | 4 ++++ midas/cli/highlighter.py | 4 ++++ midas/parser/python.py | 7 +++++++ tests/serializer/python.py | 6 ++++++ 8 files changed, 52 insertions(+) diff --git a/gen/python.py b/gen/python.py index 56a5f41..73f64c7 100644 --- a/gen/python.py +++ b/gen/python.py @@ -174,6 +174,10 @@ class SliceExpr: step: Optional[Expr] +class TupleExpr: + items: tuple[Expr, ...] + + class RawExpr: expr: ast.expr diff --git a/midas/ast/printer.py b/midas/ast/printer.py index 82cf4b3..95c4109 100644 --- a/midas/ast/printer.py +++ b/midas/ast/printer.py @@ -879,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 b67a043..6202fcc 100644 --- a/midas/ast/python.py +++ b/midas/ast/python.py @@ -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/python.py b/midas/checker/python.py index 26b6301..be3d05f 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -750,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() diff --git a/midas/checker/resolver.py b/midas/checker/resolver.py index 753f4db..43768bd 100644 --- a/midas/checker/resolver.py +++ b/midas/checker/resolver.py @@ -236,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/cli/highlighter.py b/midas/cli/highlighter.py index bced867..8d142f1 100644 --- a/midas/cli/highlighter.py +++ b/midas/cli/highlighter.py @@ -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: ... diff --git a/midas/parser/python.py b/midas/parser/python.py index 9868229..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, @@ -479,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/tests/serializer/python.py b/tests/serializer/python.py index 57e47c9..7c29616 100644 --- a/tests/serializer/python.py +++ b/tests/serializer/python.py @@ -302,6 +302,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", -- 2.49.1 From bccc96e4d09b738c7437be151999adbe37434143 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 22:41:47 +0200 Subject: [PATCH 36/39] fix: minor fixes --- midas/checker/python.py | 2 +- tests/serializer/python.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/midas/checker/python.py b/midas/checker/python.py index be3d05f..c149d8d 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -1233,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 diff --git a/tests/serializer/python.py b/tests/serializer/python.py index 7c29616..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, -- 2.49.1 From 293953a078f8a23601eed6cd05c71d2aa407780c Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 22:43:08 +0200 Subject: [PATCH 37/39] tests: update with multi-parameter generics --- .../cases/checker/04_custom_types.py.ref.json | 4 +- tests/cases/checker/07_variance.py.ref.json | 2 +- .../python-parser/01_simple_types.py.ref.json | 14 +-- .../python-parser/02_custom_types.py.ref.json | 50 ++++++----- .../python-parser/03_functions.py.ref.json | 86 ++++++++++--------- 5 files changed, 85 insertions(+), 71 deletions(-) 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/07_variance.py.ref.json b/tests/cases/checker/07_variance.py.ref.json index a4387ee..abfebf3 100644 --- a/tests/cases/checker/07_variance.py.ref.json +++ b/tests/cases/checker/07_variance.py.ref.json @@ -317,7 +317,7 @@ "pos": 0, "name": "object", "type": {}, - "required": true + "required": false } ], "args": [], 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 } -- 2.49.1 From d14f20889776bfec4533ed7e2142f6fb578153da Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 29 Jun 2026 22:44:39 +0200 Subject: [PATCH 38/39] feat(gen): add tuple expr to generator --- midas/generator/generator.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/midas/generator/generator.py b/midas/generator/generator.py index 2aaaea2..8476392 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -196,6 +196,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 -- 2.49.1 From ffccc1bedde579890a53e4ecf28a403f1ce2ec27 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 30 Jun 2026 11:19:17 +0200 Subject: [PATCH 39/39] feat(cli): generate stubs in build dir when compiling --- midas/cli/commands/compile.py | 13 ++++++++---- midas/cli/commands/stubs.py | 10 +++++---- midas/generator/generator.py | 39 +++++++++++++++++++++++++++++------ tests/generator.py | 3 ++- 4 files changed, 50 insertions(+), 15 deletions(-) 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/generator/generator.py b/midas/generator/generator.py index 8476392..c2f46ee 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -9,6 +9,7 @@ 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, @@ -30,6 +31,7 @@ from midas.checker.types import ( UnknownType, ) from midas.generator.constraints import ConstraintGenerator +from midas.generator.stubs import StubsGenerator from midas.utils import TypedAST @@ -64,8 +66,10 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): self.define_is_dataframe: bool = False self.define_is_column: bool = False - def generate_ast(self, typed_ast: TypedAST, src_path: Path) -> ast.AST: - self.rel_src_path = src_path.resolve().relative_to(self.workdir) + 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() @@ -83,10 +87,13 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): 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) @@ -98,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), 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 -- 2.49.1