"""Narwhals-level equivalent of `CompliantNamespace`.""" from __future__ import annotations from typing import ( TYPE_CHECKING, Any, Callable, Generic, Literal, Protocol, TypeVar, 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, ) if TYPE_CHECKING: from collections.abc import Collection, Sized from types import ModuleType from typing import ClassVar import duckdb import pandas as pd import polars as pl import pyarrow as pa import pyspark.sql as pyspark_sql from pyspark.sql.connect.dataframe import DataFrame as PySparkConnectDataFrame from typing_extensions import Self, TypeAlias, TypeIs 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 DataFrameLike, NativeFrame, NativeLazyFrame, NativeSeries T = TypeVar("T") _Guard: TypeAlias = "Callable[[Any], TypeIs[T]]" _Polars: TypeAlias = Literal["polars"] _Arrow: TypeAlias = Literal["pyarrow"] _Dask: TypeAlias = Literal["dask"] _DuckDB: TypeAlias = Literal["duckdb"] _PandasLike: TypeAlias = Literal["pandas", "cudf", "modin"] _Ibis: TypeAlias = Literal["ibis"] _SparkLike: TypeAlias = Literal["pyspark", "sqlframe", "pyspark[connect]"] _EagerOnly: TypeAlias = "_PandasLike | _Arrow" _EagerAllowed: TypeAlias = "_Polars | _EagerOnly" _LazyOnly: TypeAlias = "_SparkLike | _Dask | _DuckDB | _Ibis" _LazyAllowed: TypeAlias = "_Polars | _LazyOnly" Polars: TypeAlias = Literal[_Polars, Implementation.POLARS] Arrow: TypeAlias = Literal[_Arrow, Implementation.PYARROW] Dask: TypeAlias = Literal[_Dask, Implementation.DASK] DuckDB: TypeAlias = Literal[_DuckDB, Implementation.DUCKDB] Ibis: TypeAlias = Literal[_Ibis, Implementation.IBIS] PandasLike: TypeAlias = Literal[ _PandasLike, Implementation.PANDAS, Implementation.CUDF, Implementation.MODIN ] SparkLike: TypeAlias = Literal[ _SparkLike, Implementation.PYSPARK, Implementation.SQLFRAME, Implementation.PYSPARK_CONNECT, ] EagerOnly: TypeAlias = "PandasLike | Arrow" EagerAllowed: TypeAlias = "EagerOnly | Polars" LazyOnly: TypeAlias = "SparkLike | Dask | DuckDB | Ibis" LazyAllowed: TypeAlias = "LazyOnly | Polars" BackendName: TypeAlias = "_EagerAllowed | _LazyAllowed" IntoBackend: TypeAlias = "BackendName | Implementation | ModuleType" EagerAllowedNamespace: TypeAlias = "Namespace[PandasLikeNamespace] | Namespace[ArrowNamespace] | Namespace[PolarsNamespace]" EagerAllowedImplementation: TypeAlias = Literal[ Implementation.PANDAS, Implementation.CUDF, Implementation.MODIN, Implementation.PYARROW, Implementation.POLARS, ] 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, inplace: Literal[False], **kwds: Any) -> Self: """`inplace=False` is required to avoid (incorrect?) default overloads.""" ... class _BasePandasLikeFrame(NativeFrame, _BasePandasLike, Protocol): ... class _BasePandasLikeSeries(NativeSeries, _BasePandasLike, Protocol): def where(self, cond: Any, other: Any = ..., **kwds: Any) -> Any: ... class _NativeDask(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]] _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 = "pyspark_sql.DataFrame" _NativePySparkConnect: TypeAlias = "PySparkConnectDataFrame" _NativeSparkLike: TypeAlias = ( "_NativeSQLFrame | _NativePySpark | _NativePySparkConnect" ) NativeKnown: TypeAlias = "_NativePolars | _NativeArrow | _NativePandasLike | _NativeSparkLike | _NativeDuckDB | _NativeDask | _NativeIbis" NativeUnknown: TypeAlias = ( "NativeFrame | NativeSeries | NativeLazyFrame | DataFrameLike" ) NativeAny: TypeAlias = "NativeKnown | NativeUnknown" __all__ = ["Namespace"] class Namespace(Generic[CompliantNamespaceT_co]): _compliant_namespace: CompliantNamespaceT_co _version: ClassVar[Version] = Version.MAIN def __init__(self, namespace: CompliantNamespaceT_co, /) -> None: self._compliant_namespace = namespace def __init_subclass__(cls, *args: Any, version: Version, **kwds: Any) -> None: super().__init_subclass__(*args, **kwds) if isinstance(version, Version): cls._version = version else: msg = f"Expected {Version} but got {type(version).__name__!r}" raise TypeError(msg) def __repr__(self) -> str: return f"Namespace[{type(self.compliant).__name__}]" @property def compliant(self) -> CompliantNamespaceT_co: return self._compliant_namespace @property def implementation(self) -> Implementation: return self.compliant._implementation @property def version(self) -> Version: return self._version @overload @classmethod def from_backend(cls, backend: PandasLike, /) -> Namespace[PandasLikeNamespace]: ... @overload @classmethod def from_backend(cls, backend: Polars, /) -> Namespace[PolarsNamespace]: ... @overload @classmethod def from_backend(cls, backend: Arrow, /) -> Namespace[ArrowNamespace]: ... @overload @classmethod def from_backend(cls, backend: SparkLike, /) -> Namespace[SparkLikeNamespace]: ... @overload @classmethod def from_backend(cls, backend: DuckDB, /) -> Namespace[DuckDBNamespace]: ... @overload @classmethod def from_backend(cls, backend: Dask, /) -> Namespace[DaskNamespace]: ... @overload @classmethod def from_backend(cls, backend: Ibis, /) -> Namespace[IbisNamespace]: ... @overload @classmethod def from_backend(cls, backend: EagerAllowed, /) -> EagerAllowedNamespace: ... @overload @classmethod def from_backend( cls, backend: IntoBackend, / ) -> Namespace[CompliantNamespaceAny]: ... @classmethod def from_backend( cls: type[Namespace[Any]], backend: IntoBackend, / ) -> Namespace[Any]: """Instantiate from native namespace module, string, or Implementation. Arguments: backend: native namespace module, string, or Implementation. Returns: Namespace. Examples: >>> from narwhals._namespace import Namespace >>> Namespace.from_backend("polars") Namespace[PolarsNamespace] """ impl = Implementation.from_backend(backend) backend_version = impl._backend_version() # noqa: F841 version = cls._version ns: CompliantNamespaceAny if impl.is_pandas_like(): from narwhals._pandas_like.namespace import PandasLikeNamespace ns = PandasLikeNamespace(implementation=impl, version=version) elif impl.is_polars(): from narwhals._polars.namespace import PolarsNamespace ns = PolarsNamespace(version=version) elif impl.is_pyarrow(): from narwhals._arrow.namespace import ArrowNamespace ns = ArrowNamespace(version=version) elif impl.is_spark_like(): from narwhals._spark_like.namespace import SparkLikeNamespace ns = SparkLikeNamespace(implementation=impl, version=version) elif impl.is_duckdb(): from narwhals._duckdb.namespace import DuckDBNamespace ns = DuckDBNamespace(version=version) elif impl.is_dask(): from narwhals._dask.namespace import DaskNamespace ns = DaskNamespace(version=version) elif impl.is_ibis(): from narwhals._ibis.namespace import IbisNamespace ns = IbisNamespace(version=version) else: msg = "Not supported Implementation" # pragma: no cover raise AssertionError(msg) return cls(ns) @overload @classmethod def from_native_object( cls, native: _NativePolars, / ) -> Namespace[PolarsNamespace]: ... @overload @classmethod def from_native_object( cls, native: _NativePandas, / ) -> Namespace[PandasLikeNamespace[pd.DataFrame, pd.Series[Any]]]: ... @overload @classmethod def from_native_object(cls, native: _NativeArrow, /) -> Namespace[ArrowNamespace]: ... @overload @classmethod def from_native_object( cls, native: _NativeSparkLike, / ) -> Namespace[SparkLikeNamespace]: ... @overload @classmethod def from_native_object( cls, native: _NativeDuckDB, / ) -> Namespace[DuckDBNamespace]: ... @overload @classmethod def from_native_object(cls, native: _NativeDask, /) -> Namespace[DaskNamespace]: ... @overload @classmethod def from_native_object(cls, native: _NativeIbis, /) -> Namespace[IbisNamespace]: ... @overload @classmethod def from_native_object( cls, native: _NativeModin, / ) -> Namespace[PandasLikeNamespace[_ModinDataFrame, _ModinSeries]]: ... @overload @classmethod def from_native_object( cls, native: _NativeCuDF, / ) -> Namespace[PandasLikeNamespace[_CuDFDataFrame, _CuDFSeries]]: ... @overload @classmethod def from_native_object( cls, native: _NativePandasLike, / ) -> Namespace[PandasLikeNamespace[Any, Any]]: ... @overload @classmethod def from_native_object( cls, native: NativeUnknown, / ) -> Namespace[CompliantNamespaceAny]: ... @classmethod def from_native_object( # noqa: PLR0911 cls: type[Namespace[Any]], native: NativeAny, / ) -> Namespace[Any]: if is_native_polars(native): return cls.from_backend(Implementation.POLARS) elif is_native_pandas(native): return cls.from_backend(Implementation.PANDAS) elif is_native_arrow(native): return cls.from_backend(Implementation.PYARROW) elif is_native_spark_like(native): return cls.from_backend( Implementation.SQLFRAME if is_native_sqlframe(native) else Implementation.PYSPARK_CONNECT if is_native_pyspark_connect(native) else Implementation.PYSPARK ) elif is_native_dask(native): return cls.from_backend(Implementation.DASK) # pragma: no cover elif is_native_duckdb(native): return cls.from_backend(Implementation.DUCKDB) elif is_native_cudf(native): # pragma: no cover return cls.from_backend(Implementation.CUDF) elif is_native_modin(native): # pragma: no cover return cls.from_backend(Implementation.MODIN) elif is_native_ibis(native): return cls.from_backend(Implementation.IBIS) else: msg = f"Unsupported type: {type(native).__qualname__!r}" raise TypeError(msg) 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: _Guard[_NativePySpark] = is_pyspark_dataframe is_native_pyspark_connect: _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)