keep simplifying

Author: MarcoGorelli

narwhals/_arrow/group_by.py

@@ -11,7 +11,6 @@
 import pyarrow.compute as pc
 
 from narwhals._expression_parsing import is_simple_aggregation
-from narwhals._expression_parsing import parse_into_exprs
 from narwhals.exceptions import AnonymousExprError
 from narwhals.utils import generate_temporary_column_name
 from narwhals.utils import remove_prefix
@@ -20,8 +19,8 @@
     from typing_extensions import Self
 
     from narwhals._arrow.dataframe import ArrowDataFrame
+    from narwhals._arrow.expr import ArrowExpr
     from narwhals._arrow.series import ArrowSeries
-    from narwhals._arrow.typing import IntoArrowExpr
     from narwhals.typing import CompliantExpr
 
 POLARS_TO_ARROW_AGGREGATIONS = {
@@ -51,14 +50,8 @@ def __init__(
 
     def agg(
         self: Self,
-        *aggs: IntoArrowExpr,
-        **named_aggs: IntoArrowExpr,
+        *exprs: ArrowExpr,
     ) -> ArrowDataFrame:
-        exprs = parse_into_exprs(
-            *aggs,
-            namespace=self._df.__narwhals_namespace__(),
-            **named_aggs,
-        )
         for expr in exprs:
             if expr._output_names is None:
                 msg = "group_by.agg"

narwhals/_arrow/namespace.py

@@ -365,9 +365,7 @@ def concat_str(
         separator: str,
         ignore_nulls: bool,
     ) -> ArrowExpr:
-        parsed_exprs = [
-            *parse_into_exprs(*exprs, namespace=self),
-        ]
+        parsed_exprs = parse_into_exprs(*exprs, namespace=self)
         dtypes = import_dtypes_module(self._version)
 
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:

narwhals/_dask/group_by.py

@@ -88,13 +88,8 @@ def __init__(
 
     def agg(
         self: Self,
-        *aggs: DaskExpr,
-        **named_aggs: DaskExpr,
+        *exprs: DaskExpr,
     ) -> DaskLazyFrame:
-        exprs: list[DaskExpr] = [
-            *aggs,
-            *(val.alias(key) for key, val in named_aggs.items()),
-        ]
         output_names: list[str] = copy(self._keys)
         for expr in exprs:
             if expr._output_names is None:

narwhals/_duckdb/group_by.py

@@ -27,12 +27,8 @@ def __init__(
 
     def agg(
         self: Self,
-        *aggs: IntoDuckDBExpr,
-        **named_aggs: IntoDuckDBExpr,
+        *exprs: IntoDuckDBExpr,
     ) -> DuckDBLazyFrame:
-        exprs = tuple(
-            *(x for x in aggs), *(val.alias(key) for key, val in named_aggs.items())
-        )
         output_names: list[str] = copy(self._keys)
         for expr in exprs:
             if expr._output_names is None:  # pragma: no cover

narwhals/_pandas_like/group_by.py

@@ -10,7 +10,6 @@
 from typing import Sequence
 
 from narwhals._expression_parsing import is_simple_aggregation
-from narwhals._expression_parsing import parse_into_exprs
 from narwhals._pandas_like.utils import horizontal_concat
 from narwhals._pandas_like.utils import native_series_from_iterable
 from narwhals._pandas_like.utils import select_columns_by_name
@@ -24,8 +23,8 @@
     from typing_extensions import Self
 
     from narwhals._pandas_like.dataframe import PandasLikeDataFrame
+    from narwhals._pandas_like.expr import PandasLikeExpr
     from narwhals._pandas_like.series import PandasLikeSeries
-    from narwhals._pandas_like.typing import IntoPandasLikeExpr
     from narwhals.typing import CompliantExpr
 
 POLARS_TO_PANDAS_AGGREGATIONS = {
@@ -83,14 +82,8 @@ def __init__(
 
     def agg(
         self: Self,
-        *aggs: IntoPandasLikeExpr,
-        **named_aggs: IntoPandasLikeExpr,
+        *exprs: PandasLikeExpr,
     ) -> PandasLikeDataFrame:
-        exprs = parse_into_exprs(
-            *aggs,
-            namespace=self._df.__narwhals_namespace__(),
-            **named_aggs,
-        )
         implementation: Implementation = self._df._implementation
         output_names: list[str] = copy(self._keys)
         for expr in exprs:

narwhals/_pandas_like/namespace.py

@@ -381,9 +381,7 @@ def concat_str(
         separator: str,
         ignore_nulls: bool,
     ) -> PandasLikeExpr:
-        parsed_exprs = [
-            *parse_into_exprs(*exprs, namespace=self),
-        ]
+        parsed_exprs = parse_into_exprs(*exprs, namespace=self)
         dtypes = import_dtypes_module(self._version)
 
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:

narwhals/_polars/group_by.py

@@ -24,11 +24,9 @@ def __init__(
         else:
             self._grouped = df._native_frame.group_by(keys)
 
-    def agg(self: Self, *aggs: PolarsExpr, **named_aggs: PolarsExpr) -> PolarsDataFrame:
-        aggs, named_aggs = extract_args_kwargs(aggs, named_aggs)  # type: ignore[assignment]
-        return self._compliant_frame._from_native_frame(
-            self._grouped.agg(*aggs, **named_aggs),
-        )
+    def agg(self: Self, *aggs: PolarsExpr) -> PolarsDataFrame:
+        aggs, _ = extract_args_kwargs(aggs, {})  # type: ignore[assignment]
+        return self._compliant_frame._from_native_frame(self._grouped.agg(*aggs))
 
     def __iter__(self: Self) -> Iterator[tuple[tuple[str, ...], PolarsDataFrame]]:
         for key, df in self._grouped:
@@ -46,8 +44,6 @@ def __init__(
         else:
             self._grouped = df._native_frame.group_by(keys)
 
-    def agg(self: Self, *aggs: PolarsExpr, **named_aggs: PolarsExpr) -> PolarsLazyFrame:
-        aggs, named_aggs = extract_args_kwargs(aggs, named_aggs)  # type: ignore[assignment]
-        return self._compliant_frame._from_native_frame(
-            self._grouped.agg(*aggs, **named_aggs),
-        )
+    def agg(self: Self, *aggs: PolarsExpr) -> PolarsLazyFrame:
+        aggs, _ = extract_args_kwargs(aggs, {})  # type: ignore[assignment]
+        return self._compliant_frame._from_native_frame(self._grouped.agg(*aggs))

narwhals/_spark_like/group_by.py

@@ -44,12 +44,8 @@ def __init__(
 
     def agg(
         self: Self,
-        *aggs: SparkLikeExpr,
-        **named_aggs: SparkLikeExpr,
+        *exprs: SparkLikeExpr,
     ) -> SparkLikeLazyFrame:
-        exprs = tuple(
-            *(x for x in aggs), *(val.alias(key) for key, val in named_aggs.items())
-        )
         output_names: list[str] = copy(self._keys)
         for expr in exprs:
             if expr._output_names is None:  # pragma: no cover

narwhals/group_by.py

@@ -11,12 +11,13 @@
 from narwhals.dataframe import DataFrame
 from narwhals.dataframe import LazyFrame
 from narwhals.exceptions import InvalidOperationError
+from narwhals.utils import flatten
 from narwhals.utils import tupleify
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-    from narwhals.typing import IntoExpr
+    from narwhals.expr import Expr
 
 DataFrameT = TypeVar("DataFrameT")
 LazyFrameT = TypeVar("LazyFrameT")
@@ -30,9 +31,7 @@ def __init__(self: Self, df: DataFrameT, *keys: str, drop_null_keys: bool) -> No
             *self._keys, drop_null_keys=drop_null_keys
         )
 
-    def agg(
-        self: Self, *aggs: IntoExpr | Iterable[IntoExpr], **named_aggs: IntoExpr
-    ) -> DataFrameT:
+    def agg(self: Self, *aggs: Expr | Iterable[Expr], **named_aggs: Expr) -> DataFrameT:
         """Compute aggregations for each group of a group by operation.
 
         Arguments:
@@ -112,7 +111,8 @@ def agg(
             │ c   ┆ 3   ┆ 1   │
             └─────┴─────┴─────┘
         """
-        if not all(getattr(x, "_aggregates", True) for x in aggs) and all(
+        flat_aggs = tuple(flatten(aggs))
+        if not all(getattr(x, "_aggregates", True) for x in flat_aggs) and all(
             getattr(x, "_aggregates", True) for x in named_aggs.values()
         ):
             msg = (
@@ -122,11 +122,16 @@ def agg(
                 "but `df.group_by('a').agg(nw.col('b'))` is not."
             )
             raise InvalidOperationError(msg)
-        compliant_aggs, compliant_named_aggs = self._df._flatten_and_extract(
-            *aggs, **named_aggs
+        plx = self._df.__narwhals_namespace__()
+        compliant_aggs = (
+            *(x._to_compliant_expr(plx) for x in flat_aggs),
+            *(
+                value._to_compliant_expr(plx).alias(key)
+                for key, value in named_aggs.items()
+            ),
         )
         return self._df._from_compliant_dataframe(  # type: ignore[return-value]
-            self._grouped.agg(*compliant_aggs, **compliant_named_aggs),
+            self._grouped.agg(*compliant_aggs),
         )
 
     def __iter__(self: Self) -> Iterator[tuple[Any, DataFrameT]]:
@@ -144,9 +149,7 @@ def __init__(self: Self, df: LazyFrameT, *keys: str, drop_null_keys: bool) -> No
             *self._keys, drop_null_keys=drop_null_keys
         )
 
-    def agg(
-        self: Self, *aggs: IntoExpr | Iterable[IntoExpr], **named_aggs: IntoExpr
-    ) -> LazyFrameT:
+    def agg(self: Self, *aggs: Expr | Iterable[Expr], **named_aggs: Expr) -> LazyFrameT:
         """Compute aggregations for each group of a group by operation.
 
         If a library does not support lazy execution, then this is a no-op.
@@ -210,7 +213,8 @@ def agg(
             │ c   ┆ 3   ┆ 1   │
             └─────┴─────┴─────┘
         """
-        if not all(getattr(x, "_aggregates", True) for x in aggs) and all(
+        flat_aggs = tuple(flatten(aggs))
+        if not all(getattr(x, "_aggregates", True) for x in flat_aggs) and all(
             getattr(x, "_aggregates", True) for x in named_aggs.values()
         ):
             msg = (
@@ -220,9 +224,14 @@ def agg(
                 "but `df.group_by('a').agg(nw.col('b'))` is not."
             )
             raise InvalidOperationError(msg)
-        compliant_aggs, compliant_named_aggs = self._df._flatten_and_extract(
-            *aggs, **named_aggs
+        plx = self._df.__narwhals_namespace__()
+        compliant_aggs = (
+            *(x._to_compliant_expr(plx) for x in flat_aggs),
+            *(
+                value._to_compliant_expr(plx).alias(key)
+                for key, value in named_aggs.items()
+            ),
         )
         return self._df._from_compliant_dataframe(  # type: ignore[return-value]
-            self._grouped.agg(*compliant_aggs, **compliant_named_aggs),
+            self._grouped.agg(*compliant_aggs),
         )