From a59a58d21ade0976a45028a4e7eabff63b6f6489 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Wed, 1 Jul 2026 14:43:30 +0200 Subject: [PATCH 01/36] feat(gen): generate alias stubs --- midas/generator/stubs.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/midas/generator/stubs.py b/midas/generator/stubs.py index d0a9b23..5f71aac 100644 --- a/midas/generator/stubs.py +++ b/midas/generator/stubs.py @@ -91,6 +91,21 @@ class StubsGenerator: def generate_stub(self, name: str, type: Type): base_type: Type = type + # TODO: improve + match type: + case DerivedType(name=name_) | GenericType(name=name_) if name_ == name: + pass + case UnitType() if name == "None": + pass + case TopType() if name == "Any": + pass + case _: + alias = ast.Assign( + targets=[ast.Name(id=name)], value=self.dump_type(type) + ) + self.add_stub(alias) + return + members: dict[str, Member] = self.types._members.get(name, {}) if isinstance(base_type, (BaseType, TopType, UnitType)) and len(members) == 0: return From ad2fabf471539d518e130a0c06e98aa6f78caf15 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Wed, 1 Jul 2026 22:32:13 +0200 Subject: [PATCH 02/36] feat(checker): add assertion collector --- midas/checker/frame_methods.py | 5 +++++ midas/checker/python.py | 2 ++ midas/generator/collector.py | 22 ++++++++++++++++++++++ 3 files changed, 29 insertions(+) create mode 100644 midas/generator/collector.py diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py index 0f85493..2afadfd 100644 --- a/midas/checker/frame_methods.py +++ b/midas/checker/frame_methods.py @@ -18,6 +18,7 @@ from midas.checker.types import ( UnknownType, unfold_type, ) +from midas.generator.collector import AssertionCollector if TYPE_CHECKING: from midas.checker.python import PythonTyper, TypedExpr @@ -77,6 +78,10 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): def dispatcher(self) -> CallDispatcher[p.Expr]: return self.typer.dispatcher + @property + def assertions(self) -> AssertionCollector: + return self.typer.assertions + def call( self, method: str, diff --git a/midas/checker/python.py b/midas/checker/python.py index d74cbce..70cb4c9 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -36,6 +36,7 @@ from midas.checker.types import ( Variance, unfold_type, ) +from midas.generator.collector import AssertionCollector from midas.parser.python import PythonParser from midas.utils import TypedAST @@ -87,6 +88,7 @@ class PythonTyper( self.dispatcher: CallDispatcher[p.Expr] = CallDispatcher[p.Expr]( self.types, self.reporter ) + self.assertions: AssertionCollector = AssertionCollector() def set_reporter(self, reporter: FileReporter): self.reporter = reporter diff --git a/midas/generator/collector.py b/midas/generator/collector.py new file mode 100644 index 0000000..38906fc --- /dev/null +++ b/midas/generator/collector.py @@ -0,0 +1,22 @@ +import ast + +import midas.ast.python as p + + +class AssertionCollector: + def __init__(self): + self.assertions: list[tuple[p.Expr, list[ast.expr]]] = [] + self.definitions: dict[str, ast.stmt] = {} + + def add(self, assertion): + self.assertions.append(assertion) + + def define(self, name: str, stmt: ast.stmt): + if name not in self.definitions: + self.definitions[name] = stmt + + def get_definitions(self) -> list[ast.stmt]: + return list(self.definitions.values()) + + def get_assertions(self) -> list[tuple[p.Expr, list[ast.expr]]]: + return self.assertions From f7a36f61b648e14bc8d6d3291416df453e5edac5 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Wed, 1 Jul 2026 22:34:02 +0200 Subject: [PATCH 03/36] fix(checker): pass AST expression to method registry --- midas/checker/frame_methods.py | 10 ++++++++++ midas/checker/frames.py | 2 ++ midas/checker/python.py | 26 ++++++++++++++++---------- 3 files changed, 28 insertions(+), 10 deletions(-) diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py index 2afadfd..c2935f8 100644 --- a/midas/checker/frame_methods.py +++ b/midas/checker/frame_methods.py @@ -40,6 +40,7 @@ def frame_method(*names: str): class Call: location: Location frame: DataFrameType + frame_expr: p.Expr positional: list[TypedExpr] keywords: dict[str, TypedExpr] @@ -105,6 +106,7 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): by_name: dict[str, DataFrameType.Column] = {} frame2: Optional[DataFrameType] = None + # Get map of operand's columns by name, if there is at least 1 operand, which is a dataframe if len(call.positional) != 0: other: Type = call.positional[0][1] unfolded_other: Type = unfold_type(other) @@ -114,6 +116,10 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): col.name: col for col in frame2.columns if col.name is not None } + # Compute new schema: + # Step 1: for all columns in frame1: + # - if present in frame2 with equivalent type -> add to schema as is + # - if not -> add to schema as unknown in_frame1: set[str] = set() for column in call.frame.columns: if column.name is not None: @@ -134,6 +140,8 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): ) new_columns.append(new_column) + # Step 2: for all columns in frame2 + # - if not in frame1 -> add to schema as unknown if frame2 is not None: for column in frame2.columns: if column.name in in_frame1: @@ -146,6 +154,7 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): ) ) + # Build signature with new schema and generic operand signature = Function( args=[ Function.Argument( @@ -158,6 +167,7 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): returns=DataFrameType(columns=new_columns), ) + # Map arguments and compute result type result: CallResult = self.dispatcher.get_result( location=call.location, callee=signature, diff --git a/midas/checker/frames.py b/midas/checker/frames.py index da9eb0e..56c34c6 100644 --- a/midas/checker/frames.py +++ b/midas/checker/frames.py @@ -142,12 +142,14 @@ class FrameManager: method: str, location: Location, frame: DataFrameType, + frame_expr: p.Expr, positional: list[TypedExpr], keywords: dict[str, TypedExpr], ) -> Type: call: Call = Call( location=location, frame=frame, + frame_expr=frame_expr, positional=positional, keywords=keywords, ) diff --git a/midas/checker/python.py b/midas/checker/python.py index 70cb4c9..374f736 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -211,23 +211,24 @@ class PythonTyper( def call_method( self, location: Location, - obj: Type, + obj: TypedExpr, method_name: str, positional: list[TypedExpr], keywords: dict[str, TypedExpr], ) -> Optional[Type]: - unfolded: Type = unfold_type(obj) + unfolded: Type = unfold_type(obj[1]) match unfolded: case DataFrameType(): return self.frame_mgr.call( method=method_name, location=location, frame=unfolded, + frame_expr=obj[0], positional=positional, keywords=keywords, ) - method: Optional[Type] = self.types.lookup_member(obj, method_name) + method: Optional[Type] = self.types.lookup_member(obj[1], method_name) if method is None: raise UndefinedMethodException @@ -522,7 +523,9 @@ class PythonTyper( result: Optional[Type] try: - result = self.call_method(location, left, method, [(right_expr, right)], {}) + result = self.call_method( + location, (left_expr, left), method, [(right_expr, right)], {} + ) except UndefinedMethodException: self.reporter.error( location, @@ -545,7 +548,9 @@ class PythonTyper( result: Optional[Type] try: - result = self.call_method(expr.location, operand, method, [], {}) + result = self.call_method( + expr.location, (expr.right, operand), method, [], {} + ) except UndefinedMethodException: self.reporter.error( expr.location, @@ -573,11 +578,12 @@ class PythonTyper( unfolded: Type = unfold_type(obj_type) if isinstance(unfolded, DataFrameType): return self.frame_mgr.call( - method, - expr.location, - unfolded, - positional, - keywords, + method=method, + location=expr.location, + frame=unfolded, + frame_expr=obj, + positional=positional, + keywords=keywords, ) callee: Type = self.type_of(expr.callee) From bd4d793ce0a3fa89df1b83a70f30ea673bde1048 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 17:08:43 +0200 Subject: [PATCH 04/36] feat(gen): add Assertion class --- midas/generator/collector.py | 39 ++++++++++++++++++++++++++++++++---- 1 file changed, 35 insertions(+), 4 deletions(-) diff --git a/midas/generator/collector.py b/midas/generator/collector.py index 38906fc..2b96d03 100644 --- a/midas/generator/collector.py +++ b/midas/generator/collector.py @@ -1,15 +1,43 @@ import ast +from dataclasses import dataclass +from typing import Callable import midas.ast.python as p +AssertionBuilder = Callable[..., ast.expr] + + +@dataclass +class Assertion: + bound_expr: p.Expr + inputs: list[p.Expr] + builder: AssertionBuilder + message: str + + def is_bound_to(self, expr: p.Expr) -> bool: + return expr == self.bound_expr + class AssertionCollector: def __init__(self): - self.assertions: list[tuple[p.Expr, list[ast.expr]]] = [] + self.assertions: list[Assertion] = [] self.definitions: dict[str, ast.stmt] = {} - def add(self, assertion): - self.assertions.append(assertion) + def add( + self, + bound_expr: p.Expr, + inputs: list[p.Expr], + builder: AssertionBuilder, + message: str, + ): + self.assertions.append( + Assertion( + bound_expr=bound_expr, + inputs=inputs, + builder=builder, + message=message, + ) + ) def define(self, name: str, stmt: ast.stmt): if name not in self.definitions: @@ -18,5 +46,8 @@ class AssertionCollector: def get_definitions(self) -> list[ast.stmt]: return list(self.definitions.values()) - def get_assertions(self) -> list[tuple[p.Expr, list[ast.expr]]]: + def get_assertions(self) -> list[Assertion]: return self.assertions + + def get_assertions_for(self, expr: p.Expr) -> list[Assertion]: + return list(filter(lambda a: a.is_bound_to(expr), self.assertions)) From 47b2dfdd7307d9eb89d2b341ee7d4082e7c7ebaf Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 17:09:50 +0200 Subject: [PATCH 05/36] feat(gen): add assertion collector to TypedAST --- midas/checker/python.py | 1 + midas/generator/generator.py | 2 ++ midas/utils.py | 2 ++ 3 files changed, 5 insertions(+) diff --git a/midas/checker/python.py b/midas/checker/python.py index 374f736..a42115f 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -115,6 +115,7 @@ class PythonTyper( stmts=stmts, judgements=self.judgements, evaluated_casts=self.evaluated_casts, + assertions=self.assertions, ) def judge(self, expr: p.Expr, type: Type): diff --git a/midas/generator/generator.py b/midas/generator/generator.py index c2f46ee..a8b3145 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -30,6 +30,7 @@ from midas.checker.types import ( UnitType, UnknownType, ) +from midas.generator.collector import Assertion, AssertionCollector from midas.generator.constraints import ConstraintGenerator from midas.generator.stubs import StubsGenerator from midas.utils import TypedAST @@ -55,6 +56,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): stmts=[], judgements=[], evaluated_casts=[], + assertions=AssertionCollector(), ) self._alias_count: int = 0 self._predicate_count: int = 0 diff --git a/midas/utils.py b/midas/utils.py index 0ac90cf..c225a89 100644 --- a/midas/utils.py +++ b/midas/utils.py @@ -3,6 +3,7 @@ from typing import Any, Callable, Optional import midas.ast.python as p from midas.checker.types import Type +from midas.generator.collector import AssertionCollector AllowRepeat = Callable[[object], bool] @@ -63,3 +64,4 @@ class TypedAST: stmts: list[p.Stmt] judgements: list[tuple[p.Expr, Type]] evaluated_casts: list[p.CastExpr] + assertions: AssertionCollector From 8df01afd8ce1e4b4d51346d914f40bf7d290eb44 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 17:10:18 +0200 Subject: [PATCH 06/36] feat(gen): materialize assertions from collector --- midas/generator/collector.py | 6 ++ midas/generator/generator.py | 116 ++++++++++++++++++++++++----------- 2 files changed, 85 insertions(+), 37 deletions(-) diff --git a/midas/generator/collector.py b/midas/generator/collector.py index 2b96d03..fcaf545 100644 --- a/midas/generator/collector.py +++ b/midas/generator/collector.py @@ -39,6 +39,12 @@ class AssertionCollector: ) ) + def remove(self, assertion: Assertion): + try: + self.assertions.remove(assertion) + except ValueError: + pass + def define(self, name: str, stmt: ast.stmt): if name not in self.definitions: self.definitions[name] = stmt diff --git a/midas/generator/generator.py b/midas/generator/generator.py index a8b3145..bfc1d3c 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -61,6 +61,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): self._alias_count: int = 0 self._predicate_count: int = 0 self._scopes: list[Scope] = [] + self._aliases: list[tuple[p.Expr, ast.expr]] = [] self._constraint_generator: ConstraintGenerator = ConstraintGenerator(types) self._constraints: list[tuple[m.Expr, ast.expr]] = [] @@ -73,7 +74,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): def generate_ast(self, typed_ast: TypedAST) -> ast.AST: self._typed_ast = typed_ast - body: list[ast.stmt] = self._visit_body(typed_ast.stmts) + body: list[ast.stmt] = self._visit_body(typed_ast.stmts, can_be_empty=True) predicates: list[ast.stmt] = self._constraint_generator.get_definitions() body = predicates + body @@ -131,39 +132,48 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): output: str = ast.unparse(module) out_path.write_text(output) + def convert(self, expr: p.Expr) -> ast.expr: + for expr2, alias in self._aliases: + if expr2 == expr: + return alias + assertions = self._typed_ast.assertions.get_assertions_for(expr) + if len(assertions) != 0: + return self._apply_assertions(expr, assertions) + return expr.accept(self) + def visit_binary_expr(self, expr: p.BinaryExpr) -> ast.expr: return ast.BinOp( - left=expr.left.accept(self), + left=self.convert(expr.left), op=expr.operator, - right=expr.right.accept(self), + right=self.convert(expr.right), ) def visit_compare_expr(self, expr: p.CompareExpr) -> ast.expr: return ast.Compare( - left=expr.left.accept(self), + left=self.convert(expr.left), ops=[expr.operator], - comparators=[expr.right.accept(self)], + comparators=[self.convert(expr.right)], ) def visit_unary_expr(self, expr: p.UnaryExpr) -> ast.expr: return ast.UnaryOp( op=expr.operator, - operand=expr.right.accept(self), + operand=self.convert(expr.right), ) def visit_call_expr(self, expr: p.CallExpr) -> ast.expr: return ast.Call( - func=expr.callee.accept(self), - args=[arg.accept(self) for arg in expr.arguments], + func=self.convert(expr.callee), + args=[self.convert(arg) for arg in expr.arguments], keywords=[ - ast.keyword(arg=name, value=arg.accept(self)) + ast.keyword(arg=name, value=self.convert(arg)) for name, arg in expr.keywords.items() ], ) def visit_get_expr(self, expr: p.GetExpr) -> ast.expr: return ast.Attribute( - value=expr.object.accept(self), + value=self.convert(expr.object), attr=expr.name, ) @@ -176,16 +186,16 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): def visit_logical_expr(self, expr: p.LogicalExpr) -> ast.expr: return ast.BoolOp( op=expr.operator, - values=[expr.left.accept(self), expr.right.accept(self)], + values=[self.convert(expr.left), self.convert(expr.right)], ) def visit_cast_expr(self, expr: p.CastExpr) -> ast.expr: - expr2: ast.expr = expr.expr.accept(self) + expr2: ast.expr = self.convert(expr.expr) if expr in self._typed_ast.evaluated_casts or expr.unsafe: return expr2 - alias: ast.expr = self._make_alias(expr2) + alias: ast.expr = self._make_alias(expr.expr, expr2) type: Type = self._get_expr_type(expr) asserts: list[ast.stmt] = self._make_cast_asserts(expr.location, alias, type) @@ -196,38 +206,38 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): def visit_ternary_expr(self, expr: p.TernaryExpr) -> ast.expr: return ast.IfExp( - test=expr.test.accept(self), - body=expr.if_true.accept(self), - orelse=expr.if_false.accept(self), + test=self.convert(expr.test), + body=self.convert(expr.if_true), + orelse=self.convert(expr.if_false), ) def visit_list_expr(self, expr: p.ListExpr) -> ast.expr: return ast.List( - elts=[item.accept(self) for item in expr.items], + elts=[self.convert(item) for item in expr.items], ) def visit_dict_expr(self, expr: p.DictExpr) -> ast.expr: return ast.Dict( - keys=[key.accept(self) if key is not None else None for key in expr.keys], - values=[value.accept(self) for value in expr.values], + keys=[self.convert(key) if key is not None else None for key in expr.keys], + values=[self.convert(value) for value in expr.values], ) def visit_subscript_expr(self, expr: p.SubscriptExpr) -> ast.expr: return ast.Subscript( - value=expr.object.accept(self), - slice=expr.index.accept(self), + value=self.convert(expr.object), + slice=self.convert(expr.index), ) def visit_slice_expr(self, expr: p.SliceExpr) -> ast.expr: return ast.Slice( - lower=expr.lower.accept(self) if expr.lower is not None else None, - upper=expr.upper.accept(self) if expr.upper is not None else None, - step=expr.step.accept(self) if expr.step is not None else None, + lower=self.convert(expr.lower) if expr.lower is not None else None, + upper=self.convert(expr.upper) if expr.upper is not None else None, + step=self.convert(expr.step) 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], + elts=[self.convert(item) for item in expr.items], ) def visit_raw_expr(self, expr: p.RawExpr) -> ast.expr: @@ -235,7 +245,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> ast.stmt: return ast.Expr( - value=stmt.expr.accept(self), + value=self.convert(stmt.expr), ) def visit_function(self, stmt: p.Function) -> ast.stmt: @@ -248,12 +258,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): kwonlyargs=[ast.arg(arg=arg.name) for arg in stmt.kwonlyargs], kwarg=None, defaults=[ - arg.default.accept(self) + self.convert(arg.default) for arg in stmt.posonlyargs + stmt.args if arg.default is not None ], kw_defaults=[ - arg.default.accept(self) if arg.default is not None else None + self.convert(arg.default) if arg.default is not None else None for arg in stmt.kwonlyargs ], ), @@ -267,20 +277,20 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): def visit_assign_stmt(self, stmt: p.AssignStmt) -> ast.stmt: return ast.Assign( - targets=[target.accept(self) for target in stmt.targets], - value=stmt.value.accept(self), + targets=[self.convert(target) for target in stmt.targets], + value=self.convert(stmt.value), ) def visit_return_stmt(self, stmt: p.ReturnStmt) -> ast.stmt: return ast.Return( - value=stmt.value.accept(self) if stmt.value is not None else None, + value=self.convert(stmt.value) if stmt.value is not None else None, ) def visit_if_stmt(self, stmt: p.IfStmt) -> ast.stmt: return ast.If( - test=stmt.test.accept(self), + test=self.convert(stmt.test), body=self._visit_body(stmt.body), - orelse=self._visit_body(stmt.orelse), + orelse=self._visit_body(stmt.orelse, can_be_empty=True), ) def visit_pass(self, stmt: p.Pass) -> ast.stmt: @@ -288,8 +298,8 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): def visit_for_stmt(self, stmt: p.ForStmt) -> ast.stmt: return ast.For( - target=stmt.target.accept(self), - iter=stmt.iterator.accept(self), + target=self.convert(stmt.target), + iter=self.convert(stmt.iterator), body=self._visit_body(stmt.body), orelse=[], ) @@ -297,7 +307,9 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): def visit_raw_stmt(self, stmt: p.RawStmt) -> ast.stmt: return stmt.stmt - def _visit_body(self, stmts: list[p.Stmt]) -> list[ast.stmt]: + def _visit_body( + self, stmts: list[p.Stmt], can_be_empty: bool = False + ) -> list[ast.stmt]: generated: list[ast.stmt] = [] for stmt in stmts: scope = Scope() @@ -315,9 +327,11 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): # Remove redundant pass statements if len(generated) > 1: generated = [stmt for stmt in generated if not isinstance(stmt, ast.Pass)] + if len(generated) == 0 and not can_be_empty: + generated = [ast.Pass()] return generated - def _make_alias(self, expr: ast.expr) -> ast.expr: + def _make_alias(self, node: p.Expr, expr: ast.expr) -> ast.expr: name: str = f"__midas_a{self._alias_count}__" alias = ast.Name(id=name) self._alias_count += 1 @@ -328,6 +342,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): value=expr, ) ) + self._aliases.append((node, alias)) return alias def _build_assert(self, expr: ast.expr, message: str | ast.expr) -> ast.stmt: @@ -642,3 +657,30 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): body=body, orelse=[], ) + + def _convert_assertion(self, assertion: Assertion) -> ast.stmt: + inputs: list[ast.expr] = [] + + for input in assertion.inputs: + converted: ast.expr = self.convert(input) + alias: ast.expr = self._make_alias(input, converted) + inputs.append(alias) + + test: ast.expr = assertion.builder(*inputs) + location: Location = assertion.bound_expr.location + loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}" + return self._build_assert( + test, f"{loc_str}: AssertionError: {assertion.message}" + ) + + def _apply_assertions(self, expr: p.Expr, assertions: list[Assertion]) -> ast.expr: + for assertion in assertions: + assert_stmt: ast.stmt + assert_stmt = self._convert_assertion(assertion) + self._add_assert(assert_stmt) + + # Mutating list in frozen dataclass + # Not ideal but easiest way to avoid duplicate assertions + self._typed_ast.assertions.remove(assertion) + + return expr.accept(self) From ff69b651712ff098afd454620d91431a9b6dcac9 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 17:14:05 +0200 Subject: [PATCH 07/36] feat(checker): add same length assertion on frames safely adding two dataframes is only possible if the sizes are the same, or null values could be added dynamically to pad the shortest dataframe --- midas/checker/frame_methods.py | 54 ++++++++++++++++++++++++++++++++++ midas/checker/frames.py | 2 ++ midas/checker/python.py | 32 ++++++++++++++++---- 3 files changed, 83 insertions(+), 5 deletions(-) diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py index c2935f8..7472bf4 100644 --- a/midas/checker/frame_methods.py +++ b/midas/checker/frame_methods.py @@ -1,5 +1,6 @@ from __future__ import annotations +import ast from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable, Optional @@ -39,6 +40,7 @@ def frame_method(*names: str): @dataclass(frozen=True, kw_only=True) class Call: location: Location + call_expr: p.Expr frame: DataFrameType frame_expr: p.Expr positional: list[TypedExpr] @@ -174,6 +176,11 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): positional=call.positional, keywords=call.keywords, ) + if result.is_valid: + self._assert_same_length( + call.call_expr, call.frame_expr, call.positional[0][0] + ) + return result.result @frame_method() @@ -214,3 +221,50 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): keywords=call.keywords, ) return result.result + + def _assert_same_length(self, call_expr: p.Expr, frame1: p.Expr, frame2: p.Expr): + func_name: str = "__midas_frame_same_length__" + self.assertions.define( + func_name, + ast.FunctionDef( + name=func_name, + args=ast.arguments( + posonlyargs=[], + args=[ + ast.arg(arg="frame1"), + ast.arg(arg="frame2"), + ], + kwonlyargs=[], + defaults=[], + kw_defaults=[], + ), + body=[ + ast.Return( + value=ast.Compare( + left=ast.Attribute( + value=ast.Name(id="frame1"), + attr="size", + ), + ops=[ast.Eq()], + comparators=[ + ast.Attribute( + value=ast.Name(id="frame2"), + attr="size", + ) + ], + ) + ) + ], + decorator_list=[], + ), + ) + self.assertions.add( + bound_expr=call_expr, + inputs=[frame1, frame2], + builder=lambda f1, f2: ast.Call( + func=ast.Name(id=func_name), + args=[f1, f2], + keywords=[], + ), + message="DataFrames must have the same length", + ) diff --git a/midas/checker/frames.py b/midas/checker/frames.py index 56c34c6..1a1a917 100644 --- a/midas/checker/frames.py +++ b/midas/checker/frames.py @@ -141,6 +141,7 @@ class FrameManager: self, method: str, location: Location, + call_expr: p.Expr, frame: DataFrameType, frame_expr: p.Expr, positional: list[TypedExpr], @@ -148,6 +149,7 @@ class FrameManager: ) -> Type: call: Call = Call( location=location, + call_expr=call_expr, frame=frame, frame_expr=frame_expr, positional=positional, diff --git a/midas/checker/python.py b/midas/checker/python.py index a42115f..2e0a07c 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -212,6 +212,7 @@ class PythonTyper( def call_method( self, location: Location, + call_expr: p.Expr, obj: TypedExpr, method_name: str, positional: list[TypedExpr], @@ -223,6 +224,7 @@ class PythonTyper( return self.frame_mgr.call( method=method_name, location=location, + call_expr=call_expr, frame=unfolded, frame_expr=obj[0], positional=positional, @@ -503,7 +505,9 @@ class PythonTyper( ) return UnknownType() - return self._visit_binary_expr(expr.location, expr.left, expr.right, method) + return self._visit_binary_expr( + expr.location, expr, expr.left, expr.right, method + ) def visit_compare_expr(self, expr: p.CompareExpr) -> Type: method: Optional[str] = PY_COMPARATOR_METHODS.get(expr.operator.__class__) @@ -514,10 +518,17 @@ class PythonTyper( ) return UnknownType() - return self._visit_binary_expr(expr.location, expr.left, expr.right, method) + return self._visit_binary_expr( + expr.location, expr, expr.left, expr.right, method + ) def _visit_binary_expr( - self, location: Location, left_expr: p.Expr, right_expr: p.Expr, method: str + self, + location: Location, + expr: p.Expr, + left_expr: p.Expr, + right_expr: p.Expr, + method: str, ) -> Type: left: Type = self.type_of(left_expr) right: Type = self.type_of(right_expr) @@ -525,7 +536,12 @@ class PythonTyper( result: Optional[Type] try: result = self.call_method( - location, (left_expr, left), method, [(right_expr, right)], {} + location=location, + call_expr=expr, + obj=(left_expr, left), + method_name=method, + positional=[(right_expr, right)], + keywords={}, ) except UndefinedMethodException: self.reporter.error( @@ -550,7 +566,12 @@ class PythonTyper( result: Optional[Type] try: result = self.call_method( - expr.location, (expr.right, operand), method, [], {} + location=expr.location, + call_expr=expr, + obj=(expr.right, operand), + method_name=method, + positional=[], + keywords={}, ) except UndefinedMethodException: self.reporter.error( @@ -581,6 +602,7 @@ class PythonTyper( return self.frame_mgr.call( method=method, location=expr.location, + call_expr=expr, frame=unfolded, frame_expr=obj, positional=positional, From 955c2233ede5cbd29659bfe636ece0212beda3b7 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 17:14:30 +0200 Subject: [PATCH 08/36] feat(checker): statically evaluate casts to Any and None --- midas/checker/python.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/midas/checker/python.py b/midas/checker/python.py index 2e0a07c..7ab803d 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -28,6 +28,7 @@ from midas.checker.types import ( DerivedType, Function, GenericType, + TopType, TupleType, Type, TypeVar, @@ -967,6 +968,17 @@ class PythonTyper( self, expr: p.CastExpr, subject_type: Type, target_type: Type, lit_value: Any ) -> bool: match target_type: + case TopType(): + return True + + case UnitType(): + if lit_value is not None: + self.reporter.error( + expr.location, f"Value {lit_value!r} is not None" + ) + return False + return True + case DerivedType(type=base): return self._evaluate_cast_statically( expr, subject_type, base, lit_value From 5d20f8ec3e9e586553cc1574458c05b8831382ce Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 17:18:16 +0200 Subject: [PATCH 09/36] docs: mention eager evaluation in manual --- docs/manual.typ | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/docs/manual.typ b/docs/manual.typ index 0479434..c106038 100644 --- a/docs/manual.typ +++ b/docs/manual.typ @@ -678,6 +678,10 @@ In the following example, a runtime check would be generated to ensure that the caption: [Typing of `cast` expression], ) +#gc.warning[ + Assertions are statements inserted just before a statement using a `cast` expression. This means that the expression is evaluated _before_ its actual intended usage location, which might cause issues if you rely on logical operator short-circuiting. See @eager-eval for more information. +] + There may be some cases where the cost of checking a value at runtime is simply not worth the safety, for example when dealing with a big dataset. If do wish so, you can use `unsafe_cast` which will only tell the type checker the type of the value, without generating a runtime assertion. This maps to the default behavior of `typing`'s own `cast` function. If the value passed to `cast` or `unsafe_cast` is a literal (e.g. an integer, a string, a list of literals, etc.), the assertion is evaluated _at compile-time_ and no runtime assertion is generated. @@ -695,3 +699,26 @@ If the value passed to `cast` or `unsafe_cast` is a literal (e.g. an integer, a == Generating Stubs (`stubs`) == Showing Type Judgements (`types`) == Validating Definitions (`validate`) + += Known limitations + +== Eager evaluation in runtime assertions + +The process of generating assertions to ensure safety at runtime, mainly for `cast` expressions, leads to the creation of aliases for the expressions being casted. These alias definitions eagerly evaluate before the assertion, and most importantly before the real usage location. This means that you should avoid using `cast` expressions inside logical expressions like `and` or `or`, because the normal "short-circuit" behavior will be irrelevant to the evaluations of the operands. + +For example: + +#figure( + ```py + def foo(): + print("Foo") + return True + def bar(): + print("Bar") + return True + result = foo() or bar() + # Foo + # Bar + ```, + caption: [Runtime assertions may eagerly evaluate expressions and bypass logical operator's short-circuit], +) From b5acae407807f13cab159de73f27b2c140a7b458 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 17:45:18 +0200 Subject: [PATCH 10/36] feat(types): add FrameGroupBy type --- midas/checker/frame_methods.py | 62 ++++++++++++++++++++++++++++++++++ midas/checker/types.py | 18 +++++++++- midas/generator/generator.py | 7 ++-- midas/generator/stubs.py | 14 ++++++++ 4 files changed, 97 insertions(+), 4 deletions(-) diff --git a/midas/checker/frame_methods.py b/midas/checker/frame_methods.py index 7472bf4..4df00fd 100644 --- a/midas/checker/frame_methods.py +++ b/midas/checker/frame_methods.py @@ -12,6 +12,7 @@ from midas.checker.reporter import FileReporter from midas.checker.types import ( ColumnType, DataFrameType, + FrameGroupBy, Function, OverloadedFunction, TopType, @@ -222,6 +223,67 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): ) return result.result + @frame_method() + def groupby(self, call: Call) -> Type: + bool_: Type = self.types.get_type("bool") + function: Function = Function( + args=[ + Function.Argument( + pos=0, + name="by", + type=TopType(), + required=False, + ), + Function.Argument( + pos=1, + name="level", + type=TopType(), + required=False, + ), + ], + kw_args=[ + Function.Argument( + pos=2, + name="as_index", + type=bool_, + required=False, + ), + Function.Argument( + pos=3, + name="sort", + type=bool_, + required=False, + ), + Function.Argument( + pos=4, + name="group_keys", + type=bool_, + required=False, + ), + Function.Argument( + pos=5, + name="observed", + type=bool_, + required=False, + ), + Function.Argument( + pos=6, + name="dropna", + type=bool_, + required=False, + ), + ], + returns=FrameGroupBy(frame=call.frame), + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=function, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + def _assert_same_length(self, call_expr: p.Expr, frame1: p.Expr, frame2: p.Expr): func_name: str = "__midas_frame_same_length__" self.assertions.define( diff --git a/midas/checker/types.py b/midas/checker/types.py index a4279af..299b392 100644 --- a/midas/checker/types.py +++ b/midas/checker/types.py @@ -187,6 +187,14 @@ class DataFrameType: type: ColumnType +@dataclass(frozen=True, kw_only=True) +class FrameGroupBy: + frame: DataFrameType + + def __str__(self) -> str: + return f"FrameGroupBy[{self.frame}]" + + def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: def sub_argument(arg: Function.Argument): return Function.Argument( @@ -305,11 +313,15 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: columns=list(map(sub_column, columns)), ) + case FrameGroupBy(frame=frame): + return FrameGroupBy( + frame=cast(DataFrameType, substitute_typevars(frame, substitutions)) + ) + case UnknownType() | UnitType(): return type case TopType() | GenericType(): - raise NotImplementedError(f"Unsupported type {type}") # Ensure exhaustiveness @@ -382,6 +394,9 @@ def to_annotation(type: Type) -> str: case DataFrameType(): return "pd.DataFrame" + case FrameGroupBy(): + return "pd.api.typing.DataFrameGroupBy" + case _: assert_never(type) @@ -410,4 +425,5 @@ Type = ( | TupleType | ColumnType | DataFrameType + | FrameGroupBy ) diff --git a/midas/generator/generator.py b/midas/generator/generator.py index bfc1d3c..dc6af9d 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -20,6 +20,7 @@ from midas.checker.types import ( DataFrameType, DerivedType, ExtensionType, + FrameGroupBy, Function, GenericType, OverloadedFunction, @@ -366,7 +367,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): self, src_location: Location, expr: ast.expr, type: Type ) -> list[ast.stmt]: match type: - case UnknownType(): + case UnknownType() | TopType(): return [] case BaseType(name=name): @@ -497,12 +498,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): return asserts case ( - TopType() - | Function() + Function() | OverloadedFunction() | ComplexType() | ExtensionType() | GenericType() + | FrameGroupBy() ): self.logger.warning(f"Can't make assertion for type {type}") return [] diff --git a/midas/generator/stubs.py b/midas/generator/stubs.py index 5f71aac..e86d3f6 100644 --- a/midas/generator/stubs.py +++ b/midas/generator/stubs.py @@ -12,6 +12,7 @@ from midas.checker.types import ( DataFrameType, DerivedType, ExtensionType, + FrameGroupBy, Function, GenericType, OverloadedFunction, @@ -287,6 +288,19 @@ class StubsGenerator: attr="DataFrame", ) + case FrameGroupBy(): + self.import_pandas = True + return ast.Attribute( + value=ast.Attribute( + value=ast.Attribute( + value=ast.Name(id="pd"), + attr="api", + ), + attr="typing", + ), + attr="DataFrameGroupBy", + ) + case _: assert_never(type) From 3c92f0867d118786bc21726e6d75d6a4cf73a656 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 18:00:25 +0200 Subject: [PATCH 11/36] feat(types): add ColumnGroupBy --- midas/checker/frames.py | 30 +++++++++++++++++++++++++++++- midas/checker/python.py | 10 ++++++++++ midas/checker/types.py | 17 +++++++++++++++++ midas/generator/generator.py | 2 ++ midas/generator/stubs.py | 14 ++++++++++++++ 5 files changed, 72 insertions(+), 1 deletion(-) diff --git a/midas/checker/frames.py b/midas/checker/frames.py index 1a1a917..2308317 100644 --- a/midas/checker/frames.py +++ b/midas/checker/frames.py @@ -6,7 +6,15 @@ import midas.ast.python as p from midas.ast.location import Location from midas.checker.frame_methods import Call, MethodRegistry from midas.checker.reporter import FileReporter -from midas.checker.types import ColumnType, DataFrameType, TupleType, Type, UnknownType +from midas.checker.types import ( + ColumnGroupBy, + ColumnType, + DataFrameType, + FrameGroupBy, + TupleType, + Type, + UnknownType, +) if TYPE_CHECKING: from midas.checker.python import PythonTyper, TypedExpr @@ -90,6 +98,26 @@ class FrameManager: reporter.error(location, f"Invalid index type {index} on {frame}") return UnknownType() + def groupby_get( + self, + reporter: FileReporter, + location: Location, + groupby: FrameGroupBy, + index: p.Expr, + ) -> Type: + result: Type = self.get(reporter, location, groupby.frame, index) + match result: + case ColumnType(): + result = ColumnGroupBy(column=result) + case TupleType(items=columns): + result = TupleType( + items=tuple( + ColumnGroupBy(column=cast(ColumnType, column)) + for column in columns + ) + ) + return result + @classmethod def _set_column( cls, frame: DataFrameType, name: str, column: ColumnType diff --git a/midas/checker/python.py b/midas/checker/python.py index 7ab803d..6e2f30c 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -26,6 +26,7 @@ from midas.checker.types import ( ConstraintType, DataFrameType, DerivedType, + FrameGroupBy, Function, GenericType, TopType, @@ -770,6 +771,8 @@ class PythonTyper( return self._visit_tuple_subscript(unfolded, expr) case DataFrameType(): return self._visit_frame_subscript(unfolded, expr) + case FrameGroupBy(): + return self._visit_frame_groupby_subscript(unfolded, expr) operation: Optional[Type] = self.types.lookup_member(object, "__getitem__") if operation is None: @@ -1095,3 +1098,10 @@ class PythonTyper( self, frame: DataFrameType, expr: p.SubscriptExpr ) -> Type: return self.frame_mgr.get(self.reporter, expr.location, frame, expr.index) + + def _visit_frame_groupby_subscript( + self, groupby: FrameGroupBy, expr: p.SubscriptExpr + ) -> Type: + return self.frame_mgr.groupby_get( + self.reporter, expr.location, groupby, expr.index + ) diff --git a/midas/checker/types.py b/midas/checker/types.py index 299b392..51658d7 100644 --- a/midas/checker/types.py +++ b/midas/checker/types.py @@ -195,6 +195,14 @@ class FrameGroupBy: return f"FrameGroupBy[{self.frame}]" +@dataclass(frozen=True, kw_only=True) +class ColumnGroupBy: + column: ColumnType + + def __str__(self) -> str: + return f"ColumnGroupBy[{self.column}]" + + def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: def sub_argument(arg: Function.Argument): return Function.Argument( @@ -318,6 +326,11 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type: frame=cast(DataFrameType, substitute_typevars(frame, substitutions)) ) + case ColumnGroupBy(column=column): + return ColumnGroupBy( + column=cast(ColumnType, substitute_typevars(column, substitutions)) + ) + case UnknownType() | UnitType(): return type @@ -397,6 +410,9 @@ def to_annotation(type: Type) -> str: case FrameGroupBy(): return "pd.api.typing.DataFrameGroupBy" + case ColumnGroupBy(): + return "pd.api.typing.SeriesGroupBy" + case _: assert_never(type) @@ -426,4 +442,5 @@ Type = ( | ColumnType | DataFrameType | FrameGroupBy + | ColumnGroupBy ) diff --git a/midas/generator/generator.py b/midas/generator/generator.py index dc6af9d..352a544 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -14,6 +14,7 @@ from midas.checker.registry import TypesRegistry from midas.checker.types import ( AppliedType, BaseType, + ColumnGroupBy, ColumnType, ComplexType, ConstraintType, @@ -504,6 +505,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): | ExtensionType() | GenericType() | FrameGroupBy() + | ColumnGroupBy() ): self.logger.warning(f"Can't make assertion for type {type}") return [] diff --git a/midas/generator/stubs.py b/midas/generator/stubs.py index e86d3f6..07ab0db 100644 --- a/midas/generator/stubs.py +++ b/midas/generator/stubs.py @@ -6,6 +6,7 @@ from midas.checker.registry import Member, TypesRegistry from midas.checker.types import ( AppliedType, BaseType, + ColumnGroupBy, ColumnType, ComplexType, ConstraintType, @@ -301,6 +302,19 @@ class StubsGenerator: attr="DataFrameGroupBy", ) + case ColumnGroupBy(): + self.import_pandas = True + return ast.Attribute( + value=ast.Attribute( + value=ast.Attribute( + value=ast.Name(id="pd"), + attr="api", + ), + attr="typing", + ), + attr="SeriesGroupBy", + ) + case _: assert_never(type) From 0d5840a4ce9cb8fe880535e92efc75cf61419aae Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 18:19:41 +0200 Subject: [PATCH 12/36] refactor: restructure frame method registry in submodule --- .../{frames.py => frames/frame_manager.py} | 4 +- midas/checker/{ => frames}/frame_methods.py | 78 +++---------------- midas/checker/frames/utils.py | 64 +++++++++++++++ midas/checker/python.py | 2 +- 4 files changed, 77 insertions(+), 71 deletions(-) rename midas/checker/{frames.py => frames/frame_manager.py} (97%) rename midas/checker/{ => frames}/frame_methods.py (81%) create mode 100644 midas/checker/frames/utils.py diff --git a/midas/checker/frames.py b/midas/checker/frames/frame_manager.py similarity index 97% rename from midas/checker/frames.py rename to midas/checker/frames/frame_manager.py index 2308317..36fd84d 100644 --- a/midas/checker/frames.py +++ b/midas/checker/frames/frame_manager.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, MethodRegistry +from midas.checker.frames.frame_methods import Call, FrameMethodRegistry from midas.checker.reporter import FileReporter from midas.checker.types import ( ColumnGroupBy, @@ -27,7 +27,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: MethodRegistry = MethodRegistry(self.typer) + self.method_resolver: FrameMethodRegistry = FrameMethodRegistry(self.typer) def assign( self, diff --git a/midas/checker/frame_methods.py b/midas/checker/frames/frame_methods.py similarity index 81% rename from midas/checker/frame_methods.py rename to midas/checker/frames/frame_methods.py index 4df00fd..afa9e9e 100644 --- a/midas/checker/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -2,13 +2,12 @@ from __future__ import annotations import ast from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Callable, Optional +from typing import TYPE_CHECKING, Callable, Optional import midas.ast.python as p from midas.ast.location import Location -from midas.checker.dispatcher import CallDispatcher, CallResult -from midas.checker.registry import TypesRegistry -from midas.checker.reporter import FileReporter +from midas.checker.dispatcher import CallResult +from midas.checker.frames.utils import MethodRegistry, method from midas.checker.types import ( ColumnType, DataFrameType, @@ -20,22 +19,9 @@ from midas.checker.types import ( UnknownType, unfold_type, ) -from midas.generator.collector import AssertionCollector 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 + from midas.checker.python import TypedExpr @dataclass(frozen=True, kw_only=True) @@ -48,60 +34,16 @@ class Call: keywords: dict[str, TypedExpr] -class _MethodRegistryMeta(type): - _methods: dict[str, Callable[..., Type]] = {} - - def __new__( - cls, - name: str, - bases: tuple[type, ...], - namespace: dict[str, Any], - ): - new_class = super().__new__(cls, name, bases, namespace) - new_class._methods = {} - for attr in namespace.values(): - if callable(attr) and hasattr(attr, "__method_names__"): - for name in attr.__method_names__: # type: ignore - new_class._methods[name] = attr # type: ignore - return new_class - - -class MethodRegistry(metaclass=_MethodRegistryMeta): - def __init__(self, typer: PythonTyper) -> None: - self.typer: PythonTyper = typer - - @property - def reporter(self) -> FileReporter: - return self.typer.reporter - - @property - def types(self) -> TypesRegistry: - return self.typer.types - - @property - def dispatcher(self) -> CallDispatcher[p.Expr]: - return self.typer.dispatcher - - @property - def assertions(self) -> AssertionCollector: - return self.typer.assertions - - def call( - self, - method: str, - call: Call, - ) -> Type: +class FrameMethodRegistry(MethodRegistry): + def call(self, method: str, call: Call) -> Type: func: Optional[Callable[..., Type]] = self._methods.get(method) if func is None: self.reporter.warning(call.location, f"Unknown method {method}") return UnknownType() return func(self, call) - @frame_method("add", "__add__") - def add( - self, - call: Call, - ) -> Type: + @method("add", "__add__") + def add(self, call: Call) -> Type: # TODO: support add with scalar, sequence, Series, dict # TODO: check operation exists on inner column types @@ -184,7 +126,7 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): return result.result - @frame_method() + @method() def mean(self, call: Call) -> Type: with_axis = Function( kw_args=[ @@ -223,7 +165,7 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): ) return result.result - @frame_method() + @method() def groupby(self, call: Call) -> Type: bool_: Type = self.types.get_type("bool") function: Function = Function( diff --git a/midas/checker/frames/utils.py b/midas/checker/frames/utils.py new file mode 100644 index 0000000..ceb6370 --- /dev/null +++ b/midas/checker/frames/utils.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Callable + +import midas.ast.python as p +from midas.checker.dispatcher import CallDispatcher +from midas.checker.registry import TypesRegistry +from midas.checker.reporter import FileReporter +from midas.checker.types import Type +from midas.generator.collector import AssertionCollector + +if TYPE_CHECKING: + from midas.checker.python import PythonTyper + + +@staticmethod +def 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 _MethodRegistryMeta(type): + _methods: dict[str, Callable[..., Type]] = {} + + def __new__( + cls, + name: str, + bases: tuple[type, ...], + namespace: dict[str, Any], + ): + new_class = super().__new__(cls, name, bases, namespace) + new_class._methods = {} + for attr in namespace.values(): + if callable(attr) and hasattr(attr, "__method_names__"): + for name in attr.__method_names__: # type: ignore + new_class._methods[name] = attr # type: ignore + return new_class + + +class MethodRegistry(metaclass=_MethodRegistryMeta): + def __init__(self, typer: PythonTyper) -> None: + self.typer: PythonTyper = typer + + @property + def reporter(self) -> FileReporter: + return self.typer.reporter + + @property + def types(self) -> TypesRegistry: + return self.typer.types + + @property + def dispatcher(self) -> CallDispatcher[p.Expr]: + return self.typer.dispatcher + + @property + def assertions(self) -> AssertionCollector: + return self.typer.assertions diff --git a/midas/checker/python.py b/midas/checker/python.py index 6e2f30c..2031542 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -9,7 +9,7 @@ from midas.ast.printer import MidasPrinter from midas.checker.dispatcher import CallDispatcher, CallResult from midas.checker.environment import Environment from midas.checker.evaluator import Evaluator -from midas.checker.frames import FrameManager +from midas.checker.frames.frame_manager import FrameManager from midas.checker.operators import ( PY_COMPARATOR_METHODS, PY_OPERATOR_METHODS, From b48dfe5301f90574347dbd6434160a978ead5a46 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 18:27:26 +0200 Subject: [PATCH 13/36] refactor: make MethodRegistry generic on Call --- midas/checker/frames/frame_methods.py | 11 ++-------- midas/checker/frames/utils.py | 31 ++++++++++++++++++++++++--- 2 files changed, 30 insertions(+), 12 deletions(-) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index afa9e9e..059a4a8 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -2,7 +2,7 @@ from __future__ import annotations import ast from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, Optional +from typing import TYPE_CHECKING, Optional import midas.ast.python as p from midas.ast.location import Location @@ -34,14 +34,7 @@ class Call: keywords: dict[str, TypedExpr] -class FrameMethodRegistry(MethodRegistry): - def call(self, method: str, call: Call) -> Type: - func: Optional[Callable[..., Type]] = self._methods.get(method) - if func is None: - self.reporter.warning(call.location, f"Unknown method {method}") - return UnknownType() - return func(self, call) - +class FrameMethodRegistry(MethodRegistry[Call]): @method("add", "__add__") def add(self, call: Call) -> Type: # TODO: support add with scalar, sequence, Series, dict diff --git a/midas/checker/frames/utils.py b/midas/checker/frames/utils.py index ceb6370..d6ac0d8 100644 --- a/midas/checker/frames/utils.py +++ b/midas/checker/frames/utils.py @@ -1,12 +1,22 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Generic, + Optional, + Protocol, + Self, + TypeVar, +) import midas.ast.python as p +from midas.ast.location import Location from midas.checker.dispatcher import CallDispatcher from midas.checker.registry import TypesRegistry from midas.checker.reporter import FileReporter -from midas.checker.types import Type +from midas.checker.types import Type, UnknownType from midas.generator.collector import AssertionCollector if TYPE_CHECKING: @@ -43,7 +53,15 @@ class _MethodRegistryMeta(type): return new_class -class MethodRegistry(metaclass=_MethodRegistryMeta): +class HasLocation(Protocol): + @property + def location(self) -> Location: ... + + +T = TypeVar("T", bound=HasLocation) + + +class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta): def __init__(self, typer: PythonTyper) -> None: self.typer: PythonTyper = typer @@ -62,3 +80,10 @@ class MethodRegistry(metaclass=_MethodRegistryMeta): @property def assertions(self) -> AssertionCollector: return self.typer.assertions + + def call(self, method: str, call: T) -> Type: + func: Optional[Callable[[Self, T], Type]] = self._methods.get(method) + if func is None: + self.reporter.warning(call.location, f"Unknown method {method}") + return UnknownType() + return func(self, call) From 640f2d1771ae6630cf927fc669cd771aa0e1158a Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 19:22:28 +0200 Subject: [PATCH 14/36] feat(checker): support unification of frames and columns --- midas/checker/unifier.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/midas/checker/unifier.py b/midas/checker/unifier.py index aae9310..138a503 100644 --- a/midas/checker/unifier.py +++ b/midas/checker/unifier.py @@ -4,6 +4,8 @@ from typing import Optional from midas.checker.registry import TypesRegistry from midas.checker.types import ( AppliedType, + ColumnType, + DataFrameType, Function, GenericType, TopType, @@ -98,6 +100,30 @@ class Unifier: return substitutions + case ( + DataFrameType(columns=template_columns), + DataFrameType(columns=concrete_columns), + ) if len(template_columns) == len(concrete_columns): + substitutions: dict[str, Type] = {} + for template_column, concrete_column in zip( + template_columns, concrete_columns + ): + if template_column.index != concrete_column or ( + template_column.name != concrete_column.name + ): + self.logger.debug( + f"Column mismatch: template={template_column}, concrete={concrete_column}" + ) + raise UnificationError + new_substistutions: dict[str, Type] = self.match( + template_column.type, concrete_column.type + ) + substitutions = self.merge(substitutions, new_substistutions) + return substitutions + + case (ColumnType(type=template_column), ColumnType(type=concrete_column)): + return self.match(template_column, concrete_column) + case (Function(), Function()): mapped: list[tuple[Function.Argument, Function.Argument]] = ( self.map_params(template, concrete) From 2fce2f4bfc370069369813da6b693dcd85977ae2 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 19:23:23 +0200 Subject: [PATCH 15/36] feat(checker): add column method registry --- midas/checker/frames/column_manager.py | 37 ++++++++++++++++++++++++++ midas/checker/frames/column_methods.py | 27 +++++++++++++++++++ midas/checker/frames/utils.py | 27 ++++++++++--------- midas/checker/python.py | 13 +++++++++ 4 files changed, 92 insertions(+), 12 deletions(-) create mode 100644 midas/checker/frames/column_manager.py create mode 100644 midas/checker/frames/column_methods.py diff --git a/midas/checker/frames/column_manager.py b/midas/checker/frames/column_manager.py new file mode 100644 index 0000000..1ff9799 --- /dev/null +++ b/midas/checker/frames/column_manager.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import midas.ast.python as p +from midas.ast.location import Location +from midas.checker.frames.column_methods import Call, ColumnMethodRegistry +from midas.checker.types import ColumnType, Type + +if TYPE_CHECKING: + from midas.checker.python import PythonTyper, TypedExpr + + +class ColumnManager: + def __init__(self, typer: PythonTyper) -> None: + self.typer: PythonTyper = typer + self.method_resolver: ColumnMethodRegistry = ColumnMethodRegistry(self.typer) + + def call( + self, + method: str, + location: Location, + call_expr: p.Expr, + column: ColumnType, + column_expr: p.Expr, + positional: list[TypedExpr], + keywords: dict[str, TypedExpr], + ) -> Type: + call: Call = Call( + location=location, + call_expr=call_expr, + column=column, + column_expr=column_expr, + positional=positional, + keywords=keywords, + ) + return self.method_resolver.call(method, call) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py new file mode 100644 index 0000000..da22a43 --- /dev/null +++ b/midas/checker/frames/column_methods.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import midas.ast.python as p +from midas.ast.location import Location +from midas.checker.frames.utils import MethodRegistry +from midas.checker.types import ( + ColumnType, +) + +if TYPE_CHECKING: + from midas.checker.python import TypedExpr + + +@dataclass(frozen=True, kw_only=True) +class Call: + location: Location + call_expr: p.Expr + column: ColumnType + column_expr: p.Expr + positional: list[TypedExpr] + keywords: dict[str, TypedExpr] + + +class ColumnMethodRegistry(MethodRegistry[Call]): ... diff --git a/midas/checker/frames/utils.py b/midas/checker/frames/utils.py index d6ac0d8..95fa40a 100644 --- a/midas/checker/frames/utils.py +++ b/midas/checker/frames/utils.py @@ -23,18 +23,6 @@ if TYPE_CHECKING: from midas.checker.python import PythonTyper -@staticmethod -def 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 _MethodRegistryMeta(type): _methods: dict[str, Callable[..., Type]] = {} @@ -87,3 +75,18 @@ class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta): self.reporter.warning(call.location, f"Unknown method {method}") return UnknownType() return func(self, call) + + +_Self = TypeVar("_Self", bound=MethodRegistry[Any]) +Method = Callable[[_Self, T], Type] + + +def method(*names: str) -> Callable[[Method[_Self, T]], Method[_Self, T]]: + def wrapper(func: Method[_Self, T]) -> Method[_Self, T]: + names_: tuple[str, ...] = names + if len(names_) == 0: + names_ = (func.__name__,) + setattr(func, "__method_names__", names_) + return func + + return wrapper diff --git a/midas/checker/python.py b/midas/checker/python.py index 2031542..e014681 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -9,6 +9,7 @@ from midas.ast.printer import MidasPrinter from midas.checker.dispatcher import CallDispatcher, CallResult from midas.checker.environment import Environment from midas.checker.evaluator import Evaluator +from midas.checker.frames.column_manager import ColumnManager from midas.checker.frames.frame_manager import FrameManager from midas.checker.operators import ( PY_COMPARATOR_METHODS, @@ -82,6 +83,7 @@ class PythonTyper( self.reporter: FileReporter = reporter.for_file(None) self.types: TypesRegistry = types self.frame_mgr: FrameManager = FrameManager(self) + self.column_mgr: ColumnManager = ColumnManager(self) self.global_env: Environment = Preamble(self.types) self.env: Environment = self.global_env self.locals: dict[p.Expr, int] = {} @@ -233,6 +235,17 @@ class PythonTyper( keywords=keywords, ) + case ColumnType(): + return self.column_mgr.call( + method=method_name, + location=location, + call_expr=call_expr, + column=unfolded, + column_expr=obj[0], + positional=positional, + keywords=keywords, + ) + method: Optional[Type] = self.types.lookup_member(obj[1], method_name) if method is None: raise UndefinedMethodException From 8e8ed6226676edeaf5768fd19718c9c9ccd2912b Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 19:30:43 +0200 Subject: [PATCH 16/36] feat(checker): add add/mean/groupby on columns --- midas/checker/frames/column_methods.py | 195 ++++++++++++++++++++++++- 1 file changed, 192 insertions(+), 3 deletions(-) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index da22a43..91909ff 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -1,13 +1,23 @@ from __future__ import annotations +import ast from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import midas.ast.python as p from midas.ast.location import Location -from midas.checker.frames.utils import MethodRegistry +from midas.checker.dispatcher import CallResult +from midas.checker.frames.utils import MethodRegistry, method from midas.checker.types import ( + ColumnGroupBy, ColumnType, + Function, + GenericType, + TopType, + Type, + TypeVar, + UnknownType, + unfold_type, ) if TYPE_CHECKING: @@ -24,4 +34,183 @@ class Call: keywords: dict[str, TypedExpr] -class ColumnMethodRegistry(MethodRegistry[Call]): ... +class ColumnMethodRegistry(MethodRegistry[Call]): + @method("add", "__add__") + def add(self, call: Call) -> Type: + # TODO: support add with scalar + # TODO: check operation exists on inner column types + + column2: Optional[ColumnType] = None + + col_type1: Type = call.column.type + new_column: Type = ColumnType(type=UnknownType()) + if len(call.positional) != 0: + other: Type = call.positional[0][1] + unfolded_other: Type = unfold_type(other) + if isinstance(unfolded_other, ColumnType): + column2 = unfolded_other + col_type2: Type = column2.type + if self.types.are_equivalent(col_type2, col_type1): + new_column = ColumnType(type=col_type1) + + # Build signature with new column type and generic operand + param_type: TypeVar = TypeVar(name="T", bound=None) + signature = GenericType( + name="add", + params=[param_type], + body=Function( + args=[ + Function.Argument( + pos=0, + name="other", + type=ColumnType(type=param_type), + required=True, + ), + ], + returns=new_column, + ), + ) + + # Map arguments and compute result type + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + if result.is_valid: + self._assert_same_length( + call.call_expr, call.column_expr, call.positional[0][0] + ) + + return result.result + + @method() + def mean(self, call: Call) -> Type: + signature = Function( + kw_args=[ + Function.Argument( + pos=0, + name="axis", + type=TopType(), + required=False, + ) + ], + returns=ColumnType(type=TopType()), + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + + @method() + def groupby(self, call: Call) -> Type: + bool_: Type = self.types.get_type("bool") + function: Function = Function( + args=[ + Function.Argument( + pos=0, + name="by", + type=TopType(), + required=False, + ), + Function.Argument( + pos=1, + name="level", + type=TopType(), + required=False, + ), + ], + kw_args=[ + Function.Argument( + pos=2, + name="as_index", + type=bool_, + required=False, + ), + Function.Argument( + pos=3, + name="sort", + type=bool_, + required=False, + ), + Function.Argument( + pos=4, + name="group_keys", + type=bool_, + required=False, + ), + Function.Argument( + pos=5, + name="observed", + type=bool_, + required=False, + ), + Function.Argument( + pos=6, + name="dropna", + type=bool_, + required=False, + ), + ], + returns=ColumnGroupBy(column=call.column), + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=function, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + + def _assert_same_length(self, call_expr: p.Expr, column1: p.Expr, column2: p.Expr): + func_name: str = "__midas_column_same_length__" + self.assertions.define( + func_name, + ast.FunctionDef( + name=func_name, + args=ast.arguments( + posonlyargs=[], + args=[ + ast.arg(arg="column1"), + ast.arg(arg="column2"), + ], + kwonlyargs=[], + defaults=[], + kw_defaults=[], + ), + body=[ + ast.Return( + value=ast.Compare( + left=ast.Attribute( + value=ast.Name(id="column1"), + attr="size", + ), + ops=[ast.Eq()], + comparators=[ + ast.Attribute( + value=ast.Name(id="column2"), + attr="size", + ) + ], + ) + ) + ], + decorator_list=[], + ), + ) + self.assertions.add( + bound_expr=call_expr, + inputs=[column1, column2], + builder=lambda c1, c2: ast.Call( + func=ast.Name(id=func_name), + args=[c1, c2], + keywords=[], + ), + message="Columns must have the same length", + ) From b14f46d4056707e13d7b1a139c15e8319ca8dc91 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 19:53:58 +0200 Subject: [PATCH 17/36] feat(checker): handle calls on group-bys --- .../checker/frames/column_groupby_methods.py | 66 +++++++++++++++++++ midas/checker/frames/column_manager.py | 27 +++++++- midas/checker/frames/frame_groupby_methods.py | 66 +++++++++++++++++++ midas/checker/frames/frame_manager.py | 25 +++++++ midas/checker/python.py | 35 ++++++++-- 5 files changed, 212 insertions(+), 7 deletions(-) create mode 100644 midas/checker/frames/column_groupby_methods.py create mode 100644 midas/checker/frames/frame_groupby_methods.py diff --git a/midas/checker/frames/column_groupby_methods.py b/midas/checker/frames/column_groupby_methods.py new file mode 100644 index 0000000..de10c84 --- /dev/null +++ b/midas/checker/frames/column_groupby_methods.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import midas.ast.python as p +from midas.ast.location import Location +from midas.checker.dispatcher import CallResult +from midas.checker.frames.utils import MethodRegistry, method +from midas.checker.types import ColumnGroupBy, Function, Type + +if TYPE_CHECKING: + from midas.checker.python import TypedExpr + + +@dataclass(frozen=True, kw_only=True) +class Call: + location: Location + call_expr: p.Expr + groupby: ColumnGroupBy + groupby_expr: p.Expr + positional: list[TypedExpr] + keywords: dict[str, TypedExpr] + + +class ColumnGroupByMethodRegistry(MethodRegistry[Call]): + @method() + def mean(self, call: Call) -> Type: + bool_ = self.types.get_type("bool") + signature = Function( + args=[ + Function.Argument( + pos=0, + name="numeric_only", + type=bool_, + required=False, + ), + Function.Argument( + pos=1, + name="skipna", + type=bool_, + required=False, + ), + Function.Argument( + pos=2, + name="engine", + type=self.types.get_type("str"), + required=False, + ), + Function.Argument( + pos=3, + name="engine_kwargs", + type=self.types.get_type("dict"), + required=False, + ), + ], + returns=call.groupby.column, + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result diff --git a/midas/checker/frames/column_manager.py b/midas/checker/frames/column_manager.py index 1ff9799..0793034 100644 --- a/midas/checker/frames/column_manager.py +++ b/midas/checker/frames/column_manager.py @@ -4,8 +4,10 @@ from typing import TYPE_CHECKING import midas.ast.python as p from midas.ast.location import Location +from midas.checker.frames.column_groupby_methods import Call as GroupByCall +from midas.checker.frames.column_groupby_methods import ColumnGroupByMethodRegistry from midas.checker.frames.column_methods import Call, ColumnMethodRegistry -from midas.checker.types import ColumnType, Type +from midas.checker.types import ColumnGroupBy, ColumnType, Type if TYPE_CHECKING: from midas.checker.python import PythonTyper, TypedExpr @@ -15,6 +17,9 @@ class ColumnManager: def __init__(self, typer: PythonTyper) -> None: self.typer: PythonTyper = typer self.method_resolver: ColumnMethodRegistry = ColumnMethodRegistry(self.typer) + self.groupby_method_resolver: ColumnGroupByMethodRegistry = ( + ColumnGroupByMethodRegistry(self.typer) + ) def call( self, @@ -35,3 +40,23 @@ class ColumnManager: keywords=keywords, ) return self.method_resolver.call(method, call) + + def groupby_call( + self, + method: str, + location: Location, + call_expr: p.Expr, + groupby: ColumnGroupBy, + groupby_expr: p.Expr, + positional: list[TypedExpr], + keywords: dict[str, TypedExpr], + ) -> Type: + call: GroupByCall = GroupByCall( + location=location, + call_expr=call_expr, + groupby=groupby, + groupby_expr=groupby_expr, + positional=positional, + keywords=keywords, + ) + return self.groupby_method_resolver.call(method, call) diff --git a/midas/checker/frames/frame_groupby_methods.py b/midas/checker/frames/frame_groupby_methods.py new file mode 100644 index 0000000..4b0acc8 --- /dev/null +++ b/midas/checker/frames/frame_groupby_methods.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import midas.ast.python as p +from midas.ast.location import Location +from midas.checker.dispatcher import CallResult +from midas.checker.frames.utils import MethodRegistry, method +from midas.checker.types import FrameGroupBy, Function, Type + +if TYPE_CHECKING: + from midas.checker.python import TypedExpr + + +@dataclass(frozen=True, kw_only=True) +class Call: + location: Location + call_expr: p.Expr + groupby: FrameGroupBy + groupby_expr: p.Expr + positional: list[TypedExpr] + keywords: dict[str, TypedExpr] + + +class FrameGroupByMethodRegistry(MethodRegistry[Call]): + @method() + def mean(self, call: Call) -> Type: + bool_ = self.types.get_type("bool") + signature = Function( + args=[ + Function.Argument( + pos=0, + name="numeric_only", + type=bool_, + required=False, + ), + Function.Argument( + pos=1, + name="skipna", + type=bool_, + required=False, + ), + Function.Argument( + pos=2, + name="engine", + type=self.types.get_type("str"), + required=False, + ), + Function.Argument( + pos=3, + name="engine_kwargs", + type=self.types.get_type("dict"), + required=False, + ), + ], + returns=call.groupby.frame, + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result diff --git a/midas/checker/frames/frame_manager.py b/midas/checker/frames/frame_manager.py index 36fd84d..8a5794d 100644 --- a/midas/checker/frames/frame_manager.py +++ b/midas/checker/frames/frame_manager.py @@ -4,6 +4,8 @@ from typing import TYPE_CHECKING, Optional, TypeGuard, cast import midas.ast.python as p from midas.ast.location import Location +from midas.checker.frames.frame_groupby_methods import Call as GroupByCall +from midas.checker.frames.frame_groupby_methods import FrameGroupByMethodRegistry from midas.checker.frames.frame_methods import Call, FrameMethodRegistry from midas.checker.reporter import FileReporter from midas.checker.types import ( @@ -28,6 +30,9 @@ class FrameManager: def __init__(self, typer: PythonTyper) -> None: self.typer: PythonTyper = typer self.method_resolver: FrameMethodRegistry = FrameMethodRegistry(self.typer) + self.groupby_method_resolver: FrameGroupByMethodRegistry = ( + FrameGroupByMethodRegistry(self.typer) + ) def assign( self, @@ -184,3 +189,23 @@ class FrameManager: keywords=keywords, ) return self.method_resolver.call(method, call) + + def groupby_call( + self, + method: str, + location: Location, + call_expr: p.Expr, + groupby: FrameGroupBy, + groupby_expr: p.Expr, + positional: list[TypedExpr], + keywords: dict[str, TypedExpr], + ) -> Type: + call: GroupByCall = GroupByCall( + location=location, + call_expr=call_expr, + groupby=groupby, + groupby_expr=groupby_expr, + positional=positional, + keywords=keywords, + ) + return self.groupby_method_resolver.call(method, call) diff --git a/midas/checker/python.py b/midas/checker/python.py index e014681..8e4b59b 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -23,6 +23,7 @@ from midas.checker.resolver import Resolver from midas.checker.types import ( AppliedType, BaseType, + ColumnGroupBy, ColumnType, ConstraintType, DataFrameType, @@ -235,6 +236,17 @@ class PythonTyper( keywords=keywords, ) + case FrameGroupBy(): + return self.frame_mgr.groupby_call( + method=method_name, + location=location, + call_expr=call_expr, + groupby=unfolded, + groupby_expr=obj[0], + positional=positional, + keywords=keywords, + ) + case ColumnType(): return self.column_mgr.call( method=method_name, @@ -246,6 +258,17 @@ class PythonTyper( keywords=keywords, ) + case ColumnGroupBy(): + return self.column_mgr.groupby_call( + method=method_name, + location=location, + call_expr=call_expr, + groupby=unfolded, + groupby_expr=obj[0], + positional=positional, + keywords=keywords, + ) + method: Optional[Type] = self.types.lookup_member(obj[1], method_name) if method is None: raise UndefinedMethodException @@ -612,17 +635,17 @@ class PythonTyper( match expr.callee: case p.GetExpr(object=obj, name=method): obj_type: Type = self.type_of(obj) - unfolded: Type = unfold_type(obj_type) - if isinstance(unfolded, DataFrameType): - return self.frame_mgr.call( - method=method, + return ( + self.call_method( location=expr.location, call_expr=expr, - frame=unfolded, - frame_expr=obj, + obj=(obj, obj_type), + method_name=method, positional=positional, keywords=keywords, ) + or UnknownType() + ) callee: Type = self.type_of(expr.callee) result: CallResult = self.dispatcher.get_result( From 0288a05901e403e7a3ea961dbb038db3d9bc7710 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 23:29:10 +0200 Subject: [PATCH 18/36] feat(checker): handle assignment to multiple columns --- midas/checker/frames/frame_manager.py | 39 ++++++++++++++++++++++++--- 1 file changed, 35 insertions(+), 4 deletions(-) diff --git a/midas/checker/frames/frame_manager.py b/midas/checker/frames/frame_manager.py index 8a5794d..32bc782 100644 --- a/midas/checker/frames/frame_manager.py +++ b/midas/checker/frames/frame_manager.py @@ -47,12 +47,41 @@ class FrameManager: 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 + isinstance(index.value, str) for index in indices ): - raise NotImplementedError + names: list[str] = [cast(str, index.value) for index in indices] + + if not isinstance(value_type, TupleType): + reporter.error( + location, + f"Cannot assign {type} to dataframe columns. Must be a tuple of columns", + ) + return UnknownType() + + if len(names) != len(value_type.items): + reporter.error( + location, + f"Wrong number of columns. Cannot assign {len(value_type.items)} to {len(names)} targets", + ) + return UnknownType() + + new_frame: Type = frame + for name, value in zip(names, value_type.items): + new_frame = self.assign_column( + reporter, + location, + new_frame, + name, + value, + ) + if not isinstance(new_frame, DataFrameType): + return new_frame + return new_frame case _: - reporter.error(location, f"Invalid index type {index} on {frame}") + reporter.error( + location, f"Invalid index type {index} on {frame} (assignment)" + ) return UnknownType() def assign_column( @@ -100,7 +129,9 @@ class FrameManager: return TupleType(items=tuple(columns)) case _: - reporter.error(location, f"Invalid index type {index} on {frame}") + reporter.error( + location, f"Invalid index type {index} on {frame} (access)" + ) return UnknownType() def groupby_get( From 1bc4c704c36605bc2a82dc2c6b64aab445e437a7 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Thu, 2 Jul 2026 23:41:08 +0200 Subject: [PATCH 19/36] feat(checker): delegate element operation to columns delegate element-wise binary operation on frames to columns --- midas/checker/frames/frame_methods.py | 89 +++++++++++++++++++++++---- 1 file changed, 77 insertions(+), 12 deletions(-) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 059a4a8..a43f1fc 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -21,7 +21,7 @@ from midas.checker.types import ( ) if TYPE_CHECKING: - from midas.checker.python import TypedExpr + from midas.checker.python import TypedExpr, UndefinedMethodException @dataclass(frozen=True, kw_only=True) @@ -35,11 +35,70 @@ class Call: class FrameMethodRegistry(MethodRegistry[Call]): - @method("add", "__add__") - def add(self, call: Call) -> Type: - # TODO: support add with scalar, sequence, Series, dict - # TODO: check operation exists on inner column types + def _get_method_result( + self, + call: Call, + column1: ColumnType, + column2: ColumnType, + method: str, + ) -> ColumnType: + """Get the result of calling a method on a column, passing a second + This function delegates to the main typer the resolution of the method + member, as well as computing the result type. Because we don't have any + AST expression for the individual columns, the frame expressions are + used instead. + + Args: + call (Call): the call that triggered this resolution + column1 (ColumnType): the first column, i.e. left operand + column2 (ColumnType): the second column, i.e. right operand + method (str): the method name + + Returns: + ColumnType: the resulting column. + If the operation is invalid / doesn't exist, + `ColumnType(type=UnknownType())` is returned + """ + unknown: ColumnType = ColumnType(type=UnknownType()) + result: Type = unknown + try: + result = ( + self.typer.call_method( + location=call.location, + call_expr=call.call_expr, + obj=(call.frame_expr, column1), + method_name=method, + positional=[(call.positional[0][0], column2)], + keywords={}, + ) + or unknown + ) + except UndefinedMethodException: + self.reporter.error( + call.location, + f"Undefined operation {method} between {column1} and {column2}", + ) + + if not isinstance(result, ColumnType): + result = unknown + return result + + def _element_binary_op(self, call: Call, method: str) -> DataFrameType: + """Compute the result of an element-wise binary operation + + This function delegates to the matching columns for computing resulting + types. Any column only present in one of the frames is forwarded as a + generic `ColumnType(type=UnknownType())`. Columns only in the second + frame are append at the end of the schema. + + Args: + call (Call): the call that triggered this resolution + method (str): the method name + + Returns: + DataFrameType: the resulting frame type + """ new_columns: list[DataFrameType.Column] = [] by_name: dict[str, DataFrameType.Column] = {} @@ -56,20 +115,20 @@ class FrameMethodRegistry(MethodRegistry[Call]): # Compute new schema: # Step 1: for all columns in frame1: - # - if present in frame2 with equivalent type -> add to schema as is + # - if present in frame2 -> delegate operation to columns # - if not -> add to schema as unknown in_frame1: set[str] = set() for column in call.frame.columns: if column.name is not None: in_frame1.add(column.name) - col_type1: Type = column.type - col_type: Type = ColumnType(type=UnknownType()) + col_type1: ColumnType = column.type + col_type: ColumnType = 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 + col_type2: ColumnType = column2.type + + col_type = self._get_method_result(call, col_type1, col_type2, method) new_column = DataFrameType.Column( index=column.index, @@ -92,6 +151,12 @@ class FrameMethodRegistry(MethodRegistry[Call]): ) ) + return DataFrameType(columns=new_columns) + + @method("add", "__add__") + def add(self, call: Call) -> Type: + # TODO: support add with scalar, sequence, Series, dict + # Build signature with new schema and generic operand signature = Function( args=[ @@ -102,7 +167,7 @@ class FrameMethodRegistry(MethodRegistry[Call]): required=True, ), ], - returns=DataFrameType(columns=new_columns), + returns=self._element_binary_op(call, "__add__"), ) # Map arguments and compute result type From be2fd4c8374ca62a2db85c8bbf01a2ba8361d40c Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 00:05:40 +0200 Subject: [PATCH 20/36] feat(checker): delegate element operation to inner type delegate element-wise binary operation on columns to their inner types --- midas/checker/frames/column_methods.py | 34 +++++++++++++++++++------ midas/checker/frames/frame_methods.py | 35 +++++++++----------------- midas/checker/python.py | 35 ++++++++++++++++---------- 3 files changed, 61 insertions(+), 43 deletions(-) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 91909ff..e8a2ce6 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -35,11 +35,19 @@ class Call: class ColumnMethodRegistry(MethodRegistry[Call]): - @method("add", "__add__") - def add(self, call: Call) -> Type: - # TODO: support add with scalar - # TODO: check operation exists on inner column types + def _element_binary_op(self, call: Call, method: str) -> ColumnType: + """Compute the result of an element-wise binary operation + This function delegates to the inner types for computing the resulting + type. + + Args: + call (Call): the call that triggered this resolution + method (str): the method name + + Returns: + ColumnType: the resulting column type + """ column2: Optional[ColumnType] = None col_type1: Type = call.column.type @@ -50,8 +58,20 @@ class ColumnMethodRegistry(MethodRegistry[Call]): if isinstance(unfolded_other, ColumnType): column2 = unfolded_other col_type2: Type = column2.type - if self.types.are_equivalent(col_type2, col_type1): - new_column = ColumnType(type=col_type1) + + new_inner_type = self.typer.result_of_binary_op( + location=call.location, + expr=call.call_expr, + left=(call.column_expr, col_type1), + right=(call.positional[0][0], col_type2), + method=method, + ) + new_column = ColumnType(type=new_inner_type) + return new_column + + @method("add", "__add__") + def add(self, call: Call) -> Type: + # TODO: support add with scalar # Build signature with new column type and generic operand param_type: TypeVar = TypeVar(name="T", bound=None) @@ -67,7 +87,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): required=True, ), ], - returns=new_column, + returns=self._element_binary_op(call, "__add__"), ), ) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index a43f1fc..64c8d06 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -21,7 +21,7 @@ from midas.checker.types import ( ) if TYPE_CHECKING: - from midas.checker.python import TypedExpr, UndefinedMethodException + from midas.checker.python import TypedExpr @dataclass(frozen=True, kw_only=True) @@ -60,28 +60,17 @@ class FrameMethodRegistry(MethodRegistry[Call]): If the operation is invalid / doesn't exist, `ColumnType(type=UnknownType())` is returned """ - unknown: ColumnType = ColumnType(type=UnknownType()) - result: Type = unknown - try: - result = ( - self.typer.call_method( - location=call.location, - call_expr=call.call_expr, - obj=(call.frame_expr, column1), - method_name=method, - positional=[(call.positional[0][0], column2)], - keywords={}, - ) - or unknown - ) - except UndefinedMethodException: - self.reporter.error( - call.location, - f"Undefined operation {method} between {column1} and {column2}", - ) + + result: Type = self.typer.result_of_binary_op( + location=call.location, + expr=call.call_expr, + left=(call.frame_expr, column1), + right=(call.positional[0][0], column2), + method=method, + ) if not isinstance(result, ColumnType): - result = unknown + return ColumnType(type=UnknownType()) return result def _element_binary_op(self, call: Call, method: str) -> DataFrameType: @@ -105,8 +94,8 @@ class FrameMethodRegistry(MethodRegistry[Call]): frame2: Optional[DataFrameType] = None # Get map of operand's columns by name, if there is at least 1 operand, which is a dataframe if len(call.positional) != 0: - other: Type = call.positional[0][1] - unfolded_other: Type = unfold_type(other) + operand: TypedExpr = call.positional[0] + unfolded_other: Type = unfold_type(operand[1]) if isinstance(unfolded_other, DataFrameType): frame2 = unfolded_other by_name = { diff --git a/midas/checker/python.py b/midas/checker/python.py index 8e4b59b..d095162 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -543,8 +543,14 @@ class PythonTyper( ) return UnknownType() - return self._visit_binary_expr( - expr.location, expr, expr.left, expr.right, method + left: Type = self.type_of(expr.left) + right: Type = self.type_of(expr.right) + return self.result_of_binary_op( + expr.location, + expr, + (expr.left, left), + (expr.right, right), + method, ) def visit_compare_expr(self, expr: p.CompareExpr) -> Type: @@ -556,35 +562,38 @@ class PythonTyper( ) return UnknownType() - return self._visit_binary_expr( - expr.location, expr, expr.left, expr.right, method + left: Type = self.type_of(expr.left) + right: Type = self.type_of(expr.right) + return self.result_of_binary_op( + expr.location, + expr, + (expr.left, left), + (expr.right, right), + method, ) - def _visit_binary_expr( + def result_of_binary_op( self, location: Location, expr: p.Expr, - left_expr: p.Expr, - right_expr: p.Expr, + left: TypedExpr, + right: TypedExpr, method: str, ) -> Type: - left: Type = self.type_of(left_expr) - right: Type = self.type_of(right_expr) - result: Optional[Type] try: result = self.call_method( location=location, call_expr=expr, - obj=(left_expr, left), + obj=left, method_name=method, - positional=[(right_expr, right)], + positional=[right], keywords={}, ) except UndefinedMethodException: self.reporter.error( location, - f"Undefined operation {method} between {left} and {right}", + f"Undefined operation {method} between {left[1]} and {right[1]}", ) return UnknownType() From 74c07c9afb63caa5ac2807c01f5623686602f136 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 00:38:56 +0200 Subject: [PATCH 21/36] feat(checker): add arithmetic binary ops on frames --- midas/checker/frames/frame_methods.py | 36 +++++++++++++++++++++++---- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 64c8d06..02edb3c 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -142,10 +142,8 @@ class FrameMethodRegistry(MethodRegistry[Call]): return DataFrameType(columns=new_columns) - @method("add", "__add__") - def add(self, call: Call) -> Type: - # TODO: support add with scalar, sequence, Series, dict - + def _element_wise(self, call: Call, method: str) -> Type: + # TODO: support scalar, sequence, Series, dict operand # Build signature with new schema and generic operand signature = Function( args=[ @@ -156,7 +154,7 @@ class FrameMethodRegistry(MethodRegistry[Call]): required=True, ), ], - returns=self._element_binary_op(call, "__add__"), + returns=self._element_binary_op(call, method), ) # Map arguments and compute result type @@ -173,6 +171,34 @@ class FrameMethodRegistry(MethodRegistry[Call]): return result.result + @method("add", "__add__") + def add(self, call: Call) -> Type: + return self._element_wise(call, "__add__") + + @method("sub", "__sub__") + def sub(self, call: Call) -> Type: + return self._element_wise(call, "__sub__") + + @method("mul", "__mul__") + def mul(self, call: Call) -> Type: + return self._element_wise(call, "__mul__") + + @method("div", "truediv", "__truediv__") + def truediv(self, call: Call) -> Type: + return self._element_wise(call, "__truediv__") + + @method("floordiv", "__floordiv__") + def floordiv(self, call: Call) -> Type: + return self._element_wise(call, "__floordiv__") + + @method("mod", "__mod__") + def mod(self, call: Call) -> Type: + return self._element_wise(call, "__mod__") + + @method("pow", "__pow__") + def pow(self, call: Call) -> Type: + return self._element_wise(call, "__pow__") + @method() def mean(self, call: Call) -> Type: with_axis = Function( From ece2e3a6a309dcd64d117e6da1d43beb3563e6b2 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 00:42:00 +0200 Subject: [PATCH 22/36] feat(checker): add arithmetic binary ops on columns --- midas/checker/frames/column_methods.py | 33 +++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index e8a2ce6..505caff 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -69,8 +69,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): new_column = ColumnType(type=new_inner_type) return new_column - @method("add", "__add__") - def add(self, call: Call) -> Type: + def _element_wise(self, call: Call, method: str) -> Type: # TODO: support add with scalar # Build signature with new column type and generic operand @@ -87,7 +86,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): required=True, ), ], - returns=self._element_binary_op(call, "__add__"), + returns=self._element_binary_op(call, method), ), ) @@ -105,6 +104,34 @@ class ColumnMethodRegistry(MethodRegistry[Call]): return result.result + @method("add", "__add__") + def add(self, call: Call) -> Type: + return self._element_wise(call, "__add__") + + @method("sub", "__sub__") + def sub(self, call: Call) -> Type: + return self._element_wise(call, "__sub__") + + @method("mul", "__mul__") + def mul(self, call: Call) -> Type: + return self._element_wise(call, "__mul__") + + @method("div", "truediv", "__truediv__") + def truediv(self, call: Call) -> Type: + return self._element_wise(call, "__truediv__") + + @method("floordiv", "__floordiv__") + def floordiv(self, call: Call) -> Type: + return self._element_wise(call, "__floordiv__") + + @method("mod", "__mod__") + def mod(self, call: Call) -> Type: + return self._element_wise(call, "__mod__") + + @method("pow", "__pow__") + def pow(self, call: Call) -> Type: + return self._element_wise(call, "__pow__") + @method() def mean(self, call: Call) -> Type: signature = Function( From 6e06f9078e2da1213db72f47dc335b41085b2b63 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 00:57:10 +0200 Subject: [PATCH 23/36] fix(checker): improve unknown method message --- midas/checker/frames/column_groupby_methods.py | 4 ++++ midas/checker/frames/column_methods.py | 4 ++++ midas/checker/frames/frame_groupby_methods.py | 4 ++++ midas/checker/frames/frame_methods.py | 4 ++++ midas/checker/frames/utils.py | 16 ++++++++++++---- 5 files changed, 28 insertions(+), 4 deletions(-) diff --git a/midas/checker/frames/column_groupby_methods.py b/midas/checker/frames/column_groupby_methods.py index de10c84..aaed5de 100644 --- a/midas/checker/frames/column_groupby_methods.py +++ b/midas/checker/frames/column_groupby_methods.py @@ -22,6 +22,10 @@ class Call: positional: list[TypedExpr] keywords: dict[str, TypedExpr] + @property + def subject(self) -> TypedExpr: + return (self.groupby_expr, self.groupby) + class ColumnGroupByMethodRegistry(MethodRegistry[Call]): @method() diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 505caff..4b101f2 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -33,6 +33,10 @@ class Call: positional: list[TypedExpr] keywords: dict[str, TypedExpr] + @property + def subject(self) -> TypedExpr: + return (self.column_expr, self.column) + class ColumnMethodRegistry(MethodRegistry[Call]): def _element_binary_op(self, call: Call, method: str) -> ColumnType: diff --git a/midas/checker/frames/frame_groupby_methods.py b/midas/checker/frames/frame_groupby_methods.py index 4b0acc8..d3f882d 100644 --- a/midas/checker/frames/frame_groupby_methods.py +++ b/midas/checker/frames/frame_groupby_methods.py @@ -22,6 +22,10 @@ class Call: positional: list[TypedExpr] keywords: dict[str, TypedExpr] + @property + def subject(self) -> TypedExpr: + return (self.groupby_expr, self.groupby) + class FrameGroupByMethodRegistry(MethodRegistry[Call]): @method() diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 02edb3c..4c413b4 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -33,6 +33,10 @@ class Call: positional: list[TypedExpr] keywords: dict[str, TypedExpr] + @property + def subject(self) -> TypedExpr: + return (self.frame_expr, self.frame) + class FrameMethodRegistry(MethodRegistry[Call]): def _get_method_result( diff --git a/midas/checker/frames/utils.py b/midas/checker/frames/utils.py index 95fa40a..dee4696 100644 --- a/midas/checker/frames/utils.py +++ b/midas/checker/frames/utils.py @@ -20,7 +20,7 @@ from midas.checker.types import Type, UnknownType from midas.generator.collector import AssertionCollector if TYPE_CHECKING: - from midas.checker.python import PythonTyper + from midas.checker.python import PythonTyper, TypedExpr class _MethodRegistryMeta(type): @@ -41,12 +41,18 @@ class _MethodRegistryMeta(type): return new_class -class HasLocation(Protocol): +class MethodCall(Protocol): @property def location(self) -> Location: ... + @property + def call_expr(self) -> p.Expr: ... -T = TypeVar("T", bound=HasLocation) + @property + def subject(self) -> TypedExpr: ... + + +T = TypeVar("T", bound=MethodCall) class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta): @@ -72,7 +78,9 @@ class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta): def call(self, method: str, call: T) -> Type: func: Optional[Callable[[Self, T], Type]] = self._methods.get(method) if func is None: - self.reporter.warning(call.location, f"Unknown method {method}") + self.reporter.warning( + call.location, f"Unknown method {method} on {call.subject[1]}" + ) return UnknownType() return func(self, call) From d8c0b17512fbf3de0a4d5dbf870ac1dde7d3d095 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 00:57:27 +0200 Subject: [PATCH 24/36] feat(checker): add comparison binary ops on frames --- midas/checker/frames/frame_methods.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 4c413b4..c63e5b8 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -203,6 +203,30 @@ class FrameMethodRegistry(MethodRegistry[Call]): def pow(self, call: Call) -> Type: return self._element_wise(call, "__pow__") + @method("lt", "__lt__") + def lt(self, call: Call) -> Type: + return self._element_wise(call, "__lt__") + + @method("gt", "__gt__") + def gt(self, call: Call) -> Type: + return self._element_wise(call, "__gt__") + + @method("le", "__le__") + def le(self, call: Call) -> Type: + return self._element_wise(call, "__le__") + + @method("ge", "__ge__") + def ge(self, call: Call) -> Type: + return self._element_wise(call, "__ge__") + + @method("ne", "__ne__") + def ne(self, call: Call) -> Type: + return self._element_wise(call, "__ne__") + + @method("eq", "__eq__") + def eq(self, call: Call) -> Type: + return self._element_wise(call, "__eq__") + @method() def mean(self, call: Call) -> Type: with_axis = Function( From 819ab3c2bfb1d95824c0003c02048f6971a361db Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 00:57:59 +0200 Subject: [PATCH 25/36] tests: add dataframe operations test --- tests/cases/checker/09_frame_ops.py | 21 + tests/cases/checker/09_frame_ops.py.ref.json | 1500 ++++++++++++++++++ 2 files changed, 1521 insertions(+) create mode 100644 tests/cases/checker/09_frame_ops.py create mode 100644 tests/cases/checker/09_frame_ops.py.ref.json diff --git a/tests/cases/checker/09_frame_ops.py b/tests/cases/checker/09_frame_ops.py new file mode 100644 index 0000000..75c6264 --- /dev/null +++ b/tests/cases/checker/09_frame_ops.py @@ -0,0 +1,21 @@ +# type: ignore +# ruff: disable [F821] + +df1: Frame[a:int, b:float] +df2: Frame[a:int, b:float] + +_: Any +_ = df1 + df2 +_ = df1 - df2 +_ = df1 * df2 +_ = df1 / df2 +_ = df1 // df2 +_ = df1 % df2 +_ = df1**df2 + +_ = df1 < df2 +_ = df1 > df2 +_ = df1 <= df2 +_ = df1 >= df2 +_ = df1 != df2 +_ = df1 == df2 diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json new file mode 100644 index 0000000..d30278e --- /dev/null +++ b/tests/cases/checker/09_frame_ops.py.ref.json @@ -0,0 +1,1500 @@ +{ + "diagnostics": [ + { + "type": "Warning", + "location": { + "start": [ + 16, + 4 + ], + "end": [ + 16, + 13 + ] + }, + "message": "Unknown method __lt__ on Column[int]" + }, + { + "type": "Warning", + "location": { + "start": [ + 16, + 4 + ], + "end": [ + 16, + 13 + ] + }, + "message": "Unknown method __lt__ on Column[float]" + }, + { + "type": "Warning", + "location": { + "start": [ + 17, + 4 + ], + "end": [ + 17, + 13 + ] + }, + "message": "Unknown method __gt__ on Column[int]" + }, + { + "type": "Warning", + "location": { + "start": [ + 17, + 4 + ], + "end": [ + 17, + 13 + ] + }, + "message": "Unknown method __gt__ on Column[float]" + }, + { + "type": "Warning", + "location": { + "start": [ + 18, + 4 + ], + "end": [ + 18, + 14 + ] + }, + "message": "Unknown method __le__ on Column[int]" + }, + { + "type": "Warning", + "location": { + "start": [ + 18, + 4 + ], + "end": [ + 18, + 14 + ] + }, + "message": "Unknown method __le__ on Column[float]" + }, + { + "type": "Warning", + "location": { + "start": [ + 19, + 4 + ], + "end": [ + 19, + 14 + ] + }, + "message": "Unknown method __ge__ on Column[int]" + }, + { + "type": "Warning", + "location": { + "start": [ + 19, + 4 + ], + "end": [ + 19, + 14 + ] + }, + "message": "Unknown method __ge__ on Column[float]" + }, + { + "type": "Warning", + "location": { + "start": [ + 20, + 4 + ], + "end": [ + 20, + 14 + ] + }, + "message": "Unknown method __eq__ on Column[int]" + }, + { + "type": "Warning", + "location": { + "start": [ + 20, + 4 + ], + "end": [ + 20, + 14 + ] + }, + "message": "Unknown method __eq__ on Column[float]" + }, + { + "type": "Warning", + "location": { + "start": [ + 21, + 4 + ], + "end": [ + 21, + 14 + ] + }, + "message": "Unknown method __eq__ on Column[int]" + }, + { + "type": "Warning", + "location": { + "start": [ + 21, + 4 + ], + "end": [ + 21, + 14 + ] + }, + "message": "Unknown method __eq__ on Column[float]" + } + ], + "judgments": [ + { + "location": { + "from": "L8:4", + "to": "L8:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L8:10", + "to": "L8:13" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L8:4", + "to": "L8:13" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "+", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L9:4", + "to": "L9:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L9:10", + "to": "L9:13" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L9:4", + "to": "L9:13" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "-", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L10:4", + "to": "L10:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L10:10", + "to": "L10:13" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L10:4", + "to": "L10:13" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "*", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L11:4", + "to": "L11:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L11:10", + "to": "L11:13" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L11:4", + "to": "L11:13" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "/", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "float" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L12:4", + "to": "L12:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L12:11", + "to": "L12:14" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L12:4", + "to": "L12:14" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "//", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L13:4", + "to": "L13:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L13:10", + "to": "L13:13" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L13:4", + "to": "L13:13" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "%", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L14:4", + "to": "L14:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L14:9", + "to": "L14:12" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L14:4", + "to": "L14:12" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "**", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": {} + } + } + ] + } + }, + { + "location": { + "from": "L16:4", + "to": "L16:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L16:10", + "to": "L16:13" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L16:4", + "to": "L16:13" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "<", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": {} + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": {} + } + } + ] + } + }, + { + "location": { + "from": "L17:4", + "to": "L17:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L17:10", + "to": "L17:13" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L17:4", + "to": "L17:13" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": ">", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": {} + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": {} + } + } + ] + } + }, + { + "location": { + "from": "L18:4", + "to": "L18:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L18:11", + "to": "L18:14" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L18:4", + "to": "L18:14" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "<=", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": {} + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": {} + } + } + ] + } + }, + { + "location": { + "from": "L19:4", + "to": "L19:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L19:11", + "to": "L19:14" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L19:4", + "to": "L19:14" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": ">=", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": {} + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": {} + } + } + ] + } + }, + { + "location": { + "from": "L20:4", + "to": "L20:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L20:11", + "to": "L20:14" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L20:4", + "to": "L20:14" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "!=", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": {} + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": {} + } + } + ] + } + }, + { + "location": { + "from": "L21:4", + "to": "L21:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L21:11", + "to": "L21:14" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L21:4", + "to": "L21:14" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "==", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": {} + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": {} + } + } + ] + } + } + ] +} \ No newline at end of file From 66b585c3d6961a5c88366d1464def1a315fcc2f9 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 01:04:45 +0200 Subject: [PATCH 26/36] fix(checker): recursively check builtin subtypes --- midas/checker/builtins.py | 2 ++ midas/checker/registry.py | 11 ++++++++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/midas/checker/builtins.py b/midas/checker/builtins.py index ce97d9c..017f158 100644 --- a/midas/checker/builtins.py +++ b/midas/checker/builtins.py @@ -14,6 +14,8 @@ if TYPE_CHECKING: from midas.checker.registry import TypesRegistry +# Hard-coded subtype relationships between builtin types +# Circular dependencies and diamond inheritance MUST be avoided BUILTIN_SUBTYPES: dict[str, set[str]] = { "object": {"float", "list", "dict", "str", "bytes", "tuple"}, "float": {"int"}, diff --git a/midas/checker/registry.py b/midas/checker/registry.py index 02c96c3..d1074a7 100644 --- a/midas/checker/registry.py +++ b/midas/checker/registry.py @@ -113,6 +113,15 @@ class TypesRegistry: raise ValueError(f"Predicate {name} already defined") self._predicates[name] = predicate + def is_builtin_subtype(self, name1: str, name2: str) -> bool: + subtypes: set[str] = BUILTIN_SUBTYPES.get(name2, set()) + if name1 in subtypes: + return True + for subtype in subtypes: + if self.is_builtin_subtype(name1, subtype): + return True + return False + def is_subtype(self, type1: Type, type2: Type) -> bool: """Check whether `type1` is a subtype of `type2` @@ -150,7 +159,7 @@ class TypesRegistry: return self.is_subtype(base1, type2) case (BaseType(name=name1), BaseType(name=name2)): - return name1 in BUILTIN_SUBTYPES.get(name2, set()) + return self.is_builtin_subtype(name1, name2) case (ComplexType(properties=props1), ComplexType(properties=props2)): for k, t in props2.items(): From 1f6189daa46422a3e1c9e8910209e17f10c88bc6 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 01:05:24 +0200 Subject: [PATCH 27/36] feat(checker): add comparison binary ops on columns --- midas/checker/frames/column_methods.py | 24 ++ tests/cases/checker/09_frame_ops.py.ref.json | 219 ++++--------------- 2 files changed, 61 insertions(+), 182 deletions(-) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 4b101f2..c9da035 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -136,6 +136,30 @@ class ColumnMethodRegistry(MethodRegistry[Call]): def pow(self, call: Call) -> Type: return self._element_wise(call, "__pow__") + @method("lt", "__lt__") + def lt(self, call: Call) -> Type: + return self._element_wise(call, "__lt__") + + @method("gt", "__gt__") + def gt(self, call: Call) -> Type: + return self._element_wise(call, "__gt__") + + @method("le", "__le__") + def le(self, call: Call) -> Type: + return self._element_wise(call, "__le__") + + @method("ge", "__ge__") + def ge(self, call: Call) -> Type: + return self._element_wise(call, "__ge__") + + @method("ne", "__ne__") + def ne(self, call: Call) -> Type: + return self._element_wise(call, "__ne__") + + @method("eq", "__eq__") + def eq(self, call: Call) -> Type: + return self._element_wise(call, "__eq__") + @method() def mean(self, call: Call) -> Type: signature = Function( diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json index d30278e..78ca071 100644 --- a/tests/cases/checker/09_frame_ops.py.ref.json +++ b/tests/cases/checker/09_frame_ops.py.ref.json @@ -1,174 +1,5 @@ { - "diagnostics": [ - { - "type": "Warning", - "location": { - "start": [ - 16, - 4 - ], - "end": [ - 16, - 13 - ] - }, - "message": "Unknown method __lt__ on Column[int]" - }, - { - "type": "Warning", - "location": { - "start": [ - 16, - 4 - ], - "end": [ - 16, - 13 - ] - }, - "message": "Unknown method __lt__ on Column[float]" - }, - { - "type": "Warning", - "location": { - "start": [ - 17, - 4 - ], - "end": [ - 17, - 13 - ] - }, - "message": "Unknown method __gt__ on Column[int]" - }, - { - "type": "Warning", - "location": { - "start": [ - 17, - 4 - ], - "end": [ - 17, - 13 - ] - }, - "message": "Unknown method __gt__ on Column[float]" - }, - { - "type": "Warning", - "location": { - "start": [ - 18, - 4 - ], - "end": [ - 18, - 14 - ] - }, - "message": "Unknown method __le__ on Column[int]" - }, - { - "type": "Warning", - "location": { - "start": [ - 18, - 4 - ], - "end": [ - 18, - 14 - ] - }, - "message": "Unknown method __le__ on Column[float]" - }, - { - "type": "Warning", - "location": { - "start": [ - 19, - 4 - ], - "end": [ - 19, - 14 - ] - }, - "message": "Unknown method __ge__ on Column[int]" - }, - { - "type": "Warning", - "location": { - "start": [ - 19, - 4 - ], - "end": [ - 19, - 14 - ] - }, - "message": "Unknown method __ge__ on Column[float]" - }, - { - "type": "Warning", - "location": { - "start": [ - 20, - 4 - ], - "end": [ - 20, - 14 - ] - }, - "message": "Unknown method __eq__ on Column[int]" - }, - { - "type": "Warning", - "location": { - "start": [ - 20, - 4 - ], - "end": [ - 20, - 14 - ] - }, - "message": "Unknown method __eq__ on Column[float]" - }, - { - "type": "Warning", - "location": { - "start": [ - 21, - 4 - ], - "end": [ - 21, - 14 - ] - }, - "message": "Unknown method __eq__ on Column[int]" - }, - { - "type": "Warning", - "location": { - "start": [ - 21, - 4 - ], - "end": [ - 21, - 14 - ] - }, - "message": "Unknown method __eq__ on Column[float]" - } - ], + "diagnostics": [], "judgments": [ { "location": { @@ -983,14 +814,18 @@ "index": 0, "name": "a", "type": { - "type": {} + "type": { + "name": "bool" + } } }, { "index": 1, "name": "b", "type": { - "type": {} + "type": { + "name": "bool" + } } } ] @@ -1083,14 +918,18 @@ "index": 0, "name": "a", "type": { - "type": {} + "type": { + "name": "bool" + } } }, { "index": 1, "name": "b", "type": { - "type": {} + "type": { + "name": "bool" + } } } ] @@ -1183,14 +1022,18 @@ "index": 0, "name": "a", "type": { - "type": {} + "type": { + "name": "bool" + } } }, { "index": 1, "name": "b", "type": { - "type": {} + "type": { + "name": "bool" + } } } ] @@ -1283,14 +1126,18 @@ "index": 0, "name": "a", "type": { - "type": {} + "type": { + "name": "bool" + } } }, { "index": 1, "name": "b", "type": { - "type": {} + "type": { + "name": "bool" + } } } ] @@ -1383,14 +1230,18 @@ "index": 0, "name": "a", "type": { - "type": {} + "type": { + "name": "bool" + } } }, { "index": 1, "name": "b", "type": { - "type": {} + "type": { + "name": "bool" + } } } ] @@ -1483,14 +1334,18 @@ "index": 0, "name": "a", "type": { - "type": {} + "type": { + "name": "bool" + } } }, { "index": 1, "name": "b", "type": { - "type": {} + "type": { + "name": "bool" + } } } ] From 1c0c917873710706c462e1074e5939f42e1846b7 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 01:27:16 +0200 Subject: [PATCH 28/36] feat(checker): add statistical ops on frames --- midas/checker/frames/frame_methods.py | 69 +- tests/cases/checker/09_frame_ops.py | 17 + tests/cases/checker/09_frame_ops.py.ref.json | 1540 +++++++++++++----- 3 files changed, 1176 insertions(+), 450 deletions(-) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index c63e5b8..5c7b4b4 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -227,8 +227,7 @@ class FrameMethodRegistry(MethodRegistry[Call]): def eq(self, call: Call) -> Type: return self._element_wise(call, "__eq__") - @method() - def mean(self, call: Call) -> Type: + def _statistical(self, call: Call, kwargs: list[Function.Argument] = []) -> Type: with_axis = Function( kw_args=[ Function.Argument( @@ -236,7 +235,8 @@ class FrameMethodRegistry(MethodRegistry[Call]): name="axis", type=self.types.get_type("int"), required=False, - ) + ), + *kwargs, ], returns=ColumnType(type=TopType()), ) @@ -247,7 +247,8 @@ class FrameMethodRegistry(MethodRegistry[Call]): name="axis", type=self.types.get_type("None"), required=True, - ) + ), + *kwargs, ], returns=TopType(), ) @@ -266,6 +267,66 @@ class FrameMethodRegistry(MethodRegistry[Call]): ) return result.result + @method("kurtosis", "kurt") + def kurtosis(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def max(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def mean(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def median(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def min(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def mode(self, call: Call) -> Type: + return self._statistical(call) + + @method("product", "prod") + def product(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def std(self, call: Call) -> Type: + return self._statistical( + call, + [ + Function.Argument( + pos=1, + name="ddof", + type=self.types.get_type("int"), + required=False, + ) + ], + ) + + @method() + def sum(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def var(self, call: Call) -> Type: + return self._statistical( + call, + [ + Function.Argument( + pos=1, + name="var", + type=self.types.get_type("int"), + required=False, + ) + ], + ) + @method() def groupby(self, call: Call) -> Type: bool_: Type = self.types.get_type("bool") diff --git a/tests/cases/checker/09_frame_ops.py b/tests/cases/checker/09_frame_ops.py index 75c6264..3bd25cf 100644 --- a/tests/cases/checker/09_frame_ops.py +++ b/tests/cases/checker/09_frame_ops.py @@ -5,6 +5,8 @@ df1: Frame[a:int, b:float] df2: Frame[a:int, b:float] _: Any + +# Arithmetic _ = df1 + df2 _ = df1 - df2 _ = df1 * df2 @@ -13,9 +15,24 @@ _ = df1 // df2 _ = df1 % df2 _ = df1**df2 +# Comparisons _ = df1 < df2 _ = df1 > df2 _ = df1 <= df2 _ = df1 >= df2 _ = df1 != df2 _ = df1 == df2 + +# Statistical +_ = df1.kurt() +_ = df1.kurtosis() +_ = df1.max() +_ = df1.mean() +_ = df1.median() +_ = df1.min() +_ = df1.mode() +_ = df1.prod() +_ = df1.product() +_ = df1.std() +_ = df1.sum() +_ = df1.var() diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json index 78ca071..551b8dc 100644 --- a/tests/cases/checker/09_frame_ops.py.ref.json +++ b/tests/cases/checker/09_frame_ops.py.ref.json @@ -1,214 +1,6 @@ { "diagnostics": [], "judgments": [ - { - "location": { - "from": "L8:4", - "to": "L8:7" - }, - "expr": { - "_type": "VariableExpr", - "name": "df1" - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "int" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "float" - } - } - } - ] - } - }, - { - "location": { - "from": "L8:10", - "to": "L8:13" - }, - "expr": { - "_type": "VariableExpr", - "name": "df2" - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "int" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "float" - } - } - } - ] - } - }, - { - "location": { - "from": "L8:4", - "to": "L8:13" - }, - "expr": { - "_type": "BinaryExpr", - "left": { - "_type": "VariableExpr", - "name": "df1" - }, - "operator": "+", - "right": { - "_type": "VariableExpr", - "name": "df2" - } - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "int" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "float" - } - } - } - ] - } - }, - { - "location": { - "from": "L9:4", - "to": "L9:7" - }, - "expr": { - "_type": "VariableExpr", - "name": "df1" - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "int" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "float" - } - } - } - ] - } - }, - { - "location": { - "from": "L9:10", - "to": "L9:13" - }, - "expr": { - "_type": "VariableExpr", - "name": "df2" - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "int" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "float" - } - } - } - ] - } - }, - { - "location": { - "from": "L9:4", - "to": "L9:13" - }, - "expr": { - "_type": "BinaryExpr", - "left": { - "_type": "VariableExpr", - "name": "df1" - }, - "operator": "-", - "right": { - "_type": "VariableExpr", - "name": "df2" - } - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "int" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "float" - } - } - } - ] - } - }, { "location": { "from": "L10:4", @@ -284,7 +76,7 @@ "_type": "VariableExpr", "name": "df1" }, - "operator": "*", + "operator": "+", "right": { "_type": "VariableExpr", "name": "df2" @@ -388,7 +180,7 @@ "_type": "VariableExpr", "name": "df1" }, - "operator": "/", + "operator": "-", "right": { "_type": "VariableExpr", "name": "df2" @@ -401,7 +193,7 @@ "name": "a", "type": { "type": { - "name": "float" + "name": "int" } } }, @@ -451,8 +243,8 @@ }, { "location": { - "from": "L12:11", - "to": "L12:14" + "from": "L12:10", + "to": "L12:13" }, "expr": { "_type": "VariableExpr", @@ -484,7 +276,7 @@ { "location": { "from": "L12:4", - "to": "L12:14" + "to": "L12:13" }, "expr": { "_type": "BinaryExpr", @@ -492,7 +284,7 @@ "_type": "VariableExpr", "name": "df1" }, - "operator": "//", + "operator": "*", "right": { "_type": "VariableExpr", "name": "df2" @@ -596,7 +388,7 @@ "_type": "VariableExpr", "name": "df1" }, - "operator": "%", + "operator": "/", "right": { "_type": "VariableExpr", "name": "df2" @@ -609,7 +401,7 @@ "name": "a", "type": { "type": { - "name": "int" + "name": "float" } } }, @@ -659,8 +451,8 @@ }, { "location": { - "from": "L14:9", - "to": "L14:12" + "from": "L14:11", + "to": "L14:14" }, "expr": { "_type": "VariableExpr", @@ -692,7 +484,7 @@ { "location": { "from": "L14:4", - "to": "L14:12" + "to": "L14:14" }, "expr": { "_type": "BinaryExpr", @@ -700,7 +492,7 @@ "_type": "VariableExpr", "name": "df1" }, - "operator": "**", + "operator": "//", "right": { "_type": "VariableExpr", "name": "df2" @@ -721,7 +513,113 @@ "index": 1, "name": "b", "type": { - "type": {} + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L15:4", + "to": "L15:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L15:10", + "to": "L15:13" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L15:4", + "to": "L15:13" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "%", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } } } ] @@ -761,8 +659,8 @@ }, { "location": { - "from": "L16:10", - "to": "L16:13" + "from": "L16:9", + "to": "L16:12" }, "expr": { "_type": "VariableExpr", @@ -794,52 +692,20 @@ { "location": { "from": "L16:4", - "to": "L16:13" + "to": "L16:12" }, "expr": { - "_type": "CompareExpr", + "_type": "BinaryExpr", "left": { "_type": "VariableExpr", "name": "df1" }, - "operator": "<", + "operator": "**", "right": { "_type": "VariableExpr", "name": "df2" } }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "bool" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "bool" - } - } - } - ] - } - }, - { - "location": { - "from": "L17:4", - "to": "L17:7" - }, - "expr": { - "_type": "VariableExpr", - "name": "df1" - }, "type": { "columns": [ { @@ -855,185 +721,7 @@ "index": 1, "name": "b", "type": { - "type": { - "name": "float" - } - } - } - ] - } - }, - { - "location": { - "from": "L17:10", - "to": "L17:13" - }, - "expr": { - "_type": "VariableExpr", - "name": "df2" - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "int" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "float" - } - } - } - ] - } - }, - { - "location": { - "from": "L17:4", - "to": "L17:13" - }, - "expr": { - "_type": "CompareExpr", - "left": { - "_type": "VariableExpr", - "name": "df1" - }, - "operator": ">", - "right": { - "_type": "VariableExpr", - "name": "df2" - } - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "bool" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "bool" - } - } - } - ] - } - }, - { - "location": { - "from": "L18:4", - "to": "L18:7" - }, - "expr": { - "_type": "VariableExpr", - "name": "df1" - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "int" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "float" - } - } - } - ] - } - }, - { - "location": { - "from": "L18:11", - "to": "L18:14" - }, - "expr": { - "_type": "VariableExpr", - "name": "df2" - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "int" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "float" - } - } - } - ] - } - }, - { - "location": { - "from": "L18:4", - "to": "L18:14" - }, - "expr": { - "_type": "CompareExpr", - "left": { - "_type": "VariableExpr", - "name": "df1" - }, - "operator": "<=", - "right": { - "_type": "VariableExpr", - "name": "df2" - } - }, - "type": { - "columns": [ - { - "index": 0, - "name": "a", - "type": { - "type": { - "name": "bool" - } - } - }, - { - "index": 1, - "name": "b", - "type": { - "type": { - "name": "bool" - } + "type": {} } } ] @@ -1073,8 +761,8 @@ }, { "location": { - "from": "L19:11", - "to": "L19:14" + "from": "L19:10", + "to": "L19:13" }, "expr": { "_type": "VariableExpr", @@ -1106,7 +794,7 @@ { "location": { "from": "L19:4", - "to": "L19:14" + "to": "L19:13" }, "expr": { "_type": "CompareExpr", @@ -1114,7 +802,7 @@ "_type": "VariableExpr", "name": "df1" }, - "operator": ">=", + "operator": "<", "right": { "_type": "VariableExpr", "name": "df2" @@ -1177,8 +865,8 @@ }, { "location": { - "from": "L20:11", - "to": "L20:14" + "from": "L20:10", + "to": "L20:13" }, "expr": { "_type": "VariableExpr", @@ -1210,7 +898,7 @@ { "location": { "from": "L20:4", - "to": "L20:14" + "to": "L20:13" }, "expr": { "_type": "CompareExpr", @@ -1218,7 +906,7 @@ "_type": "VariableExpr", "name": "df1" }, - "operator": "!=", + "operator": ">", "right": { "_type": "VariableExpr", "name": "df2" @@ -1316,6 +1004,318 @@ "from": "L21:4", "to": "L21:14" }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "<=", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "bool" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "bool" + } + } + } + ] + } + }, + { + "location": { + "from": "L22:4", + "to": "L22:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L22:11", + "to": "L22:14" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L22:4", + "to": "L22:14" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": ">=", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "bool" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "bool" + } + } + } + ] + } + }, + { + "location": { + "from": "L23:4", + "to": "L23:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L23:11", + "to": "L23:14" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L23:4", + "to": "L23:14" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "df1" + }, + "operator": "!=", + "right": { + "_type": "VariableExpr", + "name": "df2" + } + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "bool" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "bool" + } + } + } + ] + } + }, + { + "location": { + "from": "L24:4", + "to": "L24:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L24:11", + "to": "L24:14" + }, + "expr": { + "_type": "VariableExpr", + "name": "df2" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L24:4", + "to": "L24:14" + }, "expr": { "_type": "CompareExpr", "left": { @@ -1350,6 +1350,654 @@ } ] } + }, + { + "location": { + "from": "L27:4", + "to": "L27:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L27:4", + "to": "L27:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "kurt" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L28:4", + "to": "L28:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L28:4", + "to": "L28:18" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "kurtosis" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L29:4", + "to": "L29:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L29:4", + "to": "L29:13" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "max" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L30:4", + "to": "L30:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L30:4", + "to": "L30:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "mean" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L31:4", + "to": "L31:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L31:4", + "to": "L31:16" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "median" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L32:4", + "to": "L32:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L32:4", + "to": "L32:13" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "min" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L33:4", + "to": "L33:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L33:4", + "to": "L33:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "mode" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L34:4", + "to": "L34:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L34:4", + "to": "L34:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "prod" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L35:4", + "to": "L35:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L35:4", + "to": "L35:17" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "product" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L36:4", + "to": "L36:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L36:4", + "to": "L36:13" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "std" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L37:4", + "to": "L37:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L37:4", + "to": "L37:13" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "sum" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L38:4", + "to": "L38:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L38:4", + "to": "L38:13" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "var" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } } ] } \ No newline at end of file From 0c70048b621acae133cf94c490586fe8b276f4da Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 01:34:58 +0200 Subject: [PATCH 29/36] feat(checker): add statistical ops on columns --- midas/checker/frames/column_methods.py | 66 ++++++++++++++++++++++++-- 1 file changed, 63 insertions(+), 3 deletions(-) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index c9da035..0f10fcf 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -160,8 +160,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): def eq(self, call: Call) -> Type: return self._element_wise(call, "__eq__") - @method() - def mean(self, call: Call) -> Type: + def _statistical(self, call: Call, kwargs: list[Function.Argument] = []) -> Type: signature = Function( kw_args=[ Function.Argument( @@ -169,7 +168,8 @@ class ColumnMethodRegistry(MethodRegistry[Call]): name="axis", type=TopType(), required=False, - ) + ), + *kwargs, ], returns=ColumnType(type=TopType()), ) @@ -182,6 +182,66 @@ class ColumnMethodRegistry(MethodRegistry[Call]): ) return result.result + @method("kurtosis", "kurt") + def kurtosis(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def max(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def mean(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def median(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def min(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def mode(self, call: Call) -> Type: + return self._statistical(call) + + @method("product", "prod") + def product(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def std(self, call: Call) -> Type: + return self._statistical( + call, + [ + Function.Argument( + pos=1, + name="ddof", + type=self.types.get_type("int"), + required=False, + ) + ], + ) + + @method() + def sum(self, call: Call) -> Type: + return self._statistical(call) + + @method() + def var(self, call: Call) -> Type: + return self._statistical( + call, + [ + Function.Argument( + pos=1, + name="var", + type=self.types.get_type("int"), + required=False, + ) + ], + ) + @method() def groupby(self, call: Call) -> Type: bool_: Type = self.types.get_type("bool") From a143972ef1661f2925c0fa89fb8c1d49f94086bb Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 02:20:51 +0200 Subject: [PATCH 30/36] feat(checker): add aggregation ops on frame groupby --- midas/checker/frames/frame_groupby_methods.py | 183 ++++- midas/checker/frames/frame_methods.py | 22 +- tests/cases/checker/09_frame_ops.py | 15 +- tests/cases/checker/09_frame_ops.py.ref.json | 768 ++++++++++++++++++ 4 files changed, 947 insertions(+), 41 deletions(-) diff --git a/midas/checker/frames/frame_groupby_methods.py b/midas/checker/frames/frame_groupby_methods.py index d3f882d..615b2b3 100644 --- a/midas/checker/frames/frame_groupby_methods.py +++ b/midas/checker/frames/frame_groupby_methods.py @@ -28,36 +28,37 @@ class Call: class FrameGroupByMethodRegistry(MethodRegistry[Call]): - @method() - def mean(self, call: Call) -> Type: - bool_ = self.types.get_type("bool") + NAMED_ARGS: dict[str, str] = { + "numeric_only": "bool", + "skipna": "bool", + "engine": "str", + "engine_kwargs": "dict", + } + + def _aggregate( + self, call: Call, args: list[str | tuple[str, str, bool]] = [] + ) -> Type: + real_args: list[Function.Argument] = [] + for i, arg in enumerate(args): + match arg: + case str() as name: + arg = Function.Argument( + pos=i, + name=name, + type=self.types.get_type(self.NAMED_ARGS[name]), + required=False, + ) + case (name, type, required): + arg = Function.Argument( + pos=i, + name=name, + type=self.types.get_type(type), + required=required, + ) + real_args.append(arg) + signature = Function( - args=[ - Function.Argument( - pos=0, - name="numeric_only", - type=bool_, - required=False, - ), - Function.Argument( - pos=1, - name="skipna", - type=bool_, - required=False, - ), - Function.Argument( - pos=2, - name="engine", - type=self.types.get_type("str"), - required=False, - ), - Function.Argument( - pos=3, - name="engine_kwargs", - type=self.types.get_type("dict"), - required=False, - ), - ], + args=real_args, returns=call.groupby.frame, ) @@ -68,3 +69,127 @@ class FrameGroupByMethodRegistry(MethodRegistry[Call]): keywords=call.keywords, ) return result.result + + @method() + def kurt(self, call: Call) -> Type: + return self._aggregate( + call, + [ + "skipna", + "numeric_only", + ], + ) + + @method() + def max(self, call: Call) -> Type: + return self._aggregate( + call, + [ + "numeric_only", + ( + "min_count", + "int", + False, + ), + "skipna", + "engine", + "engine_kwargs", + ], + ) + + @method() + def mean(self, call: Call) -> Type: + return self._aggregate( + call, + ["numeric_only", "skipna", "engine", "engine_kwargs"], + ) + + @method() + def median(self, call: Call) -> Type: + return self._aggregate( + call, + ["numeric_only", "skipna"], + ) + + @method() + def min(self, call: Call) -> Type: + return self._aggregate( + call, + [ + "numeric_only", + ( + "min_count", + "int", + False, + ), + "skipna", + "engine", + "engine_kwargs", + ], + ) + + @method() + def prod(self, call: Call) -> Type: + return self._aggregate( + call, + [ + "numeric_only", + ( + "min_count", + "int", + False, + ), + "skipna", + ], + ) + + @method() + def std(self, call: Call) -> Type: + return self._aggregate( + call, + [ + ( + "ddof", + "int", + False, + ), + "engine", + "engine_kwargs", + "numeric_only", + "skipna", + ], + ) + + @method() + def sum(self, call: Call) -> Type: + return self._aggregate( + call, + [ + "numeric_only", + ( + "min_count", + "int", + False, + ), + "skipna", + "engine", + "engine_kwargs", + ], + ) + + @method() + def var(self, call: Call) -> Type: + return self._aggregate( + call, + [ + ( + "var", + "int", + False, + ), + "engine", + "engine_kwargs", + "numeric_only", + "skipna", + ], + ) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 5c7b4b4..2bd8c8b 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -227,7 +227,7 @@ class FrameMethodRegistry(MethodRegistry[Call]): def eq(self, call: Call) -> Type: return self._element_wise(call, "__eq__") - def _statistical(self, call: Call, kwargs: list[Function.Argument] = []) -> Type: + def _aggregate(self, call: Call, kwargs: list[Function.Argument] = []) -> Type: with_axis = Function( kw_args=[ Function.Argument( @@ -269,35 +269,35 @@ class FrameMethodRegistry(MethodRegistry[Call]): @method("kurtosis", "kurt") def kurtosis(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def max(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def mean(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def median(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def min(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def mode(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method("product", "prod") def product(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def std(self, call: Call) -> Type: - return self._statistical( + return self._aggregate( call, [ Function.Argument( @@ -311,11 +311,11 @@ class FrameMethodRegistry(MethodRegistry[Call]): @method() def sum(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def var(self, call: Call) -> Type: - return self._statistical( + return self._aggregate( call, [ Function.Argument( diff --git a/tests/cases/checker/09_frame_ops.py b/tests/cases/checker/09_frame_ops.py index 3bd25cf..0c2f00c 100644 --- a/tests/cases/checker/09_frame_ops.py +++ b/tests/cases/checker/09_frame_ops.py @@ -23,7 +23,7 @@ _ = df1 >= df2 _ = df1 != df2 _ = df1 == df2 -# Statistical +# Aggregate _ = df1.kurt() _ = df1.kurtosis() _ = df1.max() @@ -36,3 +36,16 @@ _ = df1.product() _ = df1.std() _ = df1.sum() _ = df1.var() + +# Groupby +gb = df1.groupby(by="a") + +_ = gb.kurt() +_ = gb.max() +_ = gb.mean() +_ = gb.median() +_ = gb.min() +_ = gb.prod() +_ = gb.std() +_ = gb.sum() +_ = gb.var() diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json index 551b8dc..d7e9b4a 100644 --- a/tests/cases/checker/09_frame_ops.py.ref.json +++ b/tests/cases/checker/09_frame_ops.py.ref.json @@ -1998,6 +1998,774 @@ "type": { "type": {} } + }, + { + "location": { + "from": "L41:20", + "to": "L41:23" + }, + "expr": { + "_type": "LiteralExpr", + "value": "a" + }, + "type": { + "name": "str" + } + }, + { + "location": { + "from": "L41:5", + "to": "L41:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L41:5", + "to": "L41:24" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "groupby" + }, + "arguments": [], + "keywords": { + "by": { + "_type": "LiteralExpr", + "value": "a" + } + } + }, + "type": { + "frame": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + } + }, + { + "location": { + "from": "L43:4", + "to": "L43:6" + }, + "expr": { + "_type": "VariableExpr", + "name": "gb" + }, + "type": { + "frame": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + } + }, + { + "location": { + "from": "L43:4", + "to": "L43:13" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "gb" + }, + "name": "kurt" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L44:4", + "to": "L44:6" + }, + "expr": { + "_type": "VariableExpr", + "name": "gb" + }, + "type": { + "frame": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + } + }, + { + "location": { + "from": "L44:4", + "to": "L44:12" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "gb" + }, + "name": "max" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L45:4", + "to": "L45:6" + }, + "expr": { + "_type": "VariableExpr", + "name": "gb" + }, + "type": { + "frame": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + } + }, + { + "location": { + "from": "L45:4", + "to": "L45:13" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "gb" + }, + "name": "mean" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L46:4", + "to": "L46:6" + }, + "expr": { + "_type": "VariableExpr", + "name": "gb" + }, + "type": { + "frame": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + } + }, + { + "location": { + "from": "L46:4", + "to": "L46:15" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "gb" + }, + "name": "median" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L47:4", + "to": "L47:6" + }, + "expr": { + "_type": "VariableExpr", + "name": "gb" + }, + "type": { + "frame": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + } + }, + { + "location": { + "from": "L47:4", + "to": "L47:12" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "gb" + }, + "name": "min" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L48:4", + "to": "L48:6" + }, + "expr": { + "_type": "VariableExpr", + "name": "gb" + }, + "type": { + "frame": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + } + }, + { + "location": { + "from": "L48:4", + "to": "L48:13" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "gb" + }, + "name": "prod" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L49:4", + "to": "L49:6" + }, + "expr": { + "_type": "VariableExpr", + "name": "gb" + }, + "type": { + "frame": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + } + }, + { + "location": { + "from": "L49:4", + "to": "L49:12" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "gb" + }, + "name": "std" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L50:4", + "to": "L50:6" + }, + "expr": { + "_type": "VariableExpr", + "name": "gb" + }, + "type": { + "frame": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + } + }, + { + "location": { + "from": "L50:4", + "to": "L50:12" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "gb" + }, + "name": "sum" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L51:4", + "to": "L51:6" + }, + "expr": { + "_type": "VariableExpr", + "name": "gb" + }, + "type": { + "frame": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + } + }, + { + "location": { + "from": "L51:4", + "to": "L51:12" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "gb" + }, + "name": "var" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } } ] } \ No newline at end of file From 20173a0b07f10c62d107e4ea9eeb0c7debd0aaf9 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 10:58:28 +0200 Subject: [PATCH 31/36] feat(tests): add colors and run all tests in base module --- tests/__main__.py | 43 +++++++++++++++++++++++++++++++++++++++++++ tests/base.py | 12 +++++++++--- 2 files changed, 52 insertions(+), 3 deletions(-) create mode 100644 tests/__main__.py diff --git a/tests/__main__.py b/tests/__main__.py new file mode 100644 index 0000000..0280e3b --- /dev/null +++ b/tests/__main__.py @@ -0,0 +1,43 @@ +from typing import Type + +from midas.cli.ansi import Ansi +from tests.base import Tester +from tests.checker import CheckerTester +from tests.generator import GeneratorTester +from tests.midas import MidasTester +from tests.python import PythonTester + + +def print_banner(name: str): + horizontal: str = "+" + "-" * (len(name) + 2) + "+" + print(horizontal) + print(f"| {name} |") + print(horizontal) + + +def run_tests(tester_cls: Type[Tester]) -> bool: + print_banner(tester_cls.__name__) + tester: Tester = tester_cls() + success: bool = tester.run_all_tests() + print() + return success + + +def main(): + testers: list[Type[Tester]] = [ + PythonTester, + MidasTester, + CheckerTester, + GeneratorTester, + ] + + success: bool = all(map(run_tests, testers)) + + if success: + print(Ansi.FG(Ansi.BRIGHT_GREEN) + "All tests passed!" + Ansi.RESET) + else: + print(Ansi.FG(Ansi.BRIGHT_RED) + "Some tests failed!" + Ansi.RESET) + + +if __name__ == "__main__": + main() diff --git a/tests/base.py b/tests/base.py index c8bad17..353c3e8 100644 --- a/tests/base.py +++ b/tests/base.py @@ -7,6 +7,8 @@ from abc import ABC, abstractmethod from pathlib import Path from typing import Iterator, Protocol +from midas.cli.ansi import Ansi + class CaseResult(Protocol): def dumps(self) -> str: ... @@ -44,8 +46,11 @@ class Tester(ABC): print(rule) for i, test in enumerate(tests): - print(f"Case {i+1}/{n}: {test.resolve().relative_to(self.CASES_DIR)}") + path: Path = test.resolve().relative_to(self.CASES_DIR) + print(f"{Ansi.FG(Ansi.BRIGHT_CYAN)}Case {i+1}/{n}: {path}{Ansi.RESET}") + print(Ansi.DIM, end="") success: bool = self._run_test(test) + print(Ansi.RESET, end="") if success: successes += 1 else: @@ -146,8 +151,9 @@ class Tester(ABC): if not success: sys.exit(1) case None: - print("No subcommand provided. Available subcommands: run, update") - sys.exit(1) + success: bool = tester.run_all_tests() + if not success: + sys.exit(1) case _: print(f"Unknown subcommand '{args.subcommand}'") sys.exit(1) From 733c8736b8f571f5c80286de05e2761d5355c8b3 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 11:25:06 +0200 Subject: [PATCH 32/36] feat(checker): add aggregation ops on column groupby --- .../checker/frames/column_groupby_methods.py | 195 +- midas/checker/frames/column_methods.py | 30 +- tests/cases/checker/09_frame_ops.py | 70 +- tests/cases/checker/09_frame_ops.py.ref.json | 1749 ++++++++++++++++- 4 files changed, 1949 insertions(+), 95 deletions(-) diff --git a/midas/checker/frames/column_groupby_methods.py b/midas/checker/frames/column_groupby_methods.py index aaed5de..7470f46 100644 --- a/midas/checker/frames/column_groupby_methods.py +++ b/midas/checker/frames/column_groupby_methods.py @@ -7,7 +7,7 @@ import midas.ast.python as p from midas.ast.location import Location from midas.checker.dispatcher import CallResult from midas.checker.frames.utils import MethodRegistry, method -from midas.checker.types import ColumnGroupBy, Function, Type +from midas.checker.types import ColumnGroupBy, ColumnType, Function, TopType, Type if TYPE_CHECKING: from midas.checker.python import TypedExpr @@ -28,37 +28,46 @@ class Call: class ColumnGroupByMethodRegistry(MethodRegistry[Call]): - @method() - def mean(self, call: Call) -> Type: - bool_ = self.types.get_type("bool") + NAMED_ARGS: dict[str, str] = { + "numeric_only": "bool", + "skipna": "bool", + "engine": "str", + "engine_kwargs": "dict", + } + + def _aggregate( + self, + call: Call, + args: list[str | tuple[str, str, bool]] = [], + *, + preserve_inner_type: bool = False, + ) -> Type: + real_args: list[Function.Argument] = [] + for i, arg in enumerate(args): + match arg: + case str() as name: + arg = Function.Argument( + pos=i, + name=name, + type=self.types.get_type(self.NAMED_ARGS[name]), + required=False, + ) + case (name, type, required): + arg = Function.Argument( + pos=i, + name=name, + type=self.types.get_type(type), + required=required, + ) + real_args.append(arg) + signature = Function( - args=[ - Function.Argument( - pos=0, - name="numeric_only", - type=bool_, - required=False, - ), - Function.Argument( - pos=1, - name="skipna", - type=bool_, - required=False, - ), - Function.Argument( - pos=2, - name="engine", - type=self.types.get_type("str"), - required=False, - ), - Function.Argument( - pos=3, - name="engine_kwargs", - type=self.types.get_type("dict"), - required=False, - ), - ], - returns=call.groupby.column, + args=real_args, + returns=( + call.groupby.column + if preserve_inner_type + else ColumnType(type=TopType()) + ), ) result: CallResult = self.dispatcher.get_result( @@ -68,3 +77,127 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]): keywords=call.keywords, ) return result.result + + @method() + def kurt(self, call: Call) -> Type: + return self._aggregate( + call, + ["skipna", "numeric_only"], + ) + + @method() + def max(self, call: Call) -> Type: + return self._aggregate( + call, + [ + "numeric_only", + ( + "min_count", + "int", + False, + ), + "skipna", + "engine", + "engine_kwargs", + ], + preserve_inner_type=True, + ) + + @method() + def mean(self, call: Call) -> Type: + return self._aggregate( + call, + ["numeric_only", "skipna", "engine", "engine_kwargs"], + ) + + @method() + def median(self, call: Call) -> Type: + return self._aggregate( + call, + ["numeric_only", "skipna"], + preserve_inner_type=True, + ) + + @method() + def min(self, call: Call) -> Type: + return self._aggregate( + call, + [ + "numeric_only", + ( + "min_count", + "int", + False, + ), + "skipna", + "engine", + "engine_kwargs", + ], + preserve_inner_type=True, + ) + + @method() + def prod(self, call: Call) -> Type: + return self._aggregate( + call, + [ + "numeric_only", + ( + "min_count", + "int", + False, + ), + "skipna", + ], + ) + + @method() + def std(self, call: Call) -> Type: + return self._aggregate( + call, + [ + ( + "ddof", + "int", + False, + ), + "engine", + "engine_kwargs", + "numeric_only", + "skipna", + ], + ) + + @method() + def sum(self, call: Call) -> Type: + return self._aggregate( + call, + [ + "numeric_only", + ( + "min_count", + "int", + False, + ), + "skipna", + "engine", + "engine_kwargs", + ], + ) + + @method() + def var(self, call: Call) -> Type: + return self._aggregate( + call, + [ + ( + "var", + "int", + False, + ), + "engine", + "engine_kwargs", + "numeric_only", + "skipna", + ], + ) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 0f10fcf..d2559b8 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -160,7 +160,13 @@ class ColumnMethodRegistry(MethodRegistry[Call]): def eq(self, call: Call) -> Type: return self._element_wise(call, "__eq__") - def _statistical(self, call: Call, kwargs: list[Function.Argument] = []) -> Type: + def _aggregate( + self, + call: Call, + kwargs: list[Function.Argument] = [], + *, + preserve_inner_type: bool = False, + ) -> Type: signature = Function( kw_args=[ Function.Argument( @@ -171,7 +177,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]): ), *kwargs, ], - returns=ColumnType(type=TopType()), + returns=call.column if preserve_inner_type else ColumnType(type=TopType()), ) result: CallResult = self.dispatcher.get_result( @@ -184,35 +190,35 @@ class ColumnMethodRegistry(MethodRegistry[Call]): @method("kurtosis", "kurt") def kurtosis(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def max(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call, preserve_inner_type=True) @method() def mean(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def median(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call, preserve_inner_type=True) @method() def min(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call, preserve_inner_type=True) @method() def mode(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call, preserve_inner_type=True) @method("product", "prod") def product(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def std(self, call: Call) -> Type: - return self._statistical( + return self._aggregate( call, [ Function.Argument( @@ -226,11 +232,11 @@ class ColumnMethodRegistry(MethodRegistry[Call]): @method() def sum(self, call: Call) -> Type: - return self._statistical(call) + return self._aggregate(call) @method() def var(self, call: Call) -> Type: - return self._statistical( + return self._aggregate( call, [ Function.Argument( diff --git a/tests/cases/checker/09_frame_ops.py b/tests/cases/checker/09_frame_ops.py index 0c2f00c..18951f5 100644 --- a/tests/cases/checker/09_frame_ops.py +++ b/tests/cases/checker/09_frame_ops.py @@ -38,14 +38,64 @@ _ = df1.sum() _ = df1.var() # Groupby -gb = df1.groupby(by="a") +df_gb = df1.groupby(by="a") -_ = gb.kurt() -_ = gb.max() -_ = gb.mean() -_ = gb.median() -_ = gb.min() -_ = gb.prod() -_ = gb.std() -_ = gb.sum() -_ = gb.var() +_ = df_gb.kurt() +_ = df_gb.max() +_ = df_gb.mean() +_ = df_gb.median() +_ = df_gb.min() +_ = df_gb.prod() +_ = df_gb.std() +_ = df_gb.sum() +_ = df_gb.var() + + +# Columns + +col1 = df1["a"] +col2 = df1["a"] + +# Arithmetic +_ = col1 + col2 +_ = col1 - col2 +_ = col1 * col2 +_ = col1 / col2 +_ = col1 // col2 +_ = col1 % col2 +_ = col1**col2 + +# Comparisons +_ = col1 < col2 +_ = col1 > col2 +_ = col1 <= col2 +_ = col1 >= col2 +_ = col1 != col2 +_ = col1 == col2 + +# Aggregate +_ = col1.kurt() +_ = col1.kurtosis() +_ = col1.max() +_ = col1.mean() +_ = col1.median() +_ = col1.min() +_ = col1.mode() +_ = col1.prod() +_ = col1.product() +_ = col1.std() +_ = col1.sum() +_ = col1.var() + +# Groupby +col_gb = col1.groupby(level=0) + +_ = col_gb.kurt() +_ = col_gb.max() +_ = col_gb.mean() +_ = col_gb.median() +_ = col_gb.min() +_ = col_gb.prod() +_ = col_gb.std() +_ = col_gb.sum() +_ = col_gb.var() diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json index d7e9b4a..37e4aa9 100644 --- a/tests/cases/checker/09_frame_ops.py.ref.json +++ b/tests/cases/checker/09_frame_ops.py.ref.json @@ -2001,8 +2001,8 @@ }, { "location": { - "from": "L41:20", - "to": "L41:23" + "from": "L41:23", + "to": "L41:26" }, "expr": { "_type": "LiteralExpr", @@ -2014,8 +2014,8 @@ }, { "location": { - "from": "L41:5", - "to": "L41:8" + "from": "L41:8", + "to": "L41:11" }, "expr": { "_type": "VariableExpr", @@ -2046,8 +2046,8 @@ }, { "location": { - "from": "L41:5", - "to": "L41:24" + "from": "L41:8", + "to": "L41:27" }, "expr": { "_type": "CallExpr", @@ -2095,11 +2095,11 @@ { "location": { "from": "L43:4", - "to": "L43:6" + "to": "L43:9" }, "expr": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "type": { "frame": { @@ -2129,7 +2129,7 @@ { "location": { "from": "L43:4", - "to": "L43:13" + "to": "L43:16" }, "expr": { "_type": "CallExpr", @@ -2137,7 +2137,7 @@ "_type": "GetExpr", "object": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "name": "kurt" }, @@ -2170,11 +2170,11 @@ { "location": { "from": "L44:4", - "to": "L44:6" + "to": "L44:9" }, "expr": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "type": { "frame": { @@ -2204,7 +2204,7 @@ { "location": { "from": "L44:4", - "to": "L44:12" + "to": "L44:15" }, "expr": { "_type": "CallExpr", @@ -2212,7 +2212,7 @@ "_type": "GetExpr", "object": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "name": "max" }, @@ -2245,11 +2245,11 @@ { "location": { "from": "L45:4", - "to": "L45:6" + "to": "L45:9" }, "expr": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "type": { "frame": { @@ -2279,7 +2279,7 @@ { "location": { "from": "L45:4", - "to": "L45:13" + "to": "L45:16" }, "expr": { "_type": "CallExpr", @@ -2287,7 +2287,7 @@ "_type": "GetExpr", "object": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "name": "mean" }, @@ -2320,11 +2320,11 @@ { "location": { "from": "L46:4", - "to": "L46:6" + "to": "L46:9" }, "expr": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "type": { "frame": { @@ -2354,7 +2354,7 @@ { "location": { "from": "L46:4", - "to": "L46:15" + "to": "L46:18" }, "expr": { "_type": "CallExpr", @@ -2362,7 +2362,7 @@ "_type": "GetExpr", "object": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "name": "median" }, @@ -2395,11 +2395,11 @@ { "location": { "from": "L47:4", - "to": "L47:6" + "to": "L47:9" }, "expr": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "type": { "frame": { @@ -2429,7 +2429,7 @@ { "location": { "from": "L47:4", - "to": "L47:12" + "to": "L47:15" }, "expr": { "_type": "CallExpr", @@ -2437,7 +2437,7 @@ "_type": "GetExpr", "object": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "name": "min" }, @@ -2470,11 +2470,11 @@ { "location": { "from": "L48:4", - "to": "L48:6" + "to": "L48:9" }, "expr": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "type": { "frame": { @@ -2504,7 +2504,7 @@ { "location": { "from": "L48:4", - "to": "L48:13" + "to": "L48:16" }, "expr": { "_type": "CallExpr", @@ -2512,7 +2512,7 @@ "_type": "GetExpr", "object": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "name": "prod" }, @@ -2545,11 +2545,11 @@ { "location": { "from": "L49:4", - "to": "L49:6" + "to": "L49:9" }, "expr": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "type": { "frame": { @@ -2579,7 +2579,7 @@ { "location": { "from": "L49:4", - "to": "L49:12" + "to": "L49:15" }, "expr": { "_type": "CallExpr", @@ -2587,7 +2587,7 @@ "_type": "GetExpr", "object": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "name": "std" }, @@ -2620,11 +2620,11 @@ { "location": { "from": "L50:4", - "to": "L50:6" + "to": "L50:9" }, "expr": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "type": { "frame": { @@ -2654,7 +2654,7 @@ { "location": { "from": "L50:4", - "to": "L50:12" + "to": "L50:15" }, "expr": { "_type": "CallExpr", @@ -2662,7 +2662,7 @@ "_type": "GetExpr", "object": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "name": "sum" }, @@ -2695,11 +2695,11 @@ { "location": { "from": "L51:4", - "to": "L51:6" + "to": "L51:9" }, "expr": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "type": { "frame": { @@ -2729,7 +2729,7 @@ { "location": { "from": "L51:4", - "to": "L51:12" + "to": "L51:15" }, "expr": { "_type": "CallExpr", @@ -2737,7 +2737,7 @@ "_type": "GetExpr", "object": { "_type": "VariableExpr", - "name": "gb" + "name": "df_gb" }, "name": "var" }, @@ -2766,6 +2766,1671 @@ } ] } + }, + { + "location": { + "from": "L56:7", + "to": "L56:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L56:7", + "to": "L56:15" + }, + "expr": { + "_type": "SubscriptExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "index": { + "_type": "LiteralExpr", + "value": "a" + } + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L57:7", + "to": "L57:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L57:7", + "to": "L57:15" + }, + "expr": { + "_type": "SubscriptExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "index": { + "_type": "LiteralExpr", + "value": "a" + } + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L60:4", + "to": "L60:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L60:11", + "to": "L60:15" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L60:4", + "to": "L60:15" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "+", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L61:4", + "to": "L61:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L61:11", + "to": "L61:15" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L61:4", + "to": "L61:15" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "-", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L62:4", + "to": "L62:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L62:11", + "to": "L62:15" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L62:4", + "to": "L62:15" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "*", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L63:4", + "to": "L63:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L63:11", + "to": "L63:15" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L63:4", + "to": "L63:15" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "/", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "float" + } + } + }, + { + "location": { + "from": "L64:4", + "to": "L64:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L64:12", + "to": "L64:16" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L64:4", + "to": "L64:16" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "//", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L65:4", + "to": "L65:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L65:11", + "to": "L65:15" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L65:4", + "to": "L65:15" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "%", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L66:4", + "to": "L66:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L66:10", + "to": "L66:14" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L66:4", + "to": "L66:14" + }, + "expr": { + "_type": "BinaryExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "**", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L69:4", + "to": "L69:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L69:11", + "to": "L69:15" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L69:4", + "to": "L69:15" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "<", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "bool" + } + } + }, + { + "location": { + "from": "L70:4", + "to": "L70:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L70:11", + "to": "L70:15" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L70:4", + "to": "L70:15" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": ">", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "bool" + } + } + }, + { + "location": { + "from": "L71:4", + "to": "L71:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L71:12", + "to": "L71:16" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L71:4", + "to": "L71:16" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "<=", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "bool" + } + } + }, + { + "location": { + "from": "L72:4", + "to": "L72:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L72:12", + "to": "L72:16" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L72:4", + "to": "L72:16" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": ">=", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "bool" + } + } + }, + { + "location": { + "from": "L73:4", + "to": "L73:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L73:12", + "to": "L73:16" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L73:4", + "to": "L73:16" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "!=", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "bool" + } + } + }, + { + "location": { + "from": "L74:4", + "to": "L74:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L74:12", + "to": "L74:16" + }, + "expr": { + "_type": "VariableExpr", + "name": "col2" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L74:4", + "to": "L74:16" + }, + "expr": { + "_type": "CompareExpr", + "left": { + "_type": "VariableExpr", + "name": "col1" + }, + "operator": "==", + "right": { + "_type": "VariableExpr", + "name": "col2" + } + }, + "type": { + "type": { + "name": "bool" + } + } + }, + { + "location": { + "from": "L77:4", + "to": "L77:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L77:4", + "to": "L77:15" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "kurt" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L78:4", + "to": "L78:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L78:4", + "to": "L78:19" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "kurtosis" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L79:4", + "to": "L79:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L79:4", + "to": "L79:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "max" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L80:4", + "to": "L80:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L80:4", + "to": "L80:15" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "mean" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L81:4", + "to": "L81:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L81:4", + "to": "L81:17" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "median" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L82:4", + "to": "L82:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L82:4", + "to": "L82:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "min" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L83:4", + "to": "L83:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L83:4", + "to": "L83:15" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "mode" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L84:4", + "to": "L84:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L84:4", + "to": "L84:15" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "prod" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L85:4", + "to": "L85:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L85:4", + "to": "L85:18" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "product" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L86:4", + "to": "L86:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L86:4", + "to": "L86:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "std" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L87:4", + "to": "L87:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L87:4", + "to": "L87:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "sum" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L88:4", + "to": "L88:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L88:4", + "to": "L88:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "var" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L91:28", + "to": "L91:29" + }, + "expr": { + "_type": "LiteralExpr", + "value": 0 + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L91:9", + "to": "L91:13" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L91:9", + "to": "L91:30" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "groupby" + }, + "arguments": [], + "keywords": { + "level": { + "_type": "LiteralExpr", + "value": 0 + } + } + }, + "type": { + "column": { + "type": { + "name": "int" + } + } + } + }, + { + "location": { + "from": "L93:4", + "to": "L93:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "type": { + "column": { + "type": { + "name": "int" + } + } + } + }, + { + "location": { + "from": "L93:4", + "to": "L93:17" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "name": "kurt" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L94:4", + "to": "L94:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "type": { + "column": { + "type": { + "name": "int" + } + } + } + }, + { + "location": { + "from": "L94:4", + "to": "L94:16" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "name": "max" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L95:4", + "to": "L95:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "type": { + "column": { + "type": { + "name": "int" + } + } + } + }, + { + "location": { + "from": "L95:4", + "to": "L95:17" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "name": "mean" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L96:4", + "to": "L96:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "type": { + "column": { + "type": { + "name": "int" + } + } + } + }, + { + "location": { + "from": "L96:4", + "to": "L96:19" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "name": "median" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L97:4", + "to": "L97:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "type": { + "column": { + "type": { + "name": "int" + } + } + } + }, + { + "location": { + "from": "L97:4", + "to": "L97:16" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "name": "min" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L98:4", + "to": "L98:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "type": { + "column": { + "type": { + "name": "int" + } + } + } + }, + { + "location": { + "from": "L98:4", + "to": "L98:17" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "name": "prod" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L99:4", + "to": "L99:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "type": { + "column": { + "type": { + "name": "int" + } + } + } + }, + { + "location": { + "from": "L99:4", + "to": "L99:16" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "name": "std" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L100:4", + "to": "L100:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "type": { + "column": { + "type": { + "name": "int" + } + } + } + }, + { + "location": { + "from": "L100:4", + "to": "L100:16" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "name": "sum" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } + }, + { + "location": { + "from": "L101:4", + "to": "L101:10" + }, + "expr": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "type": { + "column": { + "type": { + "name": "int" + } + } + } + }, + { + "location": { + "from": "L101:4", + "to": "L101:16" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col_gb" + }, + "name": "var" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": {} + } } ] } \ No newline at end of file From 7a6e01cff8cd374fb67e18aae6156d762c7c5150 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 11:42:35 +0200 Subject: [PATCH 33/36] fix(checker): delegate frame aggregate methods to columns --- midas/checker/frames/frame_groupby_methods.py | 170 ++++-------------- midas/checker/python.py | 29 ++- tests/cases/checker/09_frame_ops.py.ref.json | 48 ++--- 3 files changed, 61 insertions(+), 186 deletions(-) diff --git a/midas/checker/frames/frame_groupby_methods.py b/midas/checker/frames/frame_groupby_methods.py index 615b2b3..755cab7 100644 --- a/midas/checker/frames/frame_groupby_methods.py +++ b/midas/checker/frames/frame_groupby_methods.py @@ -5,9 +5,15 @@ from typing import TYPE_CHECKING import midas.ast.python as p from midas.ast.location import Location -from midas.checker.dispatcher import CallResult from midas.checker.frames.utils import MethodRegistry, method -from midas.checker.types import FrameGroupBy, Function, Type +from midas.checker.types import ( + ColumnGroupBy, + ColumnType, + DataFrameType, + FrameGroupBy, + Type, + UnknownType, +) if TYPE_CHECKING: from midas.checker.python import TypedExpr @@ -35,161 +41,63 @@ class FrameGroupByMethodRegistry(MethodRegistry[Call]): "engine_kwargs": "dict", } - def _aggregate( - self, call: Call, args: list[str | tuple[str, str, bool]] = [] - ) -> Type: - real_args: list[Function.Argument] = [] - for i, arg in enumerate(args): - match arg: - case str() as name: - arg = Function.Argument( - pos=i, - name=name, - type=self.types.get_type(self.NAMED_ARGS[name]), - required=False, - ) - case (name, type, required): - arg = Function.Argument( - pos=i, - name=name, - type=self.types.get_type(type), - required=required, - ) - real_args.append(arg) + def _aggregate(self, call: Call, method: str) -> Type: + new_columns: list[DataFrameType.Column] = [] - signature = Function( - args=real_args, - returns=call.groupby.frame, - ) + for column in call.groupby.frame.columns: + column_groupby: ColumnGroupBy = ColumnGroupBy(column=column.type) + result_type: Type = self.typer.call_method( + location=call.location, + call_expr=call.call_expr, + obj=(call.groupby_expr, column_groupby), + method_name=method, + positional=call.positional, + keywords=call.keywords, + ) + if not isinstance(result_type, ColumnType): + result_type = ColumnType(type=UnknownType()) + new_columns.append( + DataFrameType.Column( + index=column.index, + name=column.name, + type=result_type, + ) + ) - result: CallResult = self.dispatcher.get_result( - location=call.location, - callee=signature, - positional=call.positional, - keywords=call.keywords, - ) - return result.result + return DataFrameType(columns=new_columns) @method() def kurt(self, call: Call) -> Type: - return self._aggregate( - call, - [ - "skipna", - "numeric_only", - ], - ) + return self._aggregate(call, "kurt") @method() def max(self, call: Call) -> Type: - return self._aggregate( - call, - [ - "numeric_only", - ( - "min_count", - "int", - False, - ), - "skipna", - "engine", - "engine_kwargs", - ], - ) + return self._aggregate(call, "max") @method() def mean(self, call: Call) -> Type: - return self._aggregate( - call, - ["numeric_only", "skipna", "engine", "engine_kwargs"], - ) + return self._aggregate(call, "mean") @method() def median(self, call: Call) -> Type: - return self._aggregate( - call, - ["numeric_only", "skipna"], - ) + return self._aggregate(call, "median") @method() def min(self, call: Call) -> Type: - return self._aggregate( - call, - [ - "numeric_only", - ( - "min_count", - "int", - False, - ), - "skipna", - "engine", - "engine_kwargs", - ], - ) + return self._aggregate(call, "min") @method() def prod(self, call: Call) -> Type: - return self._aggregate( - call, - [ - "numeric_only", - ( - "min_count", - "int", - False, - ), - "skipna", - ], - ) + return self._aggregate(call, "prod") @method() def std(self, call: Call) -> Type: - return self._aggregate( - call, - [ - ( - "ddof", - "int", - False, - ), - "engine", - "engine_kwargs", - "numeric_only", - "skipna", - ], - ) + return self._aggregate(call, "std") @method() def sum(self, call: Call) -> Type: - return self._aggregate( - call, - [ - "numeric_only", - ( - "min_count", - "int", - False, - ), - "skipna", - "engine", - "engine_kwargs", - ], - ) + return self._aggregate(call, "sum") @method() def var(self, call: Call) -> Type: - return self._aggregate( - call, - [ - ( - "var", - "int", - False, - ), - "engine", - "engine_kwargs", - "numeric_only", - "skipna", - ], - ) + return self._aggregate(call, "var") diff --git a/midas/checker/python.py b/midas/checker/python.py index d095162..99766d8 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -222,7 +222,7 @@ class PythonTyper( method_name: str, positional: list[TypedExpr], keywords: dict[str, TypedExpr], - ) -> Optional[Type]: + ) -> Type: unfolded: Type = unfold_type(obj[1]) match unfolded: case DataFrameType(): @@ -580,9 +580,8 @@ class PythonTyper( right: TypedExpr, method: str, ) -> Type: - result: Optional[Type] try: - result = self.call_method( + return self.call_method( location=location, call_expr=expr, obj=left, @@ -597,8 +596,6 @@ class PythonTyper( ) return UnknownType() - return result or UnknownType() - def visit_unary_expr(self, expr: p.UnaryExpr) -> Type: method: Optional[str] = PY_UNARY_METHODS.get(expr.operator.__class__) if method is None: @@ -610,9 +607,8 @@ class PythonTyper( operand: Type = self.type_of(expr.right) - result: Optional[Type] try: - result = self.call_method( + return self.call_method( location=expr.location, call_expr=expr, obj=(expr.right, operand), @@ -627,8 +623,6 @@ class PythonTyper( ) return UnknownType() - return result or UnknownType() - def visit_call_expr(self, expr: p.CallExpr) -> Type: match expr.callee: case p.VariableExpr(name="TypeVar"): @@ -644,16 +638,13 @@ class PythonTyper( match expr.callee: case p.GetExpr(object=obj, name=method): obj_type: Type = self.type_of(obj) - return ( - self.call_method( - location=expr.location, - call_expr=expr, - obj=(obj, obj_type), - method_name=method, - positional=positional, - keywords=keywords, - ) - or UnknownType() + return self.call_method( + location=expr.location, + call_expr=expr, + obj=(obj, obj_type), + method_name=method, + positional=positional, + keywords=keywords, ) callee: Type = self.type_of(expr.callee) diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json index 37e4aa9..83540db 100644 --- a/tests/cases/checker/09_frame_ops.py.ref.json +++ b/tests/cases/checker/09_frame_ops.py.ref.json @@ -2150,18 +2150,14 @@ "index": 0, "name": "a", "type": { - "type": { - "name": "int" - } + "type": {} } }, { "index": 1, "name": "b", "type": { - "type": { - "name": "float" - } + "type": {} } } ] @@ -2300,18 +2296,14 @@ "index": 0, "name": "a", "type": { - "type": { - "name": "int" - } + "type": {} } }, { "index": 1, "name": "b", "type": { - "type": { - "name": "float" - } + "type": {} } } ] @@ -2525,18 +2517,14 @@ "index": 0, "name": "a", "type": { - "type": { - "name": "int" - } + "type": {} } }, { "index": 1, "name": "b", "type": { - "type": { - "name": "float" - } + "type": {} } } ] @@ -2600,18 +2588,14 @@ "index": 0, "name": "a", "type": { - "type": { - "name": "int" - } + "type": {} } }, { "index": 1, "name": "b", "type": { - "type": { - "name": "float" - } + "type": {} } } ] @@ -2675,18 +2659,14 @@ "index": 0, "name": "a", "type": { - "type": { - "name": "int" - } + "type": {} } }, { "index": 1, "name": "b", "type": { - "type": { - "name": "float" - } + "type": {} } } ] @@ -2750,18 +2730,14 @@ "index": 0, "name": "a", "type": { - "type": { - "name": "int" - } + "type": {} } }, { "index": 1, "name": "b", "type": { - "type": { - "name": "float" - } + "type": {} } } ] From 4ea15519f38f2b6df557573d027a5c7962d9f8b5 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 12:07:36 +0200 Subject: [PATCH 34/36] feate(checker): add some frame/column attributes --- midas/checker/frames/column_manager.py | 18 +- midas/checker/frames/frame_manager.py | 13 + midas/checker/preamble.py | 4 +- midas/checker/python.py | 8 + midas/checker/registry.py | 26 ++ tests/cases/checker/09_frame_ops.py | 9 + tests/cases/checker/09_frame_ops.py.ref.json | 288 +++++++++++++++++++ 7 files changed, 363 insertions(+), 3 deletions(-) diff --git a/midas/checker/frames/column_manager.py b/midas/checker/frames/column_manager.py index 0793034..f62d025 100644 --- a/midas/checker/frames/column_manager.py +++ b/midas/checker/frames/column_manager.py @@ -1,12 +1,13 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import midas.ast.python as p from midas.ast.location import Location from midas.checker.frames.column_groupby_methods import Call as GroupByCall from midas.checker.frames.column_groupby_methods import ColumnGroupByMethodRegistry from midas.checker.frames.column_methods import Call, ColumnMethodRegistry +from midas.checker.registry import TypesRegistry from midas.checker.types import ColumnGroupBy, ColumnType, Type if TYPE_CHECKING: @@ -60,3 +61,18 @@ class ColumnManager: keywords=keywords, ) return self.groupby_method_resolver.call(method, call) + + def get_attribute(self, column: ColumnType, name: str) -> Optional[Type]: + types: TypesRegistry = self.typer.types + match name: + case "ndim" | "size": + return types.get_type("int") + + case "shape": + return types.tuple_of("int") + + case "T": + return column + + case _: + return None diff --git a/midas/checker/frames/frame_manager.py b/midas/checker/frames/frame_manager.py index 32bc782..8b99f14 100644 --- a/midas/checker/frames/frame_manager.py +++ b/midas/checker/frames/frame_manager.py @@ -7,6 +7,7 @@ from midas.ast.location import Location from midas.checker.frames.frame_groupby_methods import Call as GroupByCall from midas.checker.frames.frame_groupby_methods import FrameGroupByMethodRegistry from midas.checker.frames.frame_methods import Call, FrameMethodRegistry +from midas.checker.registry import TypesRegistry from midas.checker.reporter import FileReporter from midas.checker.types import ( ColumnGroupBy, @@ -240,3 +241,15 @@ class FrameManager: keywords=keywords, ) return self.groupby_method_resolver.call(method, call) + + def get_attribute(self, frame: DataFrameType, name: str) -> Optional[Type]: + types: TypesRegistry = self.typer.types + match name: + case "ndim" | "size": + return types.get_type("int") + + case "shape": + return types.tuple_of("int", "int") + + case _: + return None diff --git a/midas/checker/preamble.py b/midas/checker/preamble.py index b35d4dc..a60dabb 100644 --- a/midas/checker/preamble.py +++ b/midas/checker/preamble.py @@ -108,8 +108,8 @@ class Preamble(Environment): ], ) - def _list_of(self, item_type: Type) -> Type: - return self._types.apply_generic(self._types.get_type("list"), [item_type]) + def _list_of(self, item_type: str | Type) -> Type: + return self._types.list_of(item_type) def _def_type_constructor( self, name: str, py_function: Optional[Callable[..., Any]] = None diff --git a/midas/checker/python.py b/midas/checker/python.py index 99766d8..9103cf4 100644 --- a/midas/checker/python.py +++ b/midas/checker/python.py @@ -659,6 +659,14 @@ class PythonTyper( def visit_get_expr(self, expr: p.GetExpr) -> Type: object: Type = self.type_of(expr.object) member: Optional[Type] = self.types.lookup_member(object, expr.name) + + if member is None: + match object: + case DataFrameType(): + member = self.frame_mgr.get_attribute(object, expr.name) + case ColumnType(): + member = self.column_mgr.get_attribute(object, expr.name) + if member is None: self.reporter.warning( expr.location, f"Unknown member '{expr.name}' of {object}" diff --git a/midas/checker/registry.py b/midas/checker/registry.py index d1074a7..92b54a6 100644 --- a/midas/checker/registry.py +++ b/midas/checker/registry.py @@ -452,3 +452,29 @@ class TypesRegistry: def lookup_predicate(self, name: str) -> Optional[Predicate]: return self._predicates.get(name) + + def _by_name_or_type(self, name_or_type: str | Type) -> Type: + if isinstance(name_or_type, str): + return self.get_type(name_or_type) + return name_or_type + + def list_of(self, item_type: str | Type) -> Type: + list_ = self.get_type("list") + return self.apply_generic(list_, [self._by_name_or_type(item_type)]) + + def tuple_of(self, *item_types: str | Type) -> Type: + tuple_ = self.get_type("tuple") + return self.apply_generic( + tuple_, + [self._by_name_or_type(item_type) for item_type in item_types], + ) + + def dict_of(self, key_type: str | Type, value_type: str | Type) -> Type: + dict_ = self.get_type("dict") + return self.apply_generic( + dict_, + [ + self._by_name_or_type(key_type), + self._by_name_or_type(value_type), + ], + ) diff --git a/tests/cases/checker/09_frame_ops.py b/tests/cases/checker/09_frame_ops.py index 18951f5..6183698 100644 --- a/tests/cases/checker/09_frame_ops.py +++ b/tests/cases/checker/09_frame_ops.py @@ -99,3 +99,12 @@ _ = col_gb.prod() _ = col_gb.std() _ = col_gb.sum() _ = col_gb.var() + +# Attributes +_ = df1.ndim # int +_ = df1.size # int +_ = df1.shape # (int, int) +_ = col1.ndim # int +_ = col1.size # int +_ = col1.shape # (int) +_ = col1.T # Column[int] diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json index 83540db..79d7266 100644 --- a/tests/cases/checker/09_frame_ops.py.ref.json +++ b/tests/cases/checker/09_frame_ops.py.ref.json @@ -4407,6 +4407,294 @@ "type": { "type": {} } + }, + { + "location": { + "from": "L104:4", + "to": "L104:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L104:4", + "to": "L104:12" + }, + "expr": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "ndim" + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L105:4", + "to": "L105:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L105:4", + "to": "L105:12" + }, + "expr": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "size" + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L106:4", + "to": "L106:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L106:4", + "to": "L106:13" + }, + "expr": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "shape" + }, + "type": { + "items": [ + { + "name": "int" + }, + { + "name": "int" + } + ] + } + }, + { + "location": { + "from": "L107:4", + "to": "L107:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L107:4", + "to": "L107:13" + }, + "expr": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "ndim" + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L108:4", + "to": "L108:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L108:4", + "to": "L108:13" + }, + "expr": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "size" + }, + "type": { + "name": "int" + } + }, + { + "location": { + "from": "L109:4", + "to": "L109:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L109:4", + "to": "L109:14" + }, + "expr": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "shape" + }, + "type": { + "items": [ + { + "name": "int" + } + ] + } + }, + { + "location": { + "from": "L110:4", + "to": "L110:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L110:4", + "to": "L110:10" + }, + "expr": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "T" + }, + "type": { + "type": { + "name": "int" + } + } } ] } \ No newline at end of file From ec80b1e92ee5c5e69a8882a770faaae9066c7502 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 12:13:30 +0200 Subject: [PATCH 35/36] feat(checker): add head/tail methods --- midas/checker/frames/column_methods.py | 44 ++++ midas/checker/frames/frame_methods.py | 44 ++++ tests/cases/checker/09_frame_ops.py | 7 + tests/cases/checker/09_frame_ops.py.ref.json | 224 +++++++++++++++++++ 4 files changed, 319 insertions(+) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index d2559b8..02de5a3 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -248,6 +248,50 @@ class ColumnMethodRegistry(MethodRegistry[Call]): ], ) + @method() + def head(self, call: Call) -> Type: + signature = Function( + args=[ + Function.Argument( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + returns=call.column, + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + + @method() + def tail(self, call: Call) -> Type: + signature = Function( + args=[ + Function.Argument( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + returns=call.column, + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + @method() def groupby(self, call: Call) -> Type: bool_: Type = self.types.get_type("bool") diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 2bd8c8b..bd9e3ee 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -327,6 +327,50 @@ class FrameMethodRegistry(MethodRegistry[Call]): ], ) + @method() + def head(self, call: Call) -> Type: + signature = Function( + args=[ + Function.Argument( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + returns=call.frame, + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + + @method() + def tail(self, call: Call) -> Type: + signature = Function( + args=[ + Function.Argument( + pos=0, + name="n", + type=self.types.get_type("int"), + required=False, + ), + ], + returns=call.frame, + ) + + result: CallResult = self.dispatcher.get_result( + location=call.location, + callee=signature, + positional=call.positional, + keywords=call.keywords, + ) + return result.result + @method() def groupby(self, call: Call) -> Type: bool_: Type = self.types.get_type("bool") diff --git a/tests/cases/checker/09_frame_ops.py b/tests/cases/checker/09_frame_ops.py index 6183698..6d478f6 100644 --- a/tests/cases/checker/09_frame_ops.py +++ b/tests/cases/checker/09_frame_ops.py @@ -108,3 +108,10 @@ _ = col1.ndim # int _ = col1.size # int _ = col1.shape # (int) _ = col1.T # Column[int] + + +# Misc +_ = df1.head() +_ = df1.tail() +_ = col1.head() +_ = col1.tail() diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json index 79d7266..6c77757 100644 --- a/tests/cases/checker/09_frame_ops.py.ref.json +++ b/tests/cases/checker/09_frame_ops.py.ref.json @@ -4695,6 +4695,230 @@ "name": "int" } } + }, + { + "location": { + "from": "L114:4", + "to": "L114:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L114:4", + "to": "L114:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "head" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L115:4", + "to": "L115:7" + }, + "expr": { + "_type": "VariableExpr", + "name": "df1" + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L115:4", + "to": "L115:14" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "df1" + }, + "name": "tail" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "columns": [ + { + "index": 0, + "name": "a", + "type": { + "type": { + "name": "int" + } + } + }, + { + "index": 1, + "name": "b", + "type": { + "type": { + "name": "float" + } + } + } + ] + } + }, + { + "location": { + "from": "L116:4", + "to": "L116:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L116:4", + "to": "L116:15" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "head" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L117:4", + "to": "L117:8" + }, + "expr": { + "_type": "VariableExpr", + "name": "col1" + }, + "type": { + "type": { + "name": "int" + } + } + }, + { + "location": { + "from": "L117:4", + "to": "L117:15" + }, + "expr": { + "_type": "CallExpr", + "callee": { + "_type": "GetExpr", + "object": { + "_type": "VariableExpr", + "name": "col1" + }, + "name": "tail" + }, + "arguments": [], + "keywords": {} + }, + "type": { + "type": { + "name": "int" + } + } } ] } \ No newline at end of file From 4d3e3f44a1dbcf7671d93b9062f7c968d29a3588 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Fri, 3 Jul 2026 12:28:39 +0200 Subject: [PATCH 36/36] fix(checker): correctly check length of frame/column --- midas/checker/frames/column_methods.py | 25 ++++++++++++++++-------- midas/checker/frames/frame_methods.py | 27 ++++++++++++++++---------- 2 files changed, 34 insertions(+), 18 deletions(-) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index 02de5a3..95bce9e 100644 --- a/midas/checker/frames/column_methods.py +++ b/midas/checker/frames/column_methods.py @@ -355,6 +355,21 @@ class ColumnMethodRegistry(MethodRegistry[Call]): def _assert_same_length(self, call_expr: p.Expr, column1: p.Expr, column2: p.Expr): func_name: str = "__midas_column_same_length__" + + # Efficiently compute length + # https://stackoverflow.com/a/15943975/11109181 + def len_of_col(col: ast.expr) -> ast.expr: + return ast.Call( + func=ast.Name(id="len"), + args=[ + ast.Attribute( + value=col, + attr="index", + ) + ], + keywords=[], + ) + self.assertions.define( func_name, ast.FunctionDef( @@ -372,16 +387,10 @@ class ColumnMethodRegistry(MethodRegistry[Call]): body=[ ast.Return( value=ast.Compare( - left=ast.Attribute( - value=ast.Name(id="column1"), - attr="size", - ), + left=len_of_col(ast.Name(id="column1")), ops=[ast.Eq()], comparators=[ - ast.Attribute( - value=ast.Name(id="column2"), - attr="size", - ) + len_of_col(ast.Name(id="column2")), ], ) ) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index bd9e3ee..1be481f 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -434,6 +434,21 @@ class FrameMethodRegistry(MethodRegistry[Call]): def _assert_same_length(self, call_expr: p.Expr, frame1: p.Expr, frame2: p.Expr): func_name: str = "__midas_frame_same_length__" + + # Efficiently compute length + # https://stackoverflow.com/a/15943975/11109181 + def len_of_df(df: ast.expr) -> ast.expr: + return ast.Call( + func=ast.Name(id="len"), + args=[ + ast.Attribute( + value=df, + attr="index", + ) + ], + keywords=[], + ) + self.assertions.define( func_name, ast.FunctionDef( @@ -451,17 +466,9 @@ class FrameMethodRegistry(MethodRegistry[Call]): body=[ ast.Return( value=ast.Compare( - left=ast.Attribute( - value=ast.Name(id="frame1"), - attr="size", - ), + left=len_of_df(ast.Name(id="frame1")), ops=[ast.Eq()], - comparators=[ - ast.Attribute( - value=ast.Name(id="frame2"), - attr="size", - ) - ], + comparators=[len_of_df(ast.Name(id="frame2"))], ) ) ],