chore(typing): fix group_by (#2105)

- Think this is pretty straightforward
- Tiny reduction in LOC is a bonus

Author: dangotbanned

narwhals/group_by.py

@@ -6,29 +6,25 @@
 from typing import Iterable
 from typing import Iterator
 from typing import TypeVar
-from typing import cast
 
 from narwhals._expression_parsing import all_exprs_are_scalar_like
-from narwhals.dataframe import DataFrame
-from narwhals.dataframe import LazyFrame
 from narwhals.exceptions import InvalidOperationError
+from narwhals.typing import DataFrameT
 from narwhals.utils import flatten
 from narwhals.utils import tupleify
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-    from narwhals.dataframe import DataFrame
     from narwhals.dataframe import LazyFrame
     from narwhals.expr import Expr
 
-DataFrameT = TypeVar("DataFrameT")
-LazyFrameT = TypeVar("LazyFrameT")
+LazyFrameT = TypeVar("LazyFrameT", bound="LazyFrame[Any]")
 
 
 class GroupBy(Generic[DataFrameT]):
     def __init__(self: Self, df: DataFrameT, *keys: str, drop_null_keys: bool) -> None:
-        self._df = cast("DataFrame[Any]", df)
+        self._df: DataFrameT = df
         self._keys = keys
         self._grouped = self._df._compliant_frame.group_by(
             *self._keys, drop_null_keys=drop_null_keys
@@ -131,20 +127,18 @@ def agg(self: Self, *aggs: Expr | Iterable[Expr], **named_aggs: Expr) -> DataFra
                 for key, value in named_aggs.items()
             ),
         )
-        return self._df._from_compliant_dataframe(  # type: ignore[return-value]
-            self._grouped.agg(*compliant_aggs),
-        )
+        return self._df._from_compliant_dataframe(self._grouped.agg(*compliant_aggs))
 
     def __iter__(self: Self) -> Iterator[tuple[Any, DataFrameT]]:
-        yield from (  # type: ignore[misc]
+        yield from (
             (tupleify(key), self._df._from_compliant_dataframe(df))
             for (key, df) in self._grouped.__iter__()
         )
 
 
 class LazyGroupBy(Generic[LazyFrameT]):
     def __init__(self: Self, df: LazyFrameT, *keys: str, drop_null_keys: bool) -> None:
-        self._df = cast("LazyFrame[Any]", df)
+        self._df: LazyFrameT = df
         self._keys = keys
         self._grouped = self._df._compliant_frame.group_by(
             *self._keys, drop_null_keys=drop_null_keys
@@ -231,6 +225,4 @@ def agg(self: Self, *aggs: Expr | Iterable[Expr], **named_aggs: Expr) -> LazyFra
                 for key, value in named_aggs.items()
             ),
         )
-        return self._df._from_compliant_dataframe(  # type: ignore[return-value]
-            self._grouped.agg(*compliant_aggs),
-        )
+        return self._df._from_compliant_dataframe(self._grouped.agg(*compliant_aggs))