fix(typing): EagerNamespace.all_horizontal

Was not easy getting that to please both `mypy` & `pyright`

Author: dangotbanned

narwhals/_arrow/namespace.py

@@ -47,7 +47,7 @@
     _Scalar: TypeAlias = Any
 
 
-class ArrowNamespace(EagerNamespace[ArrowDataFrame, ArrowSeries]):
+class ArrowNamespace(EagerNamespace[ArrowDataFrame, ArrowSeries, ArrowExpr]):
     @property
     def _expr(self) -> type[ArrowExpr]:
         return ArrowExpr
@@ -123,9 +123,7 @@ def _lit_arrow_series(_: ArrowDataFrame) -> ArrowSeries:
             version=self._version,
         )
 
-    # NOTE: Needs to be resolved in `EagerNamespace`
-    # Probably, by adding an `EagerExprT` typevar
-    def all_horizontal(self: Self, *exprs: ArrowExpr) -> ArrowExpr:  # type: ignore[override]
+    def all_horizontal(self: Self, *exprs: ArrowExpr) -> ArrowExpr:
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:
             series = chain.from_iterable(expr(df) for expr in exprs)
             return [reduce(operator.and_, align_series_full_broadcast(*series))]

narwhals/_compliant/expr.py

@@ -285,7 +285,9 @@ def __call__(self, df: EagerDataFrameT) -> Sequence[EagerSeriesT]:
     def __repr__(self) -> str:  # pragma: no cover
         return f"{type(self).__name__}(depth={self._depth}, function_name={self._function_name})"
 
-    def __narwhals_namespace__(self) -> EagerNamespace[EagerDataFrameT, EagerSeriesT]: ...
+    def __narwhals_namespace__(
+        self,
+    ) -> EagerNamespace[EagerDataFrameT, EagerSeriesT, Self]: ...
     def __narwhals_expr__(self) -> None: ...
 
     @classmethod

narwhals/_compliant/namespace.py

@@ -7,12 +7,12 @@
 from narwhals._compliant.typing import CompliantFrameT
 from narwhals._compliant.typing import CompliantSeriesOrNativeExprT_co
 from narwhals._compliant.typing import EagerDataFrameT
-from narwhals._compliant.typing import EagerSeriesT
+from narwhals._compliant.typing import EagerExprT
+from narwhals._compliant.typing import EagerSeriesT_co
 from narwhals.utils import deprecated
 
 if TYPE_CHECKING:
     from narwhals._compliant.expr import CompliantExpr
-    from narwhals._compliant.expr import EagerExpr
     from narwhals._compliant.selectors import CompliantSelectorNamespace
     from narwhals.dtypes import DType
 
@@ -31,25 +31,23 @@ def selectors(self) -> CompliantSelectorNamespace[Any, Any]: ...
 
 
 class EagerNamespace(
-    CompliantNamespace[EagerDataFrameT, EagerSeriesT],
-    Protocol[EagerDataFrameT, EagerSeriesT],
+    CompliantNamespace[EagerDataFrameT, EagerSeriesT_co],
+    Protocol[EagerDataFrameT, EagerSeriesT_co, EagerExprT],
 ):
     # NOTE: Supporting moved ops
     # - `self_create_expr_from_callable` -> `self._expr._from_callable`
     # - `self_create_expr_from_series` -> `self._expr._from_series`
     @property
-    def _expr(self) -> type[EagerExpr[EagerDataFrameT, EagerSeriesT]]: ...
+    def _expr(self) -> type[EagerExprT]: ...
 
     # NOTE: Supporting moved ops
     # - `self._create_series_from_scalar` -> `EagerSeries()._from_scalar`
     #   - Was dependent on a `reference_series`, so is now an instance method
     # - `<class>._from_iterable` -> `self._series._from_iterable`
     @property
-    def _series(self) -> type[EagerSeriesT]: ...
+    def _series(self) -> type[EagerSeriesT_co]: ...
 
-    def all_horizontal(
-        self, *exprs: EagerExpr[EagerDataFrameT, EagerSeriesT]
-    ) -> EagerExpr[EagerDataFrameT, EagerSeriesT]: ...
+    def all_horizontal(self, *exprs: EagerExprT) -> EagerExprT: ...
 
     @deprecated("ref'd in untyped code")
-    def _create_compliant_series(self, value: Any) -> EagerSeriesT: ...
+    def _create_compliant_series(self, value: Any) -> EagerSeriesT_co: ...

narwhals/_compliant/series.py

@@ -56,7 +56,7 @@ def _from_iterable(
         cls: type[Self], data: Iterable[Any], name: str, *, context: _FullContext
     ) -> Self: ...
 
-    def __narwhals_namespace__(self) -> EagerNamespace[Any, Self]: ...
+    def __narwhals_namespace__(self) -> EagerNamespace[Any, Self, Any]: ...
 
-    def _to_expr(self) -> EagerExpr[Any, Self]:
-        return self.__narwhals_namespace__()._expr._from_series(self)
+    def _to_expr(self) -> EagerExpr[Any, Any]:
+        return self.__narwhals_namespace__()._expr._from_series(self)  # type: ignore[no-any-return]

narwhals/_pandas_like/namespace.py

@@ -40,7 +40,9 @@
     _Scalar: TypeAlias = Any
 
 
-class PandasLikeNamespace(EagerNamespace[PandasLikeDataFrame, PandasLikeSeries]):
+class PandasLikeNamespace(
+    EagerNamespace[PandasLikeDataFrame, PandasLikeSeries, PandasLikeExpr]
+):
     @property
     def _expr(self) -> type[PandasLikeExpr]:
         return PandasLikeExpr
@@ -149,9 +151,7 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
             context=self,
         )
 
-    # NOTE: Needs to be resolved in `EagerNamespace`
-    # Probably, by adding an `EagerExprT` typevar
-    def all_horizontal(self: Self, *exprs: PandasLikeExpr) -> PandasLikeExpr:  # type: ignore[override]
+    def all_horizontal(self: Self, *exprs: PandasLikeExpr) -> PandasLikeExpr:
         def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
             series = align_series_full_broadcast(
                 *(s for _expr in exprs for s in _expr(df))