feat(typing): DataFrame.from_native series typing!?!?!?!

Very related to #3086 (comment)

Author: dangotbanned

narwhals/_plan/arrow/typing.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Literal, Protocol, overload
 
 from narwhals._typing_compat import TypeVar
@@ -23,10 +23,21 @@
     )
     from typing_extensions import TypeAlias
 
+    from narwhals.typing import NativeDataFrame, NativeSeries
+
     StringScalar: TypeAlias = "Scalar[StringType | LargeStringType]"
     IntegerType: TypeAlias = "Int8Type | Int16Type | Int32Type | Int64Type | Uint8Type | Uint16Type | Uint32Type | Uint64Type"
     IntegerScalar: TypeAlias = "Scalar[IntegerType]"
 
+    class NativeArrowSeries(NativeSeries, Protocol):
+        @property
+        def chunks(self) -> list[Any]: ...
+
+    class NativeArrowDataFrame(NativeDataFrame, Protocol):
+        def column(self, *args: Any, **kwds: Any) -> NativeArrowSeries: ...
+        @property
+        def columns(self) -> Sequence[NativeArrowSeries]: ...
+
 
 ScalarT = TypeVar("ScalarT", bound="pa.Scalar[Any]", default="pa.Scalar[Any]")
 ScalarPT_contra = TypeVar(

narwhals/_plan/dataframe.py

@@ -23,8 +23,10 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+    import pyarrow as pa
     from typing_extensions import Self
 
+    from narwhals._plan.arrow.typing import NativeArrowDataFrame
     from narwhals._plan.compliant.dataframe import CompliantDataFrame, CompliantFrame
 
 
@@ -47,6 +49,7 @@ def columns(self) -> list[str]:
     def __repr__(self) -> str:  # pragma: no cover
         return generate_repr(f"nw.{type(self).__name__}", self.to_native().__repr__())
 
+    # TODO @dangotbanned: Can this be typed?
     def __init__(self, compliant: Any, /) -> None:
         self._compliant = compliant
 
@@ -100,10 +103,20 @@ class DataFrame(
     def _series(self) -> type[Series[NativeSeriesT]]:
         return Series[NativeSeriesT]
 
+    @overload
+    @classmethod
+    def from_native(
+        cls: type[DataFrame[Any, Any]], native: NativeArrowDataFrame, /
+    ) -> DataFrame[pa.Table, pa.ChunkedArray[Any]]: ...
+    @overload
+    @classmethod
+    def from_native(
+        cls: type[DataFrame[Any, Any]], native: NativeDataFrameT, /
+    ) -> DataFrame[NativeDataFrameT]: ...
     @classmethod
     def from_native(
         cls: type[DataFrame[Any, Any]], native: NativeDataFrameT, /
-    ) -> DataFrame[NativeDataFrameT]:
+    ) -> DataFrame[Any, Any]:
         if is_pyarrow_table(native):
             from narwhals._plan.arrow.dataframe import ArrowDataFrame
 

tests/plan/compliant_test.py

@@ -540,6 +540,7 @@ def test_row_is_py_literal(
 
 
 if TYPE_CHECKING:
+    from typing_extensions import assert_type
 
     def test_protocol_expr() -> None:
         """Static test for all members implemented.
@@ -554,3 +555,14 @@ def test_protocol_expr() -> None:
         scalar = ArrowScalar()
         assert expr
         assert scalar
+
+    def test_dataframe_from_native_overloads() -> None:
+        """Ensure we can reveal the `NativeSeries` **without** a dependency."""
+        data: dict[str, Any] = {}
+        native_good = pa.table(data)
+        result_good = nwp.DataFrame.from_native(native_good)
+        assert_type(result_good, "nwp.DataFrame[pa.Table, pa.ChunkedArray[Any]]")
+
+        native_bad = native_good.to_batches()[0]
+        nwp.DataFrame.from_native(native_bad)  # type: ignore[call-overload]
+        assert_type(native_bad, "pa.RecordBatch")