path: root/venv/lib/python3.8/site-packages/narwhals/dependencies.py

# pandas / Polars / etc. : if a user passes a dataframe from one of these
# libraries, it means they must already have imported the given module.
# So, we can just check sys.modules.
from __future__ import annotations

import sys
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import cudf
    import dask.dataframe as dd
    import duckdb
    import ibis
    import modin.pandas as mpd
    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 TypeGuard, TypeIs

    from narwhals._spark_like.dataframe import SQLFrameDataFrame
    from narwhals.dataframe import DataFrame, LazyFrame
    from narwhals.series import Series
    from narwhals.typing import (
        FrameT,
        IntoDataFrameT,
        IntoSeriesT,
        _1DArray,
        _1DArrayInt,
        _2DArray,
        _NDArray,
        _NumpyScalar,
        _ShapeT,
    )


# We silently allow these but - given that they claim
# to be drop-in replacements for pandas - testing is
# their responsibility.
IMPORT_HOOKS = frozenset(["fireducks"])


def get_polars() -> Any:
    """Get Polars module (if already imported - else return None)."""
    return sys.modules.get("polars", None)


def get_pandas() -> Any:
    """Get pandas module (if already imported - else return None)."""
    return sys.modules.get("pandas", None)


def get_modin() -> Any:  # pragma: no cover
    """Get modin.pandas module (if already imported - else return None)."""
    if (modin := sys.modules.get("modin", None)) is not None:
        return modin.pandas
    return None


def get_cudf() -> Any:
    """Get cudf module (if already imported - else return None)."""
    return sys.modules.get("cudf", None)


def get_cupy() -> Any:
    """Get cupy module (if already imported - else return None)."""
    return sys.modules.get("cupy", None)


def get_pyarrow() -> Any:  # pragma: no cover
    """Get pyarrow module (if already imported - else return None)."""
    return sys.modules.get("pyarrow", None)


def get_numpy() -> Any:
    """Get numpy module (if already imported - else return None)."""
    return sys.modules.get("numpy", None)


def get_dask() -> Any:
    """Get dask (if already imported - else return None)."""
    return sys.modules.get("dask", None)


def get_dask_dataframe() -> Any:
    """Get dask.dataframe module (if already imported - else return None)."""
    return sys.modules.get("dask.dataframe", None)


def get_duckdb() -> Any:
    """Get duckdb module (if already imported - else return None)."""
    return sys.modules.get("duckdb", None)


def get_ibis() -> Any:
    """Get ibis module (if already imported - else return None)."""
    return sys.modules.get("ibis", None)


def get_dask_expr() -> Any:  # pragma: no cover
    """Get dask_expr module (if already imported - else return None)."""
    if (dd := get_dask_dataframe()) is not None and hasattr(dd, "dask_expr"):
        return dd.dask_expr
    return sys.modules.get("dask_expr", None)


def get_pyspark() -> Any:  # pragma: no cover
    """Get pyspark module (if already imported - else return None)."""
    return sys.modules.get("pyspark", None)


def get_pyspark_sql() -> Any:
    """Get pyspark.sql module (if already imported - else return None)."""
    return sys.modules.get("pyspark.sql", None)


def get_pyspark_connect() -> Any:
    """Get pyspark.sql.connect module (if already imported - else return None)."""
    return sys.modules.get("pyspark.sql.connect", None)


def get_sqlframe() -> Any:
    """Get sqlframe module (if already imported - else return None)."""
    return sys.modules.get("sqlframe", None)


def is_pandas_dataframe(df: Any) -> TypeIs[pd.DataFrame]:
    """Check whether `df` is a pandas DataFrame without importing pandas."""
    return ((pd := get_pandas()) is not None and isinstance(df, pd.DataFrame)) or any(
        (mod := sys.modules.get(module_name, None)) is not None
        and isinstance(df, mod.pandas.DataFrame)
        for module_name in IMPORT_HOOKS
    )


def is_pandas_series(ser: Any) -> TypeIs[pd.Series[Any]]:
    """Check whether `ser` is a pandas Series without importing pandas."""
    return ((pd := get_pandas()) is not None and isinstance(ser, pd.Series)) or any(
        (mod := sys.modules.get(module_name, None)) is not None
        and isinstance(ser, mod.pandas.Series)
        for module_name in IMPORT_HOOKS
    )


def is_pandas_index(index: Any) -> TypeIs[pd.Index[Any]]:
    """Check whether `index` is a pandas Index without importing pandas."""
    return ((pd := get_pandas()) is not None and isinstance(index, pd.Index)) or any(
        (mod := sys.modules.get(module_name, None)) is not None
        and isinstance(index, mod.pandas.Index)
        for module_name in IMPORT_HOOKS
    )


def is_modin_dataframe(df: Any) -> TypeIs[mpd.DataFrame]:
    """Check whether `df` is a modin DataFrame without importing modin."""
    return (mpd := get_modin()) is not None and isinstance(df, mpd.DataFrame)


def is_modin_series(ser: Any) -> TypeIs[mpd.Series]:
    """Check whether `ser` is a modin Series without importing modin."""
    return (mpd := get_modin()) is not None and isinstance(ser, mpd.Series)


def is_modin_index(index: Any) -> TypeIs[mpd.Index[Any]]:  # pragma: no cover
    """Check whether `index` is a modin Index without importing modin."""
    return (mpd := get_modin()) is not None and isinstance(index, mpd.Index)


def is_cudf_dataframe(df: Any) -> TypeIs[cudf.DataFrame]:
    """Check whether `df` is a cudf DataFrame without importing cudf."""
    return (cudf := get_cudf()) is not None and isinstance(df, cudf.DataFrame)


def is_cudf_series(ser: Any) -> TypeIs[cudf.Series[Any]]:
    """Check whether `ser` is a cudf Series without importing cudf."""
    return (cudf := get_cudf()) is not None and isinstance(ser, cudf.Series)


def is_cudf_index(index: Any) -> TypeIs[cudf.Index]:
    """Check whether `index` is a cudf Index without importing cudf."""
    return (cudf := get_cudf()) is not None and isinstance(
        index, cudf.Index
    )  # pragma: no cover


def is_cupy_scalar(obj: Any) -> bool:
    return (
        (cupy := get_cupy()) is not None
        and isinstance(obj, cupy.ndarray)
        and obj.size == 1
    )  # pragma: no cover


def is_dask_dataframe(df: Any) -> TypeIs[dd.DataFrame]:
    """Check whether `df` is a Dask DataFrame without importing Dask."""
    return (dd := get_dask_dataframe()) is not None and isinstance(df, dd.DataFrame)


def is_duckdb_relation(df: Any) -> TypeIs[duckdb.DuckDBPyRelation]:
    """Check whether `df` is a DuckDB Relation without importing DuckDB."""
    return (duckdb := get_duckdb()) is not None and isinstance(
        df, duckdb.DuckDBPyRelation
    )


def is_ibis_table(df: Any) -> TypeIs[ibis.Table]:
    """Check whether `df` is a Ibis Table without importing Ibis."""
    return (ibis := get_ibis()) is not None and isinstance(df, ibis.expr.types.Table)


def is_polars_dataframe(df: Any) -> TypeIs[pl.DataFrame]:
    """Check whether `df` is a Polars DataFrame without importing Polars."""
    return (pl := get_polars()) is not None and isinstance(df, pl.DataFrame)


def is_polars_lazyframe(df: Any) -> TypeIs[pl.LazyFrame]:
    """Check whether `df` is a Polars LazyFrame without importing Polars."""
    return (pl := get_polars()) is not None and isinstance(df, pl.LazyFrame)


def is_polars_series(ser: Any) -> TypeIs[pl.Series]:
    """Check whether `ser` is a Polars Series without importing Polars."""
    return (pl := get_polars()) is not None and isinstance(ser, pl.Series)


def is_pyarrow_chunked_array(ser: Any) -> TypeIs[pa.ChunkedArray[Any]]:
    """Check whether `ser` is a PyArrow ChunkedArray without importing PyArrow."""
    return (pa := get_pyarrow()) is not None and isinstance(ser, pa.ChunkedArray)


def is_pyarrow_table(df: Any) -> TypeIs[pa.Table]:
    """Check whether `df` is a PyArrow Table without importing PyArrow."""
    return (pa := get_pyarrow()) is not None and isinstance(df, pa.Table)


def is_pyarrow_scalar(obj: Any) -> TypeIs[pa.Scalar[Any]]:
    return (pa := get_pyarrow()) is not None and isinstance(obj, pa.Scalar)


def is_pyspark_dataframe(df: Any) -> TypeIs[pyspark_sql.DataFrame]:
    """Check whether `df` is a PySpark DataFrame without importing PySpark."""
    return bool(
        (pyspark_sql := get_pyspark_sql()) is not None
        and isinstance(df, pyspark_sql.DataFrame)
    )


def is_pyspark_connect_dataframe(df: Any) -> TypeIs[PySparkConnectDataFrame]:
    """Check whether `df` is a PySpark Connect DataFrame without importing PySpark."""
    if get_pyspark_connect() is not None:  # pragma: no cover
        try:
            from pyspark.sql.connect.dataframe import DataFrame
        except ImportError:
            return False
        return isinstance(df, DataFrame)
    return False


def is_sqlframe_dataframe(df: Any) -> TypeIs[SQLFrameDataFrame]:
    """Check whether `df` is a SQLFrame DataFrame without importing SQLFrame."""
    if get_sqlframe() is not None:
        from sqlframe.base.dataframe import BaseDataFrame

        return isinstance(df, BaseDataFrame)
    return False  # pragma: no cover


def is_numpy_array(arr: Any | _NDArray[_ShapeT]) -> TypeIs[_NDArray[_ShapeT]]:
    """Check whether `arr` is a NumPy Array without importing NumPy."""
    return (np := get_numpy()) is not None and isinstance(arr, np.ndarray)


def is_numpy_array_1d(arr: Any) -> TypeIs[_1DArray]:
    """Check whether `arr` is a 1D NumPy Array without importing NumPy."""
    return is_numpy_array(arr) and arr.ndim == 1


def is_numpy_array_1d_int(arr: Any) -> TypeIs[_1DArrayInt]:
    return (
        (np := get_numpy())
        and is_numpy_array_1d(arr)
        and np.issubdtype(arr.dtype, np.integer)
    )


def is_numpy_array_2d(arr: Any) -> TypeIs[_2DArray]:
    """Check whether `arr` is a 2D NumPy Array without importing NumPy."""
    return is_numpy_array(arr) and arr.ndim == 2


def is_numpy_scalar(scalar: Any) -> TypeGuard[_NumpyScalar]:
    """Check whether `scalar` is a NumPy Scalar without importing NumPy."""
    # NOTE: Needs to stay as `TypeGuard`
    # - Used in `Series.__getitem__`, but not annotated
    # - `TypeGuard` is *hiding* that the check introduces an intersection
    return (np := get_numpy()) is not None and isinstance(scalar, np.generic)


def is_pandas_like_dataframe(df: Any) -> bool:
    """Check whether `df` is a pandas-like DataFrame without doing any imports.

    By "pandas-like", we mean: pandas, Modin, cuDF.
    """
    return is_pandas_dataframe(df) or is_modin_dataframe(df) or is_cudf_dataframe(df)


def is_pandas_like_series(ser: Any) -> bool:
    """Check whether `ser` is a pandas-like Series without doing any imports.

    By "pandas-like", we mean: pandas, Modin, cuDF.
    """
    return is_pandas_series(ser) or is_modin_series(ser) or is_cudf_series(ser)


def is_pandas_like_index(index: Any) -> bool:
    """Check whether `index` is a pandas-like Index without doing any imports.

    By "pandas-like", we mean: pandas, Modin, cuDF.
    """
    return (
        is_pandas_index(index) or is_modin_index(index) or is_cudf_index(index)
    )  # pragma: no cover


def is_into_series(native_series: Any | IntoSeriesT) -> TypeIs[IntoSeriesT]:
    """Check whether `native_series` can be converted to a Narwhals Series.

    Arguments:
        native_series: The object to check.

    Returns:
        `True` if `native_series` can be converted to a Narwhals Series, `False` otherwise.

    Examples:
        >>> import pandas as pd
        >>> import polars as pl
        >>> import numpy as np
        >>> import narwhals as nw

        >>> s_pd = pd.Series([1, 2, 3])
        >>> s_pl = pl.Series([1, 2, 3])
        >>> np_arr = np.array([1, 2, 3])

        >>> nw.dependencies.is_into_series(s_pd)
        True
        >>> nw.dependencies.is_into_series(s_pl)
        True
        >>> nw.dependencies.is_into_series(np_arr)
        False
    """
    from narwhals.series import Series

    return (
        isinstance(native_series, Series)
        or hasattr(native_series, "__narwhals_series__")
        or is_polars_series(native_series)
        or is_pyarrow_chunked_array(native_series)
        or is_pandas_like_series(native_series)
    )


def is_into_dataframe(native_dataframe: Any | IntoDataFrameT) -> TypeIs[IntoDataFrameT]:
    """Check whether `native_dataframe` can be converted to a Narwhals DataFrame.

    Arguments:
        native_dataframe: The object to check.

    Returns:
        `True` if `native_dataframe` can be converted to a Narwhals DataFrame, `False` otherwise.

    Examples:
        >>> import pandas as pd
        >>> import polars as pl
        >>> import numpy as np
        >>> from narwhals.dependencies import is_into_dataframe

        >>> df_pd = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
        >>> df_pl = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
        >>> np_arr = np.array([[1, 4], [2, 5], [3, 6]])

        >>> is_into_dataframe(df_pd)
        True
        >>> is_into_dataframe(df_pl)
        True
        >>> is_into_dataframe(np_arr)
        False
    """
    from narwhals.dataframe import DataFrame

    return (
        isinstance(native_dataframe, DataFrame)
        or hasattr(native_dataframe, "__narwhals_dataframe__")
        or is_polars_dataframe(native_dataframe)
        or is_pyarrow_table(native_dataframe)
        or is_pandas_like_dataframe(native_dataframe)
    )


def is_narwhals_dataframe(
    df: DataFrame[IntoDataFrameT] | Any,
) -> TypeIs[DataFrame[IntoDataFrameT]]:
    """Check whether `df` is a Narwhals DataFrame.

    This is useful if you expect a user to pass in a Narwhals
    DataFrame directly, and you want to catch both `narwhals.DataFrame`
    and `narwhals.stable.v1.DataFrame`.
    """
    from narwhals.dataframe import DataFrame

    return isinstance(df, DataFrame)


def is_narwhals_lazyframe(lf: Any | LazyFrame[FrameT]) -> TypeIs[LazyFrame[FrameT]]:
    """Check whether `lf` is a Narwhals LazyFrame.

    This is useful if you expect a user to pass in a Narwhals
    LazyFrame directly, and you want to catch both `narwhals.LazyFrame`
    and `narwhals.stable.v1.LazyFrame`.
    """
    from narwhals.dataframe import LazyFrame

    return isinstance(lf, LazyFrame)


def is_narwhals_series(ser: Any | Series[IntoSeriesT]) -> TypeIs[Series[IntoSeriesT]]:
    """Check whether `ser` is a Narwhals Series.

    This is useful if you expect a user to pass in a Narwhals
    Series directly, and you want to catch both `narwhals.Series`
    and `narwhals.stable.v1.Series`.
    """
    from narwhals.series import Series

    return isinstance(ser, Series)


def is_narwhals_series_int(ser: Any | Series[IntoSeriesT]) -> TypeIs[Series[IntoSeriesT]]:
    return is_narwhals_series(ser) and ser.dtype.is_integer()


__all__ = [
    "get_cudf",
    "get_ibis",
    "get_modin",
    "get_numpy",
    "get_pandas",
    "get_polars",
    "get_pyarrow",
    "is_cudf_dataframe",
    "is_cudf_series",
    "is_dask_dataframe",
    "is_ibis_table",
    "is_into_dataframe",
    "is_into_series",
    "is_modin_dataframe",
    "is_modin_series",
    "is_narwhals_dataframe",
    "is_narwhals_lazyframe",
    "is_narwhals_series",
    "is_numpy_array",
    "is_pandas_dataframe",
    "is_pandas_index",
    "is_pandas_like_dataframe",
    "is_pandas_like_series",
    "is_pandas_series",
    "is_polars_dataframe",
    "is_polars_lazyframe",
    "is_polars_series",
    "is_pyarrow_chunked_array",
    "is_pyarrow_table",
]