From 3c92f0867d118786bc21726e6d75d6a4cf73a656 Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Thu, 2 Jul 2026 18:00:25 +0200
Subject: [PATCH] feat(types): add ColumnGroupBy

---
 midas/checker/frames.py      | 30 +++++++++++++++++++++++++++++-
 midas/checker/python.py      | 10 ++++++++++
 midas/checker/types.py       | 17 +++++++++++++++++
 midas/generator/generator.py |  2 ++
 midas/generator/stubs.py     | 14 ++++++++++++++
 5 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/midas/checker/frames.py b/midas/checker/frames.py
index 1a1a917..2308317 100644
--- a/midas/checker/frames.py
+++ b/midas/checker/frames.py
@@ -6,7 +6,15 @@ import midas.ast.python as p
 from midas.ast.location import Location
 from midas.checker.frame_methods import Call, MethodRegistry
 from midas.checker.reporter import FileReporter
-from midas.checker.types import ColumnType, DataFrameType, TupleType, Type, UnknownType
+from midas.checker.types import (
+    ColumnGroupBy,
+    ColumnType,
+    DataFrameType,
+    FrameGroupBy,
+    TupleType,
+    Type,
+    UnknownType,
+)
 
 if TYPE_CHECKING:
     from midas.checker.python import PythonTyper, TypedExpr
@@ -90,6 +98,26 @@ class FrameManager:
                 reporter.error(location, f"Invalid index type {index} on {frame}")
                 return UnknownType()
 
+    def groupby_get(
+        self,
+        reporter: FileReporter,
+        location: Location,
+        groupby: FrameGroupBy,
+        index: p.Expr,
+    ) -> Type:
+        result: Type = self.get(reporter, location, groupby.frame, index)
+        match result:
+            case ColumnType():
+                result = ColumnGroupBy(column=result)
+            case TupleType(items=columns):
+                result = TupleType(
+                    items=tuple(
+                        ColumnGroupBy(column=cast(ColumnType, column))
+                        for column in columns
+                    )
+                )
+        return result
+
     @classmethod
     def _set_column(
         cls, frame: DataFrameType, name: str, column: ColumnType
diff --git a/midas/checker/python.py b/midas/checker/python.py
index 7ab803d..6e2f30c 100644
--- a/midas/checker/python.py
+++ b/midas/checker/python.py
@@ -26,6 +26,7 @@ from midas.checker.types import (
     ConstraintType,
     DataFrameType,
     DerivedType,
+    FrameGroupBy,
     Function,
     GenericType,
     TopType,
@@ -770,6 +771,8 @@ class PythonTyper(
                 return self._visit_tuple_subscript(unfolded, expr)
             case DataFrameType():
                 return self._visit_frame_subscript(unfolded, expr)
+            case FrameGroupBy():
+                return self._visit_frame_groupby_subscript(unfolded, expr)
 
         operation: Optional[Type] = self.types.lookup_member(object, "__getitem__")
         if operation is None:
@@ -1095,3 +1098,10 @@ class PythonTyper(
         self, frame: DataFrameType, expr: p.SubscriptExpr
     ) -> Type:
         return self.frame_mgr.get(self.reporter, expr.location, frame, expr.index)
+
+    def _visit_frame_groupby_subscript(
+        self, groupby: FrameGroupBy, expr: p.SubscriptExpr
+    ) -> Type:
+        return self.frame_mgr.groupby_get(
+            self.reporter, expr.location, groupby, expr.index
+        )
diff --git a/midas/checker/types.py b/midas/checker/types.py
index 299b392..51658d7 100644
--- a/midas/checker/types.py
+++ b/midas/checker/types.py
@@ -195,6 +195,14 @@ class FrameGroupBy:
         return f"FrameGroupBy[{self.frame}]"
 
 
+@dataclass(frozen=True, kw_only=True)
+class ColumnGroupBy:
+    column: ColumnType
+
+    def __str__(self) -> str:
+        return f"ColumnGroupBy[{self.column}]"
+
+
 def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
     def sub_argument(arg: Function.Argument):
         return Function.Argument(
@@ -318,6 +326,11 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
                 frame=cast(DataFrameType, substitute_typevars(frame, substitutions))
             )
 
+        case ColumnGroupBy(column=column):
+            return ColumnGroupBy(
+                column=cast(ColumnType, substitute_typevars(column, substitutions))
+            )
+
         case UnknownType() | UnitType():
             return type
 
@@ -397,6 +410,9 @@ def to_annotation(type: Type) -> str:
         case FrameGroupBy():
             return "pd.api.typing.DataFrameGroupBy"
 
+        case ColumnGroupBy():
+            return "pd.api.typing.SeriesGroupBy"
+
         case _:
             assert_never(type)
 
@@ -426,4 +442,5 @@ Type = (
     | ColumnType
     | DataFrameType
     | FrameGroupBy
+    | ColumnGroupBy
 )
diff --git a/midas/generator/generator.py b/midas/generator/generator.py
index dc6af9d..352a544 100644
--- a/midas/generator/generator.py
+++ b/midas/generator/generator.py
@@ -14,6 +14,7 @@ from midas.checker.registry import TypesRegistry
 from midas.checker.types import (
     AppliedType,
     BaseType,
+    ColumnGroupBy,
     ColumnType,
     ComplexType,
     ConstraintType,
@@ -504,6 +505,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 | ExtensionType()
                 | GenericType()
                 | FrameGroupBy()
+                | ColumnGroupBy()
             ):
                 self.logger.warning(f"Can't make assertion for type {type}")
                 return []
diff --git a/midas/generator/stubs.py b/midas/generator/stubs.py
index e86d3f6..07ab0db 100644
--- a/midas/generator/stubs.py
+++ b/midas/generator/stubs.py
@@ -6,6 +6,7 @@ from midas.checker.registry import Member, TypesRegistry
 from midas.checker.types import (
     AppliedType,
     BaseType,
+    ColumnGroupBy,
     ColumnType,
     ComplexType,
     ConstraintType,
@@ -301,6 +302,19 @@ class StubsGenerator:
                     attr="DataFrameGroupBy",
                 )
 
+            case ColumnGroupBy():
+                self.import_pandas = True
+                return ast.Attribute(
+                    value=ast.Attribute(
+                        value=ast.Attribute(
+                            value=ast.Name(id="pd"),
+                            attr="api",
+                        ),
+                        attr="typing",
+                    ),
+                    attr="SeriesGroupBy",
+                )
+
             case _:
                 assert_never(type)