feat(checker): add mean method on frames
This commit is contained in:
@@ -10,6 +10,8 @@ from midas.checker.types import (
|
|||||||
ColumnType,
|
ColumnType,
|
||||||
DataFrameType,
|
DataFrameType,
|
||||||
Function,
|
Function,
|
||||||
|
OverloadedFunction,
|
||||||
|
TopType,
|
||||||
Type,
|
Type,
|
||||||
UnknownType,
|
UnknownType,
|
||||||
unfold_type,
|
unfold_type,
|
||||||
@@ -85,6 +87,9 @@ class MethodRegistry(metaclass=_MethodRegistryMeta):
|
|||||||
self,
|
self,
|
||||||
call: Call,
|
call: Call,
|
||||||
) -> Type:
|
) -> Type:
|
||||||
|
# TODO: support add with scalar, sequence, Series, dict
|
||||||
|
# TODO: check operation exists on inner column types
|
||||||
|
|
||||||
new_columns: list[DataFrameType.Column] = []
|
new_columns: list[DataFrameType.Column] = []
|
||||||
|
|
||||||
by_name: dict[str, DataFrameType.Column] = {}
|
by_name: dict[str, DataFrameType.Column] = {}
|
||||||
@@ -151,3 +156,43 @@ class MethodRegistry(metaclass=_MethodRegistryMeta):
|
|||||||
)
|
)
|
||||||
or UnknownType()
|
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()
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user