HEL/midas
1
0
Fork 0
Code Issues Pull Requests Actions Releases 2 Activity

fix(checker): delegate frame aggregate methods to columns

Browse Source
This commit is contained in:
Louis Heredero
2026-07-03 11:42:35 +02:00
parent 733c8736b8
commit 7a6e01cff8
3 changed files with 61 additions and 186 deletions
Download Patch File Download Diff File Expand all files Collapse all files

170
midas/checker/frames/frame_groupby_methods.py
View File

@@ -5,9 +5,15 @@ from typing import TYPE_CHECKING
import midas.ast.python as p import midas.ast.python as p
from midas.ast.location import Location from midas.ast.location import Location
from midas.checker.dispatcher import CallResult
from midas.checker.frames.utils import MethodRegistry, method 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: if TYPE_CHECKING:
from midas.checker.python import TypedExpr from midas.checker.python import TypedExpr
@@ -35,161 +41,63 @@ class FrameGroupByMethodRegistry(MethodRegistry[Call]):
"engine_kwargs": "dict", "engine_kwargs": "dict",
} }
def _aggregate( def _aggregate(self, call: Call, method: str) -> Type:
self, call: Call, args: list[str | tuple[str, str, bool]] = [] new_columns: list[DataFrameType.Column] = []
) -> 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( for column in call.groupby.frame.columns:
args=real_args, column_groupby: ColumnGroupBy = ColumnGroupBy(column=column.type)
returns=call.groupby.frame, 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( return DataFrameType(columns=new_columns)
location=call.location,
callee=signature,
positional=call.positional,
keywords=call.keywords,
)
return result.result
@method() @method()
def kurt(self, call: Call) -> Type: def kurt(self, call: Call) -> Type:
return self._aggregate( return self._aggregate(call, "kurt")
call,
[
"skipna",
"numeric_only",
],
)
@method() @method()
def max(self, call: Call) -> Type: def max(self, call: Call) -> Type:
return self._aggregate( return self._aggregate(call, "max")
call,
[
"numeric_only",
(
"min_count",
"int",
False,
),
"skipna",
"engine",
"engine_kwargs",
],
)
@method() @method()
def mean(self, call: Call) -> Type: def mean(self, call: Call) -> Type:
return self._aggregate( return self._aggregate(call, "mean")
call,
["numeric_only", "skipna", "engine", "engine_kwargs"],
)
@method() @method()
def median(self, call: Call) -> Type: def median(self, call: Call) -> Type:
return self._aggregate( return self._aggregate(call, "median")
call,
["numeric_only", "skipna"],
)
@method() @method()
def min(self, call: Call) -> Type: def min(self, call: Call) -> Type:
return self._aggregate( return self._aggregate(call, "min")
call,
[
"numeric_only",
(
"min_count",
"int",
False,
),
"skipna",
"engine",
"engine_kwargs",
],
)
@method() @method()
def prod(self, call: Call) -> Type: def prod(self, call: Call) -> Type:
return self._aggregate( return self._aggregate(call, "prod")
call,
[
"numeric_only",
(
"min_count",
"int",
False,
),
"skipna",
],
)
@method() @method()
def std(self, call: Call) -> Type: def std(self, call: Call) -> Type:
return self._aggregate( return self._aggregate(call, "std")
call,
[
(
"ddof",
"int",
False,
),
"engine",
"engine_kwargs",
"numeric_only",
"skipna",
],
)
@method() @method()
def sum(self, call: Call) -> Type: def sum(self, call: Call) -> Type:
return self._aggregate( return self._aggregate(call, "sum")
call,
[
"numeric_only",
(
"min_count",
"int",
False,
),
"skipna",
"engine",
"engine_kwargs",
],
)
@method() @method()
def var(self, call: Call) -> Type: def var(self, call: Call) -> Type:
return self._aggregate( return self._aggregate(call, "var")
call,
[
(
"var",
"int",
False,
),
"engine",
"engine_kwargs",
"numeric_only",
"skipna",
],
)

29
midas/checker/python.py
View File

@@ -222,7 +222,7 @@ class PythonTyper(
method_name: str, method_name: str,
positional: list[TypedExpr], positional: list[TypedExpr],
keywords: dict[str, TypedExpr], keywords: dict[str, TypedExpr],
) -> Optional[Type]: ) -> Type:
unfolded: Type = unfold_type(obj[1]) unfolded: Type = unfold_type(obj[1])
match unfolded: match unfolded:
case DataFrameType(): case DataFrameType():
@@ -580,9 +580,8 @@ class PythonTyper(
right: TypedExpr, right: TypedExpr,
method: str, method: str,
) -> Type: ) -> Type:
result: Optional[Type]
try: try:
result = self.call_method( return self.call_method(
location=location, location=location,
call_expr=expr, call_expr=expr,
obj=left, obj=left,
@@ -597,8 +596,6 @@ class PythonTyper(
) )
return UnknownType() return UnknownType()
return result or UnknownType()
def visit_unary_expr(self, expr: p.UnaryExpr) -> Type: def visit_unary_expr(self, expr: p.UnaryExpr) -> Type:
method: Optional[str] = PY_UNARY_METHODS.get(expr.operator.__class__) method: Optional[str] = PY_UNARY_METHODS.get(expr.operator.__class__)
if method is None: if method is None:
@@ -610,9 +607,8 @@ class PythonTyper(
operand: Type = self.type_of(expr.right) operand: Type = self.type_of(expr.right)
result: Optional[Type]
try: try:
result = self.call_method( return self.call_method(
location=expr.location, location=expr.location,
call_expr=expr, call_expr=expr,
obj=(expr.right, operand), obj=(expr.right, operand),
@@ -627,8 +623,6 @@ class PythonTyper(
) )
return UnknownType() return UnknownType()
return result or UnknownType()
def visit_call_expr(self, expr: p.CallExpr) -> Type: def visit_call_expr(self, expr: p.CallExpr) -> Type:
match expr.callee: match expr.callee:
case p.VariableExpr(name="TypeVar"): case p.VariableExpr(name="TypeVar"):
@@ -644,16 +638,13 @@ class PythonTyper(
match expr.callee: match expr.callee:
case p.GetExpr(object=obj, name=method): case p.GetExpr(object=obj, name=method):
obj_type: Type = self.type_of(obj) obj_type: Type = self.type_of(obj)
return ( return self.call_method(
self.call_method( location=expr.location,
location=expr.location, call_expr=expr,
call_expr=expr, obj=(obj, obj_type),
obj=(obj, obj_type), method_name=method,
method_name=method, positional=positional,
positional=positional, keywords=keywords,
keywords=keywords,
)
or UnknownType()
) )
callee: Type = self.type_of(expr.callee) callee: Type = self.type_of(expr.callee)

48
tests/cases/checker/09_frame_ops.py.ref.json
View File

@@ -2150,18 +2150,14 @@
"index": 0, "index": 0,
"name": "a", "name": "a",
"type": { "type": {
"type": { "type": {}
"name": "int"
}
} }
}, },
{ {
"index": 1, "index": 1,
"name": "b", "name": "b",
"type": { "type": {
"type": { "type": {}
"name": "float"
}
} }
} }
] ]
@@ -2300,18 +2296,14 @@
"index": 0, "index": 0,
"name": "a", "name": "a",
"type": { "type": {
"type": { "type": {}
"name": "int"
}
} }
}, },
{ {
"index": 1, "index": 1,
"name": "b", "name": "b",
"type": { "type": {
"type": { "type": {}
"name": "float"
}
} }
} }
] ]
@@ -2525,18 +2517,14 @@
"index": 0, "index": 0,
"name": "a", "name": "a",
"type": { "type": {
"type": { "type": {}
"name": "int"
}
} }
}, },
{ {
"index": 1, "index": 1,
"name": "b", "name": "b",
"type": { "type": {
"type": { "type": {}
"name": "float"
}
} }
} }
] ]
@@ -2600,18 +2588,14 @@
"index": 0, "index": 0,
"name": "a", "name": "a",
"type": { "type": {
"type": { "type": {}
"name": "int"
}
} }
}, },
{ {
"index": 1, "index": 1,
"name": "b", "name": "b",
"type": { "type": {
"type": { "type": {}
"name": "float"
}
} }
} }
] ]
@@ -2675,18 +2659,14 @@
"index": 0, "index": 0,
"name": "a", "name": "a",
"type": { "type": {
"type": { "type": {}
"name": "int"
}
} }
}, },
{ {
"index": 1, "index": 1,
"name": "b", "name": "b",
"type": { "type": {
"type": { "type": {}
"name": "float"
}
} }
} }
] ]
@@ -2750,18 +2730,14 @@
"index": 0, "index": 0,
"name": "a", "name": "a",
"type": { "type": {
"type": { "type": {}
"name": "int"
}
} }
}, },
{ {
"index": 1, "index": 1,
"name": "b", "name": "b",
"type": { "type": {
"type": { "type": {}
"name": "float"
}
} }
} }
] ]