feat(checker): add statistical ops on columns

This commit is contained in:
2026-07-03 01:34:58 +02:00
parent 1c0c917873
commit 0c70048b62

View File

@@ -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")