feat(checker): handle calls on group-bys

midas/checker/frames/column_groupby_methods.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.dispatcher import CallResult
+from midas.checker.frames.utils import MethodRegistry, method
+from midas.checker.types import ColumnGroupBy, Function, Type
+
+if TYPE_CHECKING:
+    from midas.checker.python import TypedExpr
+
+
+@dataclass(frozen=True, kw_only=True)
+class Call:
+    location: Location
+    call_expr: p.Expr
+    groupby: ColumnGroupBy
+    groupby_expr: p.Expr
+    positional: list[TypedExpr]
+    keywords: dict[str, TypedExpr]
+
+
+class ColumnGroupByMethodRegistry(MethodRegistry[Call]):
+    @method()
+    def mean(self, call: Call) -> Type:
+        bool_ = self.types.get_type("bool")
+        signature = Function(
+            args=[
+                Function.Argument(
+                    pos=0,
+                    name="numeric_only",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=1,
+                    name="skipna",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=2,
+                    name="engine",
+                    type=self.types.get_type("str"),
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=3,
+                    name="engine_kwargs",
+                    type=self.types.get_type("dict"),
+                    required=False,
+                ),
+            ],
+            returns=call.groupby.column,
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result

midas/checker/frames/column_manager.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.frames.column_groupby_methods import Call as GroupByCall
+from midas.checker.frames.column_groupby_methods import ColumnGroupByMethodRegistry
+from midas.checker.frames.column_methods import Call, ColumnMethodRegistry
+from midas.checker.types import ColumnGroupBy, ColumnType, Type
+
+if TYPE_CHECKING:
+    from midas.checker.python import PythonTyper, TypedExpr
+
+
+class ColumnManager:
+    def __init__(self, typer: PythonTyper) -> None:
+        self.typer: PythonTyper = typer
+        self.method_resolver: ColumnMethodRegistry = ColumnMethodRegistry(self.typer)
+        self.groupby_method_resolver: ColumnGroupByMethodRegistry = (
+            ColumnGroupByMethodRegistry(self.typer)
+        )
+
+    def call(
+        self,
+        method: str,
+        location: Location,
+        call_expr: p.Expr,
+        column: ColumnType,
+        column_expr: p.Expr,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+    ) -> Type:
+        call: Call = Call(
+            location=location,
+            call_expr=call_expr,
+            column=column,
+            column_expr=column_expr,
+            positional=positional,
+            keywords=keywords,
+        )
+        return self.method_resolver.call(method, call)
+
+    def groupby_call(
+        self,
+        method: str,
+        location: Location,
+        call_expr: p.Expr,
+        groupby: ColumnGroupBy,
+        groupby_expr: p.Expr,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+    ) -> Type:
+        call: GroupByCall = GroupByCall(
+            location=location,
+            call_expr=call_expr,
+            groupby=groupby,
+            groupby_expr=groupby_expr,
+            positional=positional,
+            keywords=keywords,
+        )
+        return self.groupby_method_resolver.call(method, call)

midas/checker/frames/column_methods.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.dispatcher import CallResult
+from midas.checker.frames.utils import MethodRegistry, method
+from midas.checker.types import (
+    ColumnGroupBy,
+    ColumnType,
+    Function,
+    GenericType,
+    TopType,
+    Type,
+    TypeVar,
+    UnknownType,
+    unfold_type,
+)
+
+if TYPE_CHECKING:
+    from midas.checker.python import TypedExpr
+
+
+@dataclass(frozen=True, kw_only=True)
+class Call:
+    location: Location
+    call_expr: p.Expr
+    column: ColumnType
+    column_expr: p.Expr
+    positional: list[TypedExpr]
+    keywords: dict[str, TypedExpr]
+
+
+class ColumnMethodRegistry(MethodRegistry[Call]):
+    @method("add", "__add__")
+    def add(self, call: Call) -> Type:
+        # TODO: support add with scalar
+        # TODO: check operation exists on inner column types
+
+        column2: Optional[ColumnType] = None
+
+        col_type1: Type = call.column.type
+        new_column: Type = ColumnType(type=UnknownType())
+        if len(call.positional) != 0:
+            other: Type = call.positional[0][1]
+            unfolded_other: Type = unfold_type(other)
+            if isinstance(unfolded_other, ColumnType):
+                column2 = unfolded_other
+                col_type2: Type = column2.type
+                if self.types.are_equivalent(col_type2, col_type1):
+                    new_column = ColumnType(type=col_type1)
+
+        # Build signature with new column type and generic operand
+        param_type: TypeVar = TypeVar(name="T", bound=None)
+        signature = GenericType(
+            name="add",
+            params=[param_type],
+            body=Function(
+                args=[
+                    Function.Argument(
+                        pos=0,
+                        name="other",
+                        type=ColumnType(type=param_type),
+                        required=True,
+                    ),
+                ],
+                returns=new_column,
+            ),
+        )
+
+        # Map arguments and compute result type
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        if result.is_valid:
+            self._assert_same_length(
+                call.call_expr, call.column_expr, call.positional[0][0]
+            )
+
+        return result.result
+
+    @method()
+    def mean(self, call: Call) -> Type:
+        signature = Function(
+            kw_args=[
+                Function.Argument(
+                    pos=0,
+                    name="axis",
+                    type=TopType(),
+                    required=False,
+                )
+            ],
+            returns=ColumnType(type=TopType()),
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    @method()
+    def groupby(self, call: Call) -> Type:
+        bool_: Type = self.types.get_type("bool")
+        function: Function = Function(
+            args=[
+                Function.Argument(
+                    pos=0,
+                    name="by",
+                    type=TopType(),
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=1,
+                    name="level",
+                    type=TopType(),
+                    required=False,
+                ),
+            ],
+            kw_args=[
+                Function.Argument(
+                    pos=2,
+                    name="as_index",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=3,
+                    name="sort",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=4,
+                    name="group_keys",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=5,
+                    name="observed",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=6,
+                    name="dropna",
+                    type=bool_,
+                    required=False,
+                ),
+            ],
+            returns=ColumnGroupBy(column=call.column),
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=function,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
+    def _assert_same_length(self, call_expr: p.Expr, column1: p.Expr, column2: p.Expr):
+        func_name: str = "__midas_column_same_length__"
+        self.assertions.define(
+            func_name,
+            ast.FunctionDef(
+                name=func_name,
+                args=ast.arguments(
+                    posonlyargs=[],
+                    args=[
+                        ast.arg(arg="column1"),
+                        ast.arg(arg="column2"),
+                    ],
+                    kwonlyargs=[],
+                    defaults=[],
+                    kw_defaults=[],
+                ),
+                body=[
+                    ast.Return(
+                        value=ast.Compare(
+                            left=ast.Attribute(
+                                value=ast.Name(id="column1"),
+                                attr="size",
+                            ),
+                            ops=[ast.Eq()],
+                            comparators=[
+                                ast.Attribute(
+                                    value=ast.Name(id="column2"),
+                                    attr="size",
+                                )
+                            ],
+                        )
+                    )
+                ],
+                decorator_list=[],
+            ),
+        )
+        self.assertions.add(
+            bound_expr=call_expr,
+            inputs=[column1, column2],
+            builder=lambda c1, c2: ast.Call(
+                func=ast.Name(id=func_name),
+                args=[c1, c2],
+                keywords=[],
+            ),
+            message="Columns must have the same length",
+        )

midas/checker/frames/frame_groupby_methods.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.dispatcher import CallResult
+from midas.checker.frames.utils import MethodRegistry, method
+from midas.checker.types import FrameGroupBy, Function, Type
+
+if TYPE_CHECKING:
+    from midas.checker.python import TypedExpr
+
+
+@dataclass(frozen=True, kw_only=True)
+class Call:
+    location: Location
+    call_expr: p.Expr
+    groupby: FrameGroupBy
+    groupby_expr: p.Expr
+    positional: list[TypedExpr]
+    keywords: dict[str, TypedExpr]
+
+
+class FrameGroupByMethodRegistry(MethodRegistry[Call]):
+    @method()
+    def mean(self, call: Call) -> Type:
+        bool_ = self.types.get_type("bool")
+        signature = Function(
+            args=[
+                Function.Argument(
+                    pos=0,
+                    name="numeric_only",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=1,
+                    name="skipna",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=2,
+                    name="engine",
+                    type=self.types.get_type("str"),
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=3,
+                    name="engine_kwargs",
+                    type=self.types.get_type("dict"),
+                    required=False,
+                ),
+            ],
+            returns=call.groupby.frame,
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=signature,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result

midas/checker/frames/frame_manager.py

@@ -4,9 +4,19 @@ from typing import TYPE_CHECKING, Optional, TypeGuard, cast
 
 import midas.ast.python as p
 from midas.ast.location import Location
-from midas.checker.frame_methods import Call, MethodRegistry
+from midas.checker.frames.frame_groupby_methods import Call as GroupByCall
+from midas.checker.frames.frame_groupby_methods import FrameGroupByMethodRegistry
+from midas.checker.frames.frame_methods import Call, FrameMethodRegistry
 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
@@ -19,7 +29,10 @@ def is_list_of_literals(exprs: list[p.Expr]) -> TypeGuard[list[p.LiteralExpr]]:
 class FrameManager:
     def __init__(self, typer: PythonTyper) -> None:
         self.typer: PythonTyper = typer
-        self.method_resolver: MethodRegistry = MethodRegistry(self.typer)
+        self.method_resolver: FrameMethodRegistry = FrameMethodRegistry(self.typer)
+        self.groupby_method_resolver: FrameGroupByMethodRegistry = (
+            FrameGroupByMethodRegistry(self.typer)
+        )
 
     def assign(
         self,
@@ -90,6 +103,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
@@ -156,3 +189,23 @@ class FrameManager:
             keywords=keywords,
         )
         return self.method_resolver.call(method, call)
+
+    def groupby_call(
+        self,
+        method: str,
+        location: Location,
+        call_expr: p.Expr,
+        groupby: FrameGroupBy,
+        groupby_expr: p.Expr,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+    ) -> Type:
+        call: GroupByCall = GroupByCall(
+            location=location,
+            call_expr=call_expr,
+            groupby=groupby,
+            groupby_expr=groupby_expr,
+            positional=positional,
+            keywords=keywords,
+        )
+        return self.groupby_method_resolver.call(method, call)

midas/checker/frames/frame_methods.py

@@ -2,16 +2,16 @@ from __future__ import annotations
 
 import ast
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from typing import TYPE_CHECKING, Optional
 
 import midas.ast.python as p
 from midas.ast.location import Location
-from midas.checker.dispatcher import CallDispatcher, CallResult
-from midas.checker.registry import TypesRegistry
-from midas.checker.reporter import FileReporter
+from midas.checker.dispatcher import CallResult
+from midas.checker.frames.utils import MethodRegistry, method
 from midas.checker.types import (
     ColumnType,
     DataFrameType,
+    FrameGroupBy,
     Function,
     OverloadedFunction,
     TopType,
@@ -19,22 +19,9 @@ from midas.checker.types import (
     UnknownType,
     unfold_type,
 )
-from midas.generator.collector import AssertionCollector
 
 if TYPE_CHECKING:
-    from midas.checker.python import PythonTyper, TypedExpr
-
-
-@staticmethod
-def frame_method(*names: str):
-    def wrapper(func):
-        names_: tuple[str, ...] = names
-        if len(names_) == 0:
-            names_ = (func.__name__,)
-        setattr(func, "__method_names__", names_)
-        return func
-
-    return wrapper
+    from midas.checker.python import TypedExpr
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -47,60 +34,9 @@ class Call:
     keywords: dict[str, TypedExpr]
 
 
-class _MethodRegistryMeta(type):
-    _methods: dict[str, Callable[..., Type]] = {}
-
-    def __new__(
-        cls,
-        name: str,
-        bases: tuple[type, ...],
-        namespace: dict[str, Any],
-    ):
-        new_class = super().__new__(cls, name, bases, namespace)
-        new_class._methods = {}
-        for attr in namespace.values():
-            if callable(attr) and hasattr(attr, "__method_names__"):
-                for name in attr.__method_names__:  # type: ignore
-                    new_class._methods[name] = attr  # type: ignore
-        return new_class
-
-
-class MethodRegistry(metaclass=_MethodRegistryMeta):
-    def __init__(self, typer: PythonTyper) -> None:
-        self.typer: PythonTyper = typer
-
-    @property
-    def reporter(self) -> FileReporter:
-        return self.typer.reporter
-
-    @property
-    def types(self) -> TypesRegistry:
-        return self.typer.types
-
-    @property
-    def dispatcher(self) -> CallDispatcher[p.Expr]:
-        return self.typer.dispatcher
-
-    @property
-    def assertions(self) -> AssertionCollector:
-        return self.typer.assertions
-
-    def call(
-        self,
-        method: str,
-        call: Call,
-    ) -> Type:
-        func: Optional[Callable[..., Type]] = self._methods.get(method)
-        if func is None:
-            self.reporter.warning(call.location, f"Unknown method {method}")
-            return UnknownType()
-        return func(self, call)
-
-    @frame_method("add", "__add__")
-    def add(
-        self,
-        call: Call,
-    ) -> Type:
+class FrameMethodRegistry(MethodRegistry[Call]):
+    @method("add", "__add__")
+    def add(self, call: Call) -> Type:
         # TODO: support add with scalar, sequence, Series, dict
         # TODO: check operation exists on inner column types
 
@@ -183,7 +119,7 @@ class MethodRegistry(metaclass=_MethodRegistryMeta):
 
         return result.result
 
-    @frame_method()
+    @method()
     def mean(self, call: Call) -> Type:
         with_axis = Function(
             kw_args=[
@@ -222,6 +158,67 @@ class MethodRegistry(metaclass=_MethodRegistryMeta):
         )
         return result.result
 
+    @method()
+    def groupby(self, call: Call) -> Type:
+        bool_: Type = self.types.get_type("bool")
+        function: Function = Function(
+            args=[
+                Function.Argument(
+                    pos=0,
+                    name="by",
+                    type=TopType(),
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=1,
+                    name="level",
+                    type=TopType(),
+                    required=False,
+                ),
+            ],
+            kw_args=[
+                Function.Argument(
+                    pos=2,
+                    name="as_index",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=3,
+                    name="sort",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=4,
+                    name="group_keys",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=5,
+                    name="observed",
+                    type=bool_,
+                    required=False,
+                ),
+                Function.Argument(
+                    pos=6,
+                    name="dropna",
+                    type=bool_,
+                    required=False,
+                ),
+            ],
+            returns=FrameGroupBy(frame=call.frame),
+        )
+
+        result: CallResult = self.dispatcher.get_result(
+            location=call.location,
+            callee=function,
+            positional=call.positional,
+            keywords=call.keywords,
+        )
+        return result.result
+
     def _assert_same_length(self, call_expr: p.Expr, frame1: p.Expr, frame2: p.Expr):
         func_name: str = "__midas_frame_same_length__"
         self.assertions.define(

midas/checker/frames/utils.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generic,
+    Optional,
+    Protocol,
+    Self,
+    TypeVar,
+)
+
+import midas.ast.python as p
+from midas.ast.location import Location
+from midas.checker.dispatcher import CallDispatcher
+from midas.checker.registry import TypesRegistry
+from midas.checker.reporter import FileReporter
+from midas.checker.types import Type, UnknownType
+from midas.generator.collector import AssertionCollector
+
+if TYPE_CHECKING:
+    from midas.checker.python import PythonTyper
+
+
+class _MethodRegistryMeta(type):
+    _methods: dict[str, Callable[..., Type]] = {}
+
+    def __new__(
+        cls,
+        name: str,
+        bases: tuple[type, ...],
+        namespace: dict[str, Any],
+    ):
+        new_class = super().__new__(cls, name, bases, namespace)
+        new_class._methods = {}
+        for attr in namespace.values():
+            if callable(attr) and hasattr(attr, "__method_names__"):
+                for name in attr.__method_names__:  # type: ignore
+                    new_class._methods[name] = attr  # type: ignore
+        return new_class
+
+
+class HasLocation(Protocol):
+    @property
+    def location(self) -> Location: ...
+
+
+T = TypeVar("T", bound=HasLocation)
+
+
+class MethodRegistry(Generic[T], metaclass=_MethodRegistryMeta):
+    def __init__(self, typer: PythonTyper) -> None:
+        self.typer: PythonTyper = typer
+
+    @property
+    def reporter(self) -> FileReporter:
+        return self.typer.reporter
+
+    @property
+    def types(self) -> TypesRegistry:
+        return self.typer.types
+
+    @property
+    def dispatcher(self) -> CallDispatcher[p.Expr]:
+        return self.typer.dispatcher
+
+    @property
+    def assertions(self) -> AssertionCollector:
+        return self.typer.assertions
+
+    def call(self, method: str, call: T) -> Type:
+        func: Optional[Callable[[Self, T], Type]] = self._methods.get(method)
+        if func is None:
+            self.reporter.warning(call.location, f"Unknown method {method}")
+            return UnknownType()
+        return func(self, call)
+
+
+_Self = TypeVar("_Self", bound=MethodRegistry[Any])
+Method = Callable[[_Self, T], Type]
+
+
+def method(*names: str) -> Callable[[Method[_Self, T]], Method[_Self, T]]:
+    def wrapper(func: Method[_Self, T]) -> Method[_Self, T]:
+        names_: tuple[str, ...] = names
+        if len(names_) == 0:
+            names_ = (func.__name__,)
+        setattr(func, "__method_names__", names_)
+        return func
+
+    return wrapper

midas/checker/python.py

@@ -9,7 +9,8 @@ from midas.ast.printer import MidasPrinter
 from midas.checker.dispatcher import CallDispatcher, CallResult
 from midas.checker.environment import Environment
 from midas.checker.evaluator import Evaluator
-from midas.checker.frames import FrameManager
+from midas.checker.frames.column_manager import ColumnManager
+from midas.checker.frames.frame_manager import FrameManager
 from midas.checker.operators import (
     PY_COMPARATOR_METHODS,
     PY_OPERATOR_METHODS,
@@ -22,10 +23,12 @@ from midas.checker.resolver import Resolver
 from midas.checker.types import (
     AppliedType,
     BaseType,
+    ColumnGroupBy,
     ColumnType,
     ConstraintType,
     DataFrameType,
     DerivedType,
+    FrameGroupBy,
     Function,
     GenericType,
     TopType,
@@ -81,6 +84,7 @@ class PythonTyper(
         self.reporter: FileReporter = reporter.for_file(None)
         self.types: TypesRegistry = types
         self.frame_mgr: FrameManager = FrameManager(self)
+        self.column_mgr: ColumnManager = ColumnManager(self)
         self.global_env: Environment = Preamble(self.types)
         self.env: Environment = self.global_env
         self.locals: dict[p.Expr, int] = {}
@@ -232,6 +236,39 @@ class PythonTyper(
                     keywords=keywords,
                 )
 
+            case FrameGroupBy():
+                return self.frame_mgr.groupby_call(
+                    method=method_name,
+                    location=location,
+                    call_expr=call_expr,
+                    groupby=unfolded,
+                    groupby_expr=obj[0],
+                    positional=positional,
+                    keywords=keywords,
+                )
+
+            case ColumnType():
+                return self.column_mgr.call(
+                    method=method_name,
+                    location=location,
+                    call_expr=call_expr,
+                    column=unfolded,
+                    column_expr=obj[0],
+                    positional=positional,
+                    keywords=keywords,
+                )
+
+            case ColumnGroupBy():
+                return self.column_mgr.groupby_call(
+                    method=method_name,
+                    location=location,
+                    call_expr=call_expr,
+                    groupby=unfolded,
+                    groupby_expr=obj[0],
+                    positional=positional,
+                    keywords=keywords,
+                )
+
         method: Optional[Type] = self.types.lookup_member(obj[1], method_name)
         if method is None:
             raise UndefinedMethodException
@@ -598,17 +635,17 @@ class PythonTyper(
         match expr.callee:
             case p.GetExpr(object=obj, name=method):
                 obj_type: Type = self.type_of(obj)
-                unfolded: Type = unfold_type(obj_type)
-                if isinstance(unfolded, DataFrameType):
-                    return self.frame_mgr.call(
-                        method=method,
+                return (
+                    self.call_method(
                         location=expr.location,
                         call_expr=expr,
-                        frame=unfolded,
-                        frame_expr=obj,
+                        obj=(obj, obj_type),
+                        method_name=method,
                         positional=positional,
                         keywords=keywords,
                     )
+                    or UnknownType()
+                )
 
         callee: Type = self.type_of(expr.callee)
         result: CallResult = self.dispatcher.get_result(
@@ -770,6 +807,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 +1134,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
+        )

midas/checker/types.py

@@ -187,6 +187,22 @@ class DataFrameType:
         type: ColumnType
 
 
+@dataclass(frozen=True, kw_only=True)
+class FrameGroupBy:
+    frame: DataFrameType
+
+    def __str__(self) -> str:
+        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(
@@ -305,11 +321,20 @@ def substitute_typevars(type: Type, substitutions: dict[str, Type]) -> Type:
                 columns=list(map(sub_column, columns)),
             )
 
+        case FrameGroupBy(frame=frame):
+            return FrameGroupBy(
+                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
 
         case TopType() | GenericType():
-
             raise NotImplementedError(f"Unsupported type {type}")
 
         # Ensure exhaustiveness
@@ -382,6 +407,12 @@ def to_annotation(type: Type) -> str:
         case DataFrameType():
             return "pd.DataFrame"
 
+        case FrameGroupBy():
+            return "pd.api.typing.DataFrameGroupBy"
+
+        case ColumnGroupBy():
+            return "pd.api.typing.SeriesGroupBy"
+
         case _:
             assert_never(type)
 
@@ -410,4 +441,6 @@ Type = (
     | TupleType
     | ColumnType
     | DataFrameType
+    | FrameGroupBy
+    | ColumnGroupBy
 )

midas/checker/unifier.py

@@ -4,6 +4,8 @@ from typing import Optional
 from midas.checker.registry import TypesRegistry
 from midas.checker.types import (
     AppliedType,
+    ColumnType,
+    DataFrameType,
     Function,
     GenericType,
     TopType,
@@ -98,6 +100,30 @@ class Unifier:
 
                 return substitutions
 
+            case (
+                DataFrameType(columns=template_columns),
+                DataFrameType(columns=concrete_columns),
+            ) if len(template_columns) == len(concrete_columns):
+                substitutions: dict[str, Type] = {}
+                for template_column, concrete_column in zip(
+                    template_columns, concrete_columns
+                ):
+                    if template_column.index != concrete_column or (
+                        template_column.name != concrete_column.name
+                    ):
+                        self.logger.debug(
+                            f"Column mismatch: template={template_column}, concrete={concrete_column}"
+                        )
+                        raise UnificationError
+                    new_substistutions: dict[str, Type] = self.match(
+                        template_column.type, concrete_column.type
+                    )
+                    substitutions = self.merge(substitutions, new_substistutions)
+                return substitutions
+
+            case (ColumnType(type=template_column), ColumnType(type=concrete_column)):
+                return self.match(template_column, concrete_column)
+
             case (Function(), Function()):
                 mapped: list[tuple[Function.Argument, Function.Argument]] = (
                     self.map_params(template, concrete)

midas/generator/generator.py

@@ -14,12 +14,14 @@ from midas.checker.registry import TypesRegistry
 from midas.checker.types import (
     AppliedType,
     BaseType,
+    ColumnGroupBy,
     ColumnType,
     ComplexType,
     ConstraintType,
     DataFrameType,
     DerivedType,
     ExtensionType,
+    FrameGroupBy,
     Function,
     GenericType,
     OverloadedFunction,
@@ -366,7 +368,7 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
         self, src_location: Location, expr: ast.expr, type: Type
     ) -> list[ast.stmt]:
         match type:
-            case UnknownType():
+            case UnknownType() | TopType():
                 return []
 
             case BaseType(name=name):
@@ -497,12 +499,13 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
                 return asserts
 
             case (
-                TopType()
-                | Function()
+                Function()
                 | OverloadedFunction()
                 | ComplexType()
                 | ExtensionType()
                 | GenericType()
+                | FrameGroupBy()
+                | ColumnGroupBy()
             ):
                 self.logger.warning(f"Can't make assertion for type {type}")
                 return []

midas/generator/stubs.py

@@ -6,12 +6,14 @@ from midas.checker.registry import Member, TypesRegistry
 from midas.checker.types import (
     AppliedType,
     BaseType,
+    ColumnGroupBy,
     ColumnType,
     ComplexType,
     ConstraintType,
     DataFrameType,
     DerivedType,
     ExtensionType,
+    FrameGroupBy,
     Function,
     GenericType,
     OverloadedFunction,
@@ -287,6 +289,32 @@ class StubsGenerator:
                     attr="DataFrame",
                 )
 
+            case FrameGroupBy():
+                self.import_pandas = True
+                return ast.Attribute(
+                    value=ast.Attribute(
+                        value=ast.Attribute(
+                            value=ast.Name(id="pd"),
+                            attr="api",
+                        ),
+                        attr="typing",
+                    ),
+                    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)