From 25c11c3a5362c2d78d8661e5ef844128422cccd0 Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Tue, 7 Jul 2026 17:07:20 +0200
Subject: [PATCH] fix(checker): compute aggregation of column groupby

---
 .../checker/frames/column_groupby_methods.py  | 35 +++++++++++++------
 midas/checker/frames/column_methods.py        |  7 +++-
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/midas/checker/frames/column_groupby_methods.py b/midas/checker/frames/column_groupby_methods.py
index f84d8df..c74bed4 100644
--- a/midas/checker/frames/column_groupby_methods.py
+++ b/midas/checker/frames/column_groupby_methods.py
@@ -12,8 +12,8 @@ from midas.checker.types import (
     ColumnType,
     Function,
     ParamSpec,
-    TopType,
     Type,
+    UnknownType,
 )
 
 if TYPE_CHECKING:
@@ -49,9 +49,8 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def _aggregate(
         self,
         call: Call,
+        method: str,
         params: list[str | tuple[str, str, bool]] = [],
-        *,
-        preserve_inner_type: bool = False,
     ) -> Type:
         """Compute the result type of an aggregate method call
 
@@ -87,13 +86,21 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
                     )
             real_params.append(param)
 
+        # TODO: maybe better to filter arguments and pass some, in case the
+        # return type depends on them
+        returns: Type = self.typer.call_method(
+            location=call.location,
+            call_expr=call.call_expr,
+            obj=(call.groupby_expr, call.groupby.column),
+            method_name=method,
+            positional=[],
+            keywords={},
+        )
+        if not isinstance(returns, ColumnType):
+            returns = ColumnType(type=UnknownType())
         signature = Function(
             params=ParamSpec(mixed=real_params),
-            returns=(
-                call.groupby.column
-                if preserve_inner_type
-                else ColumnType(type=TopType())
-            ),
+            returns=returns,
         )
 
         result: CallResult = self.dispatcher.get_result(
@@ -108,6 +115,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def kurt(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "kurt",
             ["skipna", "numeric_only"],
         )
 
@@ -115,6 +123,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def max(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "max",
             [
                 "numeric_only",
                 (
@@ -126,13 +135,13 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
                 "engine",
                 "engine_kwargs",
             ],
-            preserve_inner_type=True,
         )
 
     @method()
     def mean(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "mean",
             ["numeric_only", "skipna", "engine", "engine_kwargs"],
         )
 
@@ -140,14 +149,15 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def median(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "median",
             ["numeric_only", "skipna"],
-            preserve_inner_type=True,
         )
 
     @method()
     def min(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "min",
             [
                 "numeric_only",
                 (
@@ -159,13 +169,13 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
                 "engine",
                 "engine_kwargs",
             ],
-            preserve_inner_type=True,
         )
 
     @method()
     def prod(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "prod",
             [
                 "numeric_only",
                 (
@@ -181,6 +191,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def std(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "std",
             [
                 (
                     "ddof",
@@ -198,6 +209,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def sum(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "sum",
             [
                 "numeric_only",
                 (
@@ -215,6 +227,7 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
     def var(self, call: Call) -> Type:
         return self._aggregate(
             call,
+            "var",
             [
                 (
                     "var",
diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py
index 4cbb776..4a31ed1 100644
--- a/midas/checker/frames/column_methods.py
+++ b/midas/checker/frames/column_methods.py
@@ -318,7 +318,12 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
 
         returns: Type = ColumnType(type=TopType())
         if formula:
-            returns = self._resolve_formula_type(call, formula(call.column.type))
+            returns = ColumnType(
+                type=self._resolve_formula_type(
+                    call,
+                    formula(call.column.type),
+                )
+            )
         signature = Function(
             params=ParamSpec(
                 kw=[