feat(checker): add statistical ops on frames

This commit is contained in:
2026-07-03 01:27:16 +02:00
parent 1f6189daa4
commit 1c0c917873
3 changed files with 1176 additions and 450 deletions

View File

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

View File

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

File diff suppressed because it is too large Load Diff