feat(checker): add mean method on frames

This commit is contained in:
2026-06-26 11:21:38 +02:00
parent 43e40396a1
commit 5b0c5c01ad

View File

@@ -10,6 +10,8 @@ from midas.checker.types import (
ColumnType,
DataFrameType,
Function,
OverloadedFunction,
TopType,
Type,
UnknownType,
unfold_type,
@@ -85,6 +87,9 @@ class MethodRegistry(metaclass=_MethodRegistryMeta):
self,
call: Call,
) -> Type:
# TODO: support add with scalar, sequence, Series, dict
# TODO: check operation exists on inner column types
new_columns: list[DataFrameType.Column] = []
by_name: dict[str, DataFrameType.Column] = {}
@@ -151,3 +156,43 @@ class MethodRegistry(metaclass=_MethodRegistryMeta):
)
or UnknownType()
)
@frame_method()
def mean(self, call: Call) -> Type:
with_axis = Function(
kw_args=[
Function.Argument(
pos=0,
name="axis",
type=self.types.get_type("int"),
required=False,
)
],
returns=ColumnType(type=TopType()),
)
without_axis = Function(
kw_args=[
Function.Argument(
pos=0,
name="axis",
type=self.types.get_type("None"),
required=True,
)
],
returns=TopType(),
)
overload = OverloadedFunction(
overloads=[
with_axis,
without_axis,
]
)
return (
self.typer._get_call_result(
location=call.location,
callee=overload,
positional=call.positional,
keywords=call.keywords,
)
or UnknownType()
)