fix(typing): Mostly resolve invariance issues

The problem is that type checkers special-case `__init__` to allow covariance
... despite the same thing not being allowed on a classmethod (`from_<...>`)

Author: dangotbanned

narwhals/_plan/arrow/dataframe.py

@@ -24,10 +24,9 @@
 
     from typing_extensions import Self
 
-    from narwhals._arrow.typing import ChunkedArrayAny
+    from narwhals._arrow.typing import ChunkedArrayAny  # noqa: F401
     from narwhals._plan.arrow.expr import ArrowExpr as Expr, ArrowScalar as Scalar
     from narwhals._plan.arrow.namespace import ArrowNamespace
-    from narwhals._plan.dataframe import DataFrame as NwDataFrame
     from narwhals._plan.expressions import ExprIR, NamedIR
     from narwhals._plan.options import SortMultipleOptions
     from narwhals._plan.typing import Seq
@@ -63,11 +62,6 @@ def schema(self) -> dict[str, DType]:
     def __len__(self) -> int:
         return self.native.num_rows
 
-    def to_narwhals(self) -> NwDataFrame[pa.Table, ChunkedArrayAny]:
-        from narwhals._plan.dataframe import DataFrame
-
-        return DataFrame[pa.Table, "ChunkedArrayAny"]._from_compliant(self)
-
     @classmethod
     def from_dict(
         cls, data: Mapping[str, Any], /, *, schema: IntoSchema | None = None

narwhals/_plan/compliant/dataframe.py

@@ -114,7 +114,11 @@ def to_dict(
     def to_dict(
         self, *, as_series: bool
     ) -> dict[str, SeriesT] | dict[str, list[Any]]: ...
-    def to_narwhals(self) -> DataFrame[NativeDataFrameT, NativeSeriesT]: ...
+    def to_narwhals(self) -> DataFrame[NativeDataFrameT, NativeSeriesT]:
+        from narwhals._plan.dataframe import DataFrame
+
+        return DataFrame[NativeDataFrameT, NativeSeriesT](self)
+
     def with_row_index(self, name: str) -> Self: ...
 
 

narwhals/_plan/compliant/series.py

@@ -70,6 +70,6 @@ def to_list(self) -> list[Any]: ...
     def to_narwhals(self) -> Series[NativeSeriesT]:
         from narwhals._plan.series import Series
 
-        return Series[NativeSeriesT]._from_compliant(self)
+        return Series[NativeSeriesT](self)
 
     def to_numpy(self, dtype: Any = None, *, copy: bool | None = None) -> _1DArray: ...

narwhals/_plan/dataframe.py

@@ -10,7 +10,9 @@
 from narwhals._plan.typing import (
     IntoExpr,
     NativeDataFrameT,
+    NativeDataFrameT_co,
     NativeFrameT,
+    NativeFrameT_co,
     NativeSeriesT,
     OneOrIterable,
 )
@@ -21,15 +23,13 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    import pyarrow as pa
     from typing_extensions import Self
 
     from narwhals._plan.compliant.dataframe import CompliantDataFrame, CompliantFrame
-    from narwhals.typing import NativeFrame
 
 
-class BaseFrame(Generic[NativeFrameT]):
-    _compliant: CompliantFrame[Any, NativeFrameT]
+class BaseFrame(Generic[NativeFrameT_co]):
+    _compliant: CompliantFrame[Any, NativeFrameT_co]
     _version: ClassVar[Version] = Version.MAIN
 
     @property
@@ -47,30 +47,26 @@ def columns(self) -> list[str]:
     def __repr__(self) -> str:  # pragma: no cover
         return generate_repr(f"nw.{type(self).__name__}", self.to_native().__repr__())
 
-    @classmethod
-    def from_native(cls, native: Any, /) -> Self:
-        raise NotImplementedError
+    def __init__(self, compliant: Any, /) -> None:
+        self._compliant = compliant
 
-    @classmethod
-    def _from_compliant(cls, compliant: CompliantFrame[Any, NativeFrameT], /) -> Self:
-        obj = cls.__new__(cls)
-        obj._compliant = compliant
-        return obj
+    def _with_compliant(self, compliant: CompliantFrame[Any, NativeFrameT], /) -> Self:
+        return type(self)(compliant)
 
-    def to_native(self) -> NativeFrameT:
+    def to_native(self) -> NativeFrameT_co:
         return self._compliant.native
 
     def select(self, *exprs: OneOrIterable[IntoExpr], **named_exprs: Any) -> Self:
         named_irs, schema = prepare_projection(
             _parse.parse_into_seq_of_expr_ir(*exprs, **named_exprs), schema=self
         )
-        return self._from_compliant(self._compliant.select(schema.select_irs(named_irs)))
+        return self._with_compliant(self._compliant.select(schema.select_irs(named_irs)))
 
     def with_columns(self, *exprs: OneOrIterable[IntoExpr], **named_exprs: Any) -> Self:
         named_irs, schema = prepare_projection(
             _parse.parse_into_seq_of_expr_ir(*exprs, **named_exprs), schema=self
         )
-        return self._from_compliant(
+        return self._with_compliant(
             self._compliant.with_columns(schema.with_columns_irs(named_irs))
         )
 
@@ -85,32 +81,33 @@ def sort(
             by, *more_by, descending=descending, nulls_last=nulls_last
         )
         named_irs, _ = prepare_projection(sort, schema=self)
-        return self._from_compliant(self._compliant.sort(named_irs, opts))
+        return self._with_compliant(self._compliant.sort(named_irs, opts))
 
     def drop(self, columns: Sequence[str], *, strict: bool = True) -> Self:
-        return self._from_compliant(self._compliant.drop(columns, strict=strict))
+        return self._with_compliant(self._compliant.drop(columns, strict=strict))
 
     def drop_nulls(self, subset: str | Sequence[str] | None = None) -> Self:
         subset = [subset] if isinstance(subset, str) else subset
-        return self._from_compliant(self._compliant.drop_nulls(subset))
+        return self._with_compliant(self._compliant.drop_nulls(subset))
 
 
-class DataFrame(BaseFrame[NativeDataFrameT], Generic[NativeDataFrameT, NativeSeriesT]):
-    _compliant: CompliantDataFrame[Any, NativeDataFrameT, NativeSeriesT]
+class DataFrame(
+    BaseFrame[NativeDataFrameT_co], Generic[NativeDataFrameT_co, NativeSeriesT]
+):
+    _compliant: CompliantDataFrame[Any, NativeDataFrameT_co, NativeSeriesT]
 
     @property
     def _series(self) -> type[Series[NativeSeriesT]]:
         return Series[NativeSeriesT]
 
-    # NOTE: Gave up on trying to get typing working for now
     @classmethod
-    def from_native(  # type: ignore[override]
-        cls, native: NativeFrame, /
-    ) -> DataFrame[pa.Table, pa.ChunkedArray[Any]]:
+    def from_native(
+        cls: type[DataFrame[Any, Any]], native: NativeDataFrameT, /
+    ) -> DataFrame[NativeDataFrameT]:
         if is_pyarrow_table(native):
             from narwhals._plan.arrow.dataframe import ArrowDataFrame
 
-            return ArrowDataFrame.from_native(native, cls._version).to_narwhals()
+            return cls(ArrowDataFrame.from_native(native, cls._version))
 
         raise NotImplementedError(type(native))
 
@@ -129,7 +126,7 @@ def to_dict(
     ) -> dict[str, Series[NativeSeriesT]] | dict[str, list[Any]]:
         if as_series:
             return {
-                key: self._series._from_compliant(value)
+                key: self._series(value)
                 for key, value in self._compliant.to_dict(as_series=as_series).items()
             }
         return self._compliant.to_dict(as_series=as_series)

narwhals/_plan/group_by.py

@@ -25,7 +25,7 @@ def __init__(self, frame: DataFrameT, grouper: Grouped, /) -> None:
 
     def agg(self, *aggs: OneOrIterable[IntoExpr], **named_aggs: IntoExpr) -> DataFrameT:
         frame = self._frame
-        return frame._from_compliant(
+        return frame._with_compliant(
             self._grouper.agg(*aggs, **named_aggs)
             .resolve(frame)
             .evaluate(frame._compliant)
@@ -35,7 +35,7 @@ def __iter__(self) -> Iterator[tuple[Any, DataFrameT]]:
         frame = self._frame
         resolver = self._grouper.agg().resolve(frame)
         for key, df in frame._compliant.group_by_resolver(resolver):
-            yield key, frame._from_compliant(df)
+            yield key, frame._with_compliant(df)
 
 
 class Grouped(Grouper["Resolved"]):

narwhals/_plan/series.py

@@ -2,23 +2,19 @@
 
 from typing import TYPE_CHECKING, Any, ClassVar, Generic
 
-from narwhals._plan.typing import NativeSeriesT
+from narwhals._plan.typing import NativeSeriesT, NativeSeriesT_co
 from narwhals._utils import Version
 from narwhals.dependencies import is_pyarrow_chunked_array
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
-    import pyarrow as pa
-    from typing_extensions import Self
-
     from narwhals._plan.compliant.series import CompliantSeries
     from narwhals.dtypes import DType
-    from narwhals.typing import NativeSeries
 
 
-class Series(Generic[NativeSeriesT]):
-    _compliant: CompliantSeries[NativeSeriesT]
+class Series(Generic[NativeSeriesT_co]):
+    _compliant: CompliantSeries[NativeSeriesT_co]
     _version: ClassVar[Version] = Version.MAIN
 
     @property
@@ -33,27 +29,21 @@ def dtype(self) -> DType:
     def name(self) -> str:
         return self._compliant.name
 
-    # NOTE: Gave up on trying to get typing working for now
+    def __init__(self, compliant: CompliantSeries[NativeSeriesT_co], /) -> None:
+        self._compliant = compliant
+
     @classmethod
     def from_native(
-        cls, native: NativeSeries, name: str = "", /
-    ) -> Series[pa.ChunkedArray[Any]]:
+        cls: type[Series[Any]], native: NativeSeriesT, name: str = "", /
+    ) -> Series[NativeSeriesT]:
         if is_pyarrow_chunked_array(native):
             from narwhals._plan.arrow.series import ArrowSeries
 
-            return ArrowSeries.from_native(
-                native, name, version=cls._version
-            ).to_narwhals()
+            return cls(ArrowSeries.from_native(native, name, version=cls._version))
 
         raise NotImplementedError(type(native))
 
-    @classmethod
-    def _from_compliant(cls, compliant: CompliantSeries[NativeSeriesT], /) -> Self:
-        obj = cls.__new__(cls)
-        obj._compliant = compliant
-        return obj
-
-    def to_native(self) -> NativeSeriesT:
+    def to_native(self) -> NativeSeriesT_co:
         return self._compliant.native
 
     def to_list(self) -> list[Any]:
@@ -63,5 +53,5 @@ def __iter__(self) -> Iterator[Any]:
         yield from self.to_native()
 
 
-class SeriesV1(Series[NativeSeriesT]):
+class SeriesV1(Series[NativeSeriesT_co]):
     _version: ClassVar[Version] = Version.V1

narwhals/_plan/typing.py

@@ -87,10 +87,22 @@
     "NonNestedLiteralT", bound="NonNestedLiteral", default="NonNestedLiteral"
 )
 NativeSeriesT = TypeVar("NativeSeriesT", bound="NativeSeries", default="NativeSeries")
+NativeSeriesT_co = TypeVar(
+    "NativeSeriesT_co", bound="NativeSeries", covariant=True, default="NativeSeries"
+)
 NativeFrameT = TypeVar("NativeFrameT", bound="NativeFrame", default="NativeFrame")
+NativeFrameT_co = TypeVar(
+    "NativeFrameT_co", bound="NativeFrame", covariant=True, default="NativeFrame"
+)
 NativeDataFrameT = TypeVar(
     "NativeDataFrameT", bound="NativeDataFrame", default="NativeDataFrame"
 )
+NativeDataFrameT_co = TypeVar(
+    "NativeDataFrameT_co",
+    bound="NativeDataFrame",
+    covariant=True,
+    default="NativeDataFrame",
+)
 LiteralT = TypeVar("LiteralT", bound="NonNestedLiteral | Series[t.Any]", default=t.Any)
 MapIR: TypeAlias = "t.Callable[[ExprIR], ExprIR]"
 """A function to apply to all nodes in this tree."""