From 607ff5398718883b81e5af3a0a046b8e1d479148 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Wed, 8 Jul 2026 10:12:54 +0200 Subject: [PATCH] feat(checker): handle single string literal in groupby --- midas/checker/frames/frame_methods.py | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index 93355d6..24a41f7 100644 --- a/midas/checker/frames/frame_methods.py +++ b/midas/checker/frames/frame_methods.py @@ -601,17 +601,31 @@ class FrameMethodRegistry(MethodRegistry[Call]): def _filter_groupby_columns( self, frame: DataFrameType, by: TypedExpr ) -> DataFrameType: + """Remove columns passed as string literals in groupby's `by` argument + + Args: + frame (DataFrameType): the original dataframe + by (TypedExpr): the by argument + + Returns: + DataFrameType: the filtered dataframe + """ by_columns: list[str] = [] by_expr, _ = by - if not isinstance(by_expr, p.ListExpr): - return frame + match by_expr: + case p.ListExpr(items=items): + for item in items: + match item: + case p.LiteralExpr(value=str() as name): + by_columns.append(name) - for item in by_expr.items: - match item: - case p.LiteralExpr(value=str() as name): - by_columns.append(name) + case p.LiteralExpr(value=str() as name): + by_columns.append(name) + + if len(by_columns) == 0: + return frame new_columns: list[DataFrameType.Column] = [] for column in frame.columns: