refactor(typing): Add _native.py (#3086)

Co-authored-by: Francesco Bruzzesi <[email protected]>

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/_namespace.py

@@ -2,50 +2,50 @@
 
 from __future__ import annotations
 
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Generic,
-    Protocol,
-    TypeVar,
-    cast,
-    overload,
-)
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, overload
 
 from narwhals._compliant.typing import CompliantNamespaceAny, CompliantNamespaceT_co
-from narwhals._utils import Implementation, Version
-from narwhals.dependencies import (
-    get_cudf,
-    get_modin,
-    get_pandas,
-    get_polars,
-    get_pyarrow,
-    is_dask_dataframe,
-    is_duckdb_relation,
-    is_ibis_table,
-    is_pyspark_connect_dataframe,
-    is_pyspark_dataframe,
-    is_sqlframe_dataframe,
+from narwhals._native import (
+    NativeAny,
+    NativeArrow,
+    NativeCuDF,
+    NativeDask,
+    NativeDuckDB,
+    NativeIbis,
+    NativeModin,
+    NativePandas,
+    NativePandasLike,
+    NativePolars,
+    NativeSparkLike,
+    NativeUnknown,
+    _CuDFDataFrame,
+    _CuDFSeries,
+    _ModinDataFrame,
+    _ModinSeries,
+    is_native_arrow,
+    is_native_cudf,
+    is_native_dask,
+    is_native_duckdb,
+    is_native_ibis,
+    is_native_modin,
+    is_native_pandas,
+    is_native_polars,
+    is_native_pyspark_connect,
+    is_native_spark_like,
+    is_native_sqlframe,
 )
+from narwhals._utils import Implementation, Version
 
 if TYPE_CHECKING:
-    from collections.abc import Collection, Sized
-    from typing import ClassVar
-
-    import duckdb
     import pandas as pd
-    import polars as pl
-    import pyarrow as pa
-    from typing_extensions import Self, TypeAlias, TypeIs
+    from typing_extensions import TypeAlias
 
     from narwhals._arrow.namespace import ArrowNamespace
     from narwhals._dask.namespace import DaskNamespace
     from narwhals._duckdb.namespace import DuckDBNamespace
     from narwhals._ibis.namespace import IbisNamespace
     from narwhals._pandas_like.namespace import PandasLikeNamespace
     from narwhals._polars.namespace import PolarsNamespace
-    from narwhals._spark_like.dataframe import SQLFrameDataFrame
     from narwhals._spark_like.namespace import SparkLikeNamespace
     from narwhals._typing import (
         Arrow,
@@ -59,87 +59,8 @@
         Polars,
         SparkLike,
     )
-    from narwhals.typing import NativeDataFrame, NativeLazyFrame, NativeSeries
-
-    T = TypeVar("T")
-
-    _Guard: TypeAlias = "Callable[[Any], TypeIs[T]]"
 
     EagerAllowedNamespace: TypeAlias = "Namespace[PandasLikeNamespace] | Namespace[ArrowNamespace] | Namespace[PolarsNamespace]"
-    Incomplete: TypeAlias = Any
-
-    class _BasePandasLike(Sized, Protocol):
-        index: Any
-        """`mypy` doesn't like the asymmetric `property` setter in `pandas`."""
-
-        def __getitem__(self, key: Any, /) -> Any: ...
-        def __mul__(self, other: float | Collection[float] | Self, /) -> Self: ...
-        def __floordiv__(self, other: float | Collection[float] | Self, /) -> Self: ...
-        @property
-        def loc(self) -> Any: ...
-        @property
-        def shape(self) -> tuple[int, ...]: ...
-        def set_axis(self, labels: Any, *, axis: Any = ..., copy: bool = ...) -> Self: ...
-        def copy(self, deep: bool = ...) -> Self: ...  # noqa: FBT001
-        def rename(self, *args: Any, **kwds: Any) -> Self | Incomplete:
-            """`mypy` & `pyright` disagree on overloads.
-
-            `Incomplete` used to fix [more important issue](https://github.com/narwhals-dev/narwhals/pull/3016#discussion_r2296139744).
-            """
-
-    class _BasePandasLikeFrame(NativeDataFrame, _BasePandasLike, Protocol): ...
-
-    class _BasePandasLikeSeries(NativeSeries, _BasePandasLike, Protocol):
-        def where(self, cond: Any, other: Any = ..., /) -> Self | Incomplete: ...
-
-    class _NativeDask(NativeLazyFrame, Protocol):
-        _partition_type: type[pd.DataFrame]
-
-    class _CuDFDataFrame(_BasePandasLikeFrame, Protocol):
-        def to_pylibcudf(self, *args: Any, **kwds: Any) -> Any: ...
-
-    class _CuDFSeries(_BasePandasLikeSeries, Protocol):
-        def to_pylibcudf(self, *args: Any, **kwds: Any) -> Any: ...
-
-    class _NativeIbis(Protocol):
-        def sql(self, *args: Any, **kwds: Any) -> Any: ...
-        def __pyarrow_result__(self, *args: Any, **kwds: Any) -> Any: ...
-        def __pandas_result__(self, *args: Any, **kwds: Any) -> Any: ...
-        def __polars_result__(self, *args: Any, **kwds: Any) -> Any: ...
-
-    class _ModinDataFrame(_BasePandasLikeFrame, Protocol):
-        _pandas_class: type[pd.DataFrame]
-
-    class _ModinSeries(_BasePandasLikeSeries, Protocol):
-        _pandas_class: type[pd.Series[Any]]
-
-    # NOTE: Using `pyspark.sql.DataFrame` creates false positives in overloads when not installed
-    class _PySparkDataFrame(NativeLazyFrame, Protocol):
-        # Arbitrary method that `sqlframe` doesn't have and unlikely to appear anywhere else
-        # https://github.com/apache/spark/blob/8530444e25b83971da4314c608aa7d763adeceb3/python/pyspark/sql/dataframe.py#L4875
-        def dropDuplicatesWithinWatermark(self, *arg: Any, **kwargs: Any) -> Any: ...  # noqa: N802
-
-    _NativePolars: TypeAlias = "pl.DataFrame | pl.LazyFrame | pl.Series"
-    _NativeArrow: TypeAlias = "pa.Table | pa.ChunkedArray[Any]"
-    _NativeDuckDB: TypeAlias = "duckdb.DuckDBPyRelation"
-    _NativePandas: TypeAlias = "pd.DataFrame | pd.Series[Any]"
-    _NativeModin: TypeAlias = "_ModinDataFrame | _ModinSeries"
-    _NativeCuDF: TypeAlias = "_CuDFDataFrame | _CuDFSeries"
-    _NativePandasLikeSeries: TypeAlias = "pd.Series[Any] | _CuDFSeries | _ModinSeries"
-    _NativePandasLikeDataFrame: TypeAlias = (
-        "pd.DataFrame | _CuDFDataFrame | _ModinDataFrame"
-    )
-    _NativePandasLike: TypeAlias = "_NativePandasLikeDataFrame |_NativePandasLikeSeries"
-    _NativeSQLFrame: TypeAlias = "SQLFrameDataFrame"
-    _NativePySpark: TypeAlias = _PySparkDataFrame
-    _NativePySparkConnect: TypeAlias = _PySparkDataFrame
-    _NativeSparkLike: TypeAlias = (
-        "_NativeSQLFrame | _NativePySpark | _NativePySparkConnect"
-    )
-
-    NativeKnown: TypeAlias = "_NativePolars | _NativeArrow | _NativePandasLike | _NativeSparkLike | _NativeDuckDB | _NativeDask | _NativeIbis"
-    NativeUnknown: TypeAlias = "NativeDataFrame | NativeSeries | NativeLazyFrame"
-    NativeAny: TypeAlias = "NativeKnown | NativeUnknown"
 
 __all__ = ["Namespace"]
 
@@ -268,55 +189,55 @@ def from_backend(
     @overload
     @classmethod
     def from_native_object(
-        cls, native: _NativePolars, /
+        cls, native: NativePolars, /
     ) -> Namespace[PolarsNamespace]: ...
 
     @overload
     @classmethod
     def from_native_object(
-        cls, native: _NativePandas, /
+        cls, native: NativePandas, /
     ) -> Namespace[PandasLikeNamespace[pd.DataFrame, pd.Series[Any]]]: ...
 
     @overload
     @classmethod
-    def from_native_object(cls, native: _NativeArrow, /) -> Namespace[ArrowNamespace]: ...
+    def from_native_object(cls, native: NativeArrow, /) -> Namespace[ArrowNamespace]: ...
 
     @overload
     @classmethod
     def from_native_object(
-        cls, native: _NativeSparkLike, /
+        cls, native: NativeSparkLike, /
     ) -> Namespace[SparkLikeNamespace]: ...
 
     @overload
     @classmethod
     def from_native_object(
-        cls, native: _NativeDuckDB, /
+        cls, native: NativeDuckDB, /
     ) -> Namespace[DuckDBNamespace]: ...
 
     @overload
     @classmethod
-    def from_native_object(cls, native: _NativeDask, /) -> Namespace[DaskNamespace]: ...
+    def from_native_object(cls, native: NativeDask, /) -> Namespace[DaskNamespace]: ...
 
     @overload
     @classmethod
-    def from_native_object(cls, native: _NativeIbis, /) -> Namespace[IbisNamespace]: ...
+    def from_native_object(cls, native: NativeIbis, /) -> Namespace[IbisNamespace]: ...
 
     @overload
     @classmethod
     def from_native_object(
-        cls, native: _NativeModin, /
+        cls, native: NativeModin, /
     ) -> Namespace[PandasLikeNamespace[_ModinDataFrame, _ModinSeries]]: ...
 
     @overload
     @classmethod
     def from_native_object(
-        cls, native: _NativeCuDF, /
+        cls, native: NativeCuDF, /
     ) -> Namespace[PandasLikeNamespace[_CuDFDataFrame, _CuDFSeries]]: ...
 
     @overload
     @classmethod
     def from_native_object(
-        cls, native: _NativePandasLike, /
+        cls, native: NativePandasLike, /
     ) -> Namespace[PandasLikeNamespace[Any, Any]]: ...
 
     @overload
@@ -358,61 +279,3 @@ def from_native_object(
             msg = f"Unsupported type: {type(native).__qualname__!r}"
             raise TypeError(msg)
         return cls.from_backend(impl)
-
-
-def is_native_polars(obj: Any) -> TypeIs[_NativePolars]:
-    return (pl := get_polars()) is not None and isinstance(
-        obj, (pl.DataFrame, pl.Series, pl.LazyFrame)
-    )
-
-
-def is_native_arrow(obj: Any) -> TypeIs[_NativeArrow]:
-    return (pa := get_pyarrow()) is not None and isinstance(
-        obj, (pa.Table, pa.ChunkedArray)
-    )
-
-
-def is_native_dask(obj: Any) -> TypeIs[_NativeDask]:
-    return is_dask_dataframe(obj)
-
-
-is_native_duckdb: _Guard[_NativeDuckDB] = is_duckdb_relation
-is_native_sqlframe: _Guard[_NativeSQLFrame] = is_sqlframe_dataframe
-is_native_pyspark = cast("_Guard[_NativePySpark]", is_pyspark_dataframe)
-is_native_pyspark_connect = cast(
-    "_Guard[_NativePySparkConnect]", is_pyspark_connect_dataframe
-)
-
-
-def is_native_pandas(obj: Any) -> TypeIs[_NativePandas]:
-    return (pd := get_pandas()) is not None and isinstance(obj, (pd.DataFrame, pd.Series))
-
-
-def is_native_modin(obj: Any) -> TypeIs[_NativeModin]:
-    return (mpd := get_modin()) is not None and isinstance(
-        obj, (mpd.DataFrame, mpd.Series)
-    )  # pragma: no cover
-
-
-def is_native_cudf(obj: Any) -> TypeIs[_NativeCuDF]:
-    return (cudf := get_cudf()) is not None and isinstance(
-        obj, (cudf.DataFrame, cudf.Series)
-    )  # pragma: no cover
-
-
-def is_native_pandas_like(obj: Any) -> TypeIs[_NativePandasLike]:
-    return (
-        is_native_pandas(obj) or is_native_cudf(obj) or is_native_modin(obj)
-    )  # pragma: no cover
-
-
-def is_native_spark_like(obj: Any) -> TypeIs[_NativeSparkLike]:
-    return (
-        is_native_sqlframe(obj)
-        or is_native_pyspark(obj)
-        or is_native_pyspark_connect(obj)
-    )
-
-
-def is_native_ibis(obj: Any) -> TypeIs[_NativeIbis]:
-    return is_ibis_table(obj)