refactor: Replace imports from typing

Author: dangotbanned

narwhals/_compliant/typing.py

@@ -23,13 +23,11 @@
     from narwhals._compliant.namespace import CompliantNamespace, EagerNamespace
     from narwhals._compliant.series import CompliantSeries, EagerSeries
     from narwhals._compliant.window import WindowInputs
+    from narwhals._native import NativeDataFrame, NativeFrame, NativeSeries
     from narwhals.typing import (
         FillNullStrategy,
         IntoLazyFrame,
         ModeKeepStrategy,
-        NativeDataFrame,
-        NativeFrame,
-        NativeSeries,
         RankMethod,
         RollingInterpolationMethod,
     )

narwhals/functions.py

@@ -41,6 +41,7 @@
     from typing_extensions import TypeAlias, TypeIs
 
     from narwhals._compliant import CompliantExpr, CompliantNamespace
+    from narwhals._native import NativeDataFrame, NativeLazyFrame, NativeSeries
     from narwhals._translate import IntoArrowTable
     from narwhals._typing import Backend, EagerAllowed, IntoBackend
     from narwhals.dataframe import DataFrame, LazyFrame
@@ -50,9 +51,6 @@
         IntoDType,
         IntoExpr,
         IntoSchema,
-        NativeDataFrame,
-        NativeLazyFrame,
-        NativeSeries,
         NonNestedLiteral,
         _1DArray,
         _2DArray,

narwhals/typing.py

@@ -8,7 +8,7 @@
 
 if TYPE_CHECKING:
     import datetime as dt
-    from collections.abc import Iterable, Sequence, Sized
+    from collections.abc import Sequence
     from decimal import Decimal
     from types import ModuleType
 
@@ -19,29 +19,17 @@
     from typing_extensions import TypeAlias
 
     from narwhals import dtypes
-    from narwhals._native import NativeIbis
+    from narwhals._native import (
+        NativeDataFrame,
+        NativeIbis,
+        NativeLazyFrame,
+        NativeSeries,
+    )
     from narwhals.dataframe import DataFrame, LazyFrame
     from narwhals.expr import Expr
     from narwhals.schema import Schema
     from narwhals.series import Series
 
-    # All dataframes supported by Narwhals have a
-    # `columns` property. Their similarities don't extend
-    # _that_ much further unfortunately...
-    class NativeFrame(Protocol):
-        @property
-        def columns(self) -> Any: ...
-
-        def join(self, *args: Any, **kwargs: Any) -> Any: ...
-
-    class NativeDataFrame(Sized, NativeFrame, Protocol): ...
-
-    class NativeLazyFrame(NativeFrame, Protocol):
-        def explain(self, *args: Any, **kwargs: Any) -> Any: ...
-
-    class NativeSeries(Sized, Iterable[Any], Protocol):
-        def filter(self, *args: Any, **kwargs: Any) -> Any: ...
-
     class SupportsNativeNamespace(Protocol):
         def __native_namespace__(self) -> ModuleType: ...
 

tests/conftest.py

@@ -15,18 +15,16 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    import duckdb
     import ibis
     import pandas as pd
     import polars as pl
     import pyarrow as pa
     from ibis.backends.duckdb import Backend as IbisDuckDBBackend
-    from pyspark.sql import DataFrame as PySparkDataFrame
     from typing_extensions import TypeAlias
 
-    from narwhals._spark_like.dataframe import SQLFrameDataFrame
+    from narwhals._native import NativeDask, NativeDuckDB, NativePySpark, NativeSQLFrame
     from narwhals._typing import EagerAllowed
-    from narwhals.typing import NativeDataFrame, NativeLazyFrame
+    from narwhals.typing import IntoDataFrame
     from tests.utils import Constructor, ConstructorEager, ConstructorLazy
 
     Data: TypeAlias = "dict[str, list[Any]]"
@@ -100,27 +98,27 @@ def pandas_pyarrow_constructor(obj: Data) -> pd.DataFrame:
     return pd.DataFrame(obj).convert_dtypes(dtype_backend="pyarrow")
 
 
-def modin_constructor(obj: Data) -> NativeDataFrame:  # pragma: no cover
+def modin_constructor(obj: Data) -> IntoDataFrame:  # pragma: no cover
     import modin.pandas as mpd
     import pandas as pd
 
     df = mpd.DataFrame(pd.DataFrame(obj))
-    return cast("NativeDataFrame", df)
+    return cast("IntoDataFrame", df)
 
 
-def modin_pyarrow_constructor(obj: Data) -> NativeDataFrame:  # pragma: no cover
+def modin_pyarrow_constructor(obj: Data) -> IntoDataFrame:  # pragma: no cover
     import modin.pandas as mpd
     import pandas as pd
 
     df = mpd.DataFrame(pd.DataFrame(obj)).convert_dtypes(dtype_backend="pyarrow")
-    return cast("NativeDataFrame", df)
+    return cast("IntoDataFrame", df)
 
 
-def cudf_constructor(obj: Data) -> NativeDataFrame:  # pragma: no cover
+def cudf_constructor(obj: Data) -> IntoDataFrame:  # pragma: no cover
     import cudf
 
     df = cudf.DataFrame(obj)
-    return cast("NativeDataFrame", df)
+    return cast("IntoDataFrame", df)
 
 
 def polars_eager_constructor(obj: Data) -> pl.DataFrame:
@@ -135,7 +133,7 @@ def polars_lazy_constructor(obj: Data) -> pl.LazyFrame:
     return pl.LazyFrame(obj)
 
 
-def duckdb_lazy_constructor(obj: Data) -> duckdb.DuckDBPyRelation:
+def duckdb_lazy_constructor(obj: Data) -> NativeDuckDB:
     import duckdb
     import polars as pl
 
@@ -145,16 +143,16 @@ def duckdb_lazy_constructor(obj: Data) -> duckdb.DuckDBPyRelation:
     return duckdb.table("_df")
 
 
-def dask_lazy_p1_constructor(obj: Data) -> NativeLazyFrame:  # pragma: no cover
+def dask_lazy_p1_constructor(obj: Data) -> NativeDask:  # pragma: no cover
     import dask.dataframe as dd
 
-    return cast("NativeLazyFrame", dd.from_dict(obj, npartitions=1))
+    return cast("NativeDask", dd.from_dict(obj, npartitions=1))
 
 
-def dask_lazy_p2_constructor(obj: Data) -> NativeLazyFrame:  # pragma: no cover
+def dask_lazy_p2_constructor(obj: Data) -> NativeDask:  # pragma: no cover
     import dask.dataframe as dd
 
-    return cast("NativeLazyFrame", dd.from_dict(obj, npartitions=2))
+    return cast("NativeDask", dd.from_dict(obj, npartitions=2))
 
 
 def pyarrow_table_constructor(obj: dict[str, Any]) -> pa.Table:
@@ -163,7 +161,7 @@ def pyarrow_table_constructor(obj: dict[str, Any]) -> pa.Table:
     return pa.table(obj)
 
 
-def pyspark_lazy_constructor() -> Callable[[Data], PySparkDataFrame]:  # pragma: no cover
+def pyspark_lazy_constructor() -> Callable[[Data], NativePySpark]:  # pragma: no cover
     pytest.importorskip("pyspark")
     import warnings
     from atexit import register
@@ -178,22 +176,22 @@ def pyspark_lazy_constructor() -> Callable[[Data], PySparkDataFrame]:  # pragma:
 
         register(session.stop)
 
-        def _constructor(obj: Data) -> PySparkDataFrame:
+        def _constructor(obj: Data) -> NativePySpark:
             _obj = deepcopy(obj)
             index_col_name = generate_temporary_column_name(n_bytes=8, columns=list(_obj))
             _obj[index_col_name] = list(range(len(_obj[next(iter(_obj))])))
-
-            return (
+            result = (
                 session.createDataFrame([*zip(*_obj.values())], schema=[*_obj.keys()])
                 .repartition(2)
                 .orderBy(index_col_name)
                 .drop(index_col_name)
             )
+            return cast("NativePySpark", result)
 
         return _constructor
 
 
-def sqlframe_pyspark_lazy_constructor(obj: Data) -> SQLFrameDataFrame:  # pragma: no cover
+def sqlframe_pyspark_lazy_constructor(obj: Data) -> NativeSQLFrame:  # pragma: no cover
     session = sqlframe_session()
     return session.createDataFrame([*zip(*obj.values())], schema=[*obj.keys()])
 

tests/expr_and_series/cast_test.py

@@ -17,7 +17,7 @@
 )
 
 if TYPE_CHECKING:
-    from narwhals.typing import NativeLazyFrame
+    from narwhals._native import NativeSQLFrame
 
 DATA = {
     "a": [1],
@@ -282,12 +282,12 @@ def test_cast_struct(request: pytest.FixtureRequest, constructor: Constructor) -
     if "spark" in str(constructor):  # pragma: no cover
         # Special handling for pyspark as it natively maps the input to
         # a column of type MAP<STRING, STRING>
-        native_ldf = cast("NativeLazyFrame", native_df)
+        native_ldf = cast("NativeSQLFrame", native_df)
         _tmp_nw_compliant_frame = nw.from_native(native_ldf)._compliant_frame
         F = _tmp_nw_compliant_frame._F  # type: ignore[attr-defined]
         T = _tmp_nw_compliant_frame._native_dtypes  # type: ignore[attr-defined] # noqa: N806
 
-        native_ldf = native_ldf.withColumn(  # type: ignore[attr-defined]
+        native_ldf = native_ldf.withColumn(
             "a",
             F.struct(
                 F.col("a.movie ").cast(T.StringType()).alias("movie "),

tests/frame/join_test.py

@@ -17,11 +17,11 @@
 )
 
 if TYPE_CHECKING:
-    from narwhals.typing import IntoLazyFrameT, JoinStrategy, NativeDataFrame
+    from narwhals.typing import IntoDataFrame, IntoLazyFrameT, JoinStrategy
 
 
 def from_native_lazy(
-    native: IntoLazyFrameT | NativeDataFrame,
+    native: IntoLazyFrameT | IntoDataFrame,
 ) -> nw.LazyFrame[IntoLazyFrameT] | nw.LazyFrame[Any]:
     """Every join test [needs to use `.lazy()` for typing]*.
 

tests/utils.py

@@ -22,7 +22,8 @@
     from sqlframe.duckdb import DuckDBSession
     from typing_extensions import TypeAlias
 
-    from narwhals.typing import Frame, NativeDataFrame, NativeLazyFrame
+    from narwhals._native import NativeLazyFrame
+    from narwhals.typing import Frame, IntoDataFrame
 
 
 def get_module_version_as_tuple(module_name: str) -> tuple[int, ...]:
@@ -42,8 +43,8 @@ def get_module_version_as_tuple(module_name: str) -> tuple[int, ...]:
 PYSPARK_VERSION: tuple[int, ...] = get_module_version_as_tuple("pyspark")
 CUDF_VERSION: tuple[int, ...] = get_module_version_as_tuple("cudf")
 
-Constructor: TypeAlias = Callable[[Any], "NativeLazyFrame | NativeDataFrame"]
-ConstructorEager: TypeAlias = Callable[[Any], "NativeDataFrame"]
+Constructor: TypeAlias = Callable[[Any], "NativeLazyFrame | IntoDataFrame"]
+ConstructorEager: TypeAlias = Callable[[Any], "IntoDataFrame"]
 ConstructorLazy: TypeAlias = Callable[[Any], "NativeLazyFrame"]
 ConstructorPandasLike: TypeAlias = Callable[[Any], "pd.DataFrame"]