refactor: Reuse eager-backend iter_columns

Utilizing (#2115) and (#2104)

Author: dangotbanned

narwhals/_arrow/selectors.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
-from typing import Iterator
 
 from narwhals._arrow.expr import ArrowExpr
 from narwhals._selectors import CompliantSelector
-from narwhals._selectors import CompliantSelectorNamespace
+from narwhals._selectors import EagerSelectorNamespace
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -17,15 +16,7 @@
     from narwhals.utils import _FullContext
 
 
-class ArrowSelectorNamespace(CompliantSelectorNamespace["ArrowDataFrame", "ArrowSeries"]):
-    def _iter_columns(self, df: ArrowDataFrame) -> Iterator[ArrowSeries]:
-        from narwhals._arrow.series import ArrowSeries
-
-        for col, ser in zip(df.columns, df._native_frame.itercolumns()):
-            yield ArrowSeries(
-                ser, name=col, backend_version=df._backend_version, version=df._version
-            )
-
+class ArrowSelectorNamespace(EagerSelectorNamespace["ArrowDataFrame", "ArrowSeries"]):
     def _selector(
         self,
         call: EvalSeries[ArrowDataFrame, ArrowSeries],

narwhals/_pandas_like/selectors.py

@@ -1,14 +1,12 @@
 from __future__ import annotations
 
-from functools import partial
 from typing import TYPE_CHECKING
-from typing import Iterator
 
 from narwhals._pandas_like.dataframe import PandasLikeDataFrame
 from narwhals._pandas_like.expr import PandasLikeExpr
 from narwhals._pandas_like.series import PandasLikeSeries
 from narwhals._selectors import CompliantSelector
-from narwhals._selectors import CompliantSelectorNamespace
+from narwhals._selectors import EagerSelectorNamespace
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -21,20 +19,8 @@
 
 
 class PandasSelectorNamespace(
-    CompliantSelectorNamespace["PandasLikeDataFrame", "PandasLikeSeries"]
+    EagerSelectorNamespace["PandasLikeDataFrame", "PandasLikeSeries"]
 ):
-    def _iter_columns(self, df: PandasLikeDataFrame) -> Iterator[PandasLikeSeries]:
-        from narwhals._pandas_like.series import PandasLikeSeries
-
-        series = partial(
-            PandasLikeSeries,
-            implementation=df._implementation,
-            backend_version=df._backend_version,
-            version=df._version,
-        )
-        for _col, ser in df._native_frame.items():  # noqa: PERF102
-            yield series(ser)
-
     def _selector(
         self,
         call: EvalSeries[PandasLikeDataFrame, PandasLikeSeries],

narwhals/_selectors.py

@@ -52,6 +52,7 @@
 
 SeriesT = TypeVar("SeriesT", bound="CompliantSeries")
 FrameT = TypeVar("FrameT", bound="CompliantDataFrame[Any] | CompliantLazyFrame")
+DataFrameT = TypeVar("DataFrameT", bound="CompliantDataFrame[Any]")
 SelectorOrExpr: TypeAlias = (
     "CompliantSelector[FrameT, SeriesT] | CompliantExpr[FrameT, SeriesT]"
 )
@@ -168,6 +169,13 @@ def names(df: FrameT) -> Sequence[str]:
         return self._selector(series, names)
 
 
+class EagerSelectorNamespace(
+    CompliantSelectorNamespace[DataFrameT, SeriesT], Protocol[DataFrameT, SeriesT]
+):
+    def _iter_columns(self, df: DataFrameT, /) -> Iterator[SeriesT]:
+        yield from df.iter_columns()
+
+
 class LazySelectorNamespace(
     CompliantSelectorNamespace[FrameT, SeriesT], Protocol[FrameT, SeriesT]
 ):

narwhals/typing.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
+from typing import Iterator
 from typing import Literal
 from typing import Protocol
 from typing import Sequence
@@ -83,6 +84,7 @@ def columns(self) -> Sequence[str]: ...
     @property
     def schema(self) -> Mapping[str, DType]: ...
     def get_column(self, name: str) -> CompliantSeriesT_co: ...
+    def iter_columns(self) -> Iterator[CompliantSeriesT_co]: ...
 
 
 class CompliantLazyFrame(Protocol):