From 205d19fb721f9049b30e1c35f3c377345192c907 Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Wed, 8 Jul 2026 10:03:14 +0200
Subject: [PATCH] feat(checker): try to filter groupby columns

---
 midas/checker/frames/frame_methods.py | 42 ++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py
index dfe0698..93355d6 100644
--- a/midas/checker/frames/frame_methods.py
+++ b/midas/checker/frames/frame_methods.py
@@ -598,8 +598,48 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         )
         return result.result
 
+    def _filter_groupby_columns(
+        self, frame: DataFrameType, by: TypedExpr
+    ) -> DataFrameType:
+        by_columns: list[str] = []
+
+        by_expr, _ = by
+
+        if not isinstance(by_expr, p.ListExpr):
+            return frame
+
+        for item in by_expr.items:
+            match item:
+                case p.LiteralExpr(value=str() as name):
+                    by_columns.append(name)
+
+        new_columns: list[DataFrameType.Column] = []
+        for column in frame.columns:
+            if column.name in by_columns:
+                continue
+            new_columns.append(
+                DataFrameType.Column(
+                    index=len(new_columns),
+                    name=column.name,
+                    type=column.type,
+                )
+            )
+
+        return DataFrameType(columns=new_columns)
+
     @method()
     def groupby(self, call: Call) -> Type:
+        new_frame: DataFrameType = call.frame
+
+        by: Optional[TypedExpr] = None
+        if len(call.positional) != 0:
+            by = call.positional[0]
+        elif "by" in call.keywords:
+            by = call.keywords["by"]
+
+        if by is not None:
+            new_frame = self._filter_groupby_columns(call.frame, by)
+
         bool_: Type = self.types.get_type("bool")
         function: Function = Function(
             params=ParamSpec(
@@ -629,7 +669,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
                     )
                 ],
             ),
-            returns=FrameGroupBy(frame=call.frame),
+            returns=FrameGroupBy(frame=new_frame),
         )
 
         result: CallResult = self.dispatcher.get_result(