feat(checker): handle single string literal in groupby
This commit is contained in:
@@ -601,17 +601,31 @@ class FrameMethodRegistry(MethodRegistry[Call]):
|
|||||||
def _filter_groupby_columns(
|
def _filter_groupby_columns(
|
||||||
self, frame: DataFrameType, by: TypedExpr
|
self, frame: DataFrameType, by: TypedExpr
|
||||||
) -> DataFrameType:
|
) -> 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_columns: list[str] = []
|
||||||
|
|
||||||
by_expr, _ = by
|
by_expr, _ = by
|
||||||
|
|
||||||
if not isinstance(by_expr, p.ListExpr):
|
match by_expr:
|
||||||
return frame
|
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:
|
case p.LiteralExpr(value=str() as name):
|
||||||
match item:
|
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] = []
|
new_columns: list[DataFrameType.Column] = []
|
||||||
for column in frame.columns:
|
for column in frame.columns:
|
||||||
|
|||||||
Reference in New Issue
Block a user