From e37432ad571ecf3ed2bbc32607367d6ec3aeba0f Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Fri, 24 Jul 2026 00:03:26 +0200
Subject: [PATCH] fix(checker): correct return type for column aggregation

---
 .../checker/frames/column_groupby_methods.py  |  5 +-
 midas/checker/frames/column_methods.py        | 10 ++--
 tests/cases/checker/09_frame_ops.py.ref.json  | 48 +++++--------------
 3 files changed, 17 insertions(+), 46 deletions(-)

diff --git a/midas/checker/frames/column_groupby_methods.py b/midas/checker/frames/column_groupby_methods.py
index c92b08e..2cd1c8e 100644
--- a/midas/checker/frames/column_groupby_methods.py
+++ b/midas/checker/frames/column_groupby_methods.py
@@ -13,7 +13,6 @@ from midas.checker.types import (
     Function,
     ParamSpec,
     Type,
-    UnknownType,
 )
 
 if TYPE_CHECKING:
@@ -97,11 +96,9 @@ class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
             positional=[],
             keywords={},
         )
-        if not isinstance(returns, ColumnType):
-            returns = ColumnType(type=UnknownType())
         signature = Function(
             params=ParamSpec(mixed=real_params),
-            returns=returns,
+            returns=ColumnType(type=returns),
         )
 
         result: CallResult = self.dispatcher.get_result(
diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py
index e957a55..3aa0136 100644
--- a/midas/checker/frames/column_methods.py
+++ b/midas/checker/frames/column_methods.py
@@ -364,13 +364,11 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
             Type: the result type
         """
 
-        returns: Type = ColumnType(type=TopType())
+        returns: Type = TopType()
         if formula:
-            returns = ColumnType(
-                type=self._resolve_formula_type(
-                    call,
-                    formula(call.column.type),
-                )
+            returns = self._resolve_formula_type(
+                call,
+                formula(call.column.type),
             )
         signature = Function(
             params=ParamSpec(
diff --git a/tests/cases/checker/09_frame_ops.py.ref.json b/tests/cases/checker/09_frame_ops.py.ref.json
index 2e6c448..c48c697 100644
--- a/tests/cases/checker/09_frame_ops.py.ref.json
+++ b/tests/cases/checker/09_frame_ops.py.ref.json
@@ -4305,9 +4305,7 @@
         "arguments": [],
         "keywords": {}
       },
-      "type": {
-        "type": {}
-      }
+      "type": {}
     },
     {
       "location": {
@@ -4342,9 +4340,7 @@
         "arguments": [],
         "keywords": {}
       },
-      "type": {
-        "type": {}
-      }
+      "type": {}
     },
     {
       "location": {
@@ -4380,9 +4376,7 @@
         "keywords": {}
       },
       "type": {
-        "type": {
-          "name": "int"
-        }
+        "name": "int"
       }
     },
     {
@@ -4419,9 +4413,7 @@
         "keywords": {}
       },
       "type": {
-        "type": {
-          "name": "float"
-        }
+        "name": "float"
       }
     },
     {
@@ -4458,9 +4450,7 @@
         "keywords": {}
       },
       "type": {
-        "type": {
-          "name": "int"
-        }
+        "name": "int"
       }
     },
     {
@@ -4497,9 +4487,7 @@
         "keywords": {}
       },
       "type": {
-        "type": {
-          "name": "int"
-        }
+        "name": "int"
       }
     },
     {
@@ -4536,9 +4524,7 @@
         "keywords": {}
       },
       "type": {
-        "type": {
-          "name": "int"
-        }
+        "name": "int"
       }
     },
     {
@@ -4575,9 +4561,7 @@
         "keywords": {}
       },
       "type": {
-        "type": {
-          "name": "int"
-        }
+        "name": "int"
       }
     },
     {
@@ -4614,9 +4598,7 @@
         "keywords": {}
       },
       "type": {
-        "type": {
-          "name": "int"
-        }
+        "name": "int"
       }
     },
     {
@@ -4652,9 +4634,7 @@
         "arguments": [],
         "keywords": {}
       },
-      "type": {
-        "type": {}
-      }
+      "type": {}
     },
     {
       "location": {
@@ -4690,9 +4670,7 @@
         "keywords": {}
       },
       "type": {
-        "type": {
-          "name": "int"
-        }
+        "name": "int"
       }
     },
     {
@@ -4728,9 +4706,7 @@
         "arguments": [],
         "keywords": {}
       },
-      "type": {
-        "type": {}
-      }
+      "type": {}
     },
     {
       "location": {