435 lines
14 KiB
Python
435 lines
14 KiB
Python
# 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
|
|
from typing import Any
|
|
|
|
if TYPE_CHECKING:
|
|
import cudf
|
|
import dask.dataframe as dd
|
|
import duckdb
|
|
import ibis
|
|
import modin.pandas as mpd
|
|
import numpy as np
|
|
import pandas as pd
|
|
import polars as pl
|
|
import pyarrow as pa
|
|
import pyspark.sql as pyspark_sql
|
|
from typing_extensions import TypeGuard
|
|
from typing_extensions import TypeIs
|
|
|
|
from narwhals._arrow.typing import ArrowChunkedArray
|
|
from narwhals._spark_like.dataframe import SQLFrameDataFrame
|
|
from narwhals.dataframe import DataFrame
|
|
from narwhals.dataframe import LazyFrame
|
|
from narwhals.series import Series
|
|
from narwhals.typing import DataFrameT
|
|
from narwhals.typing import FrameT
|
|
from narwhals.typing import IntoDataFrameT
|
|
from narwhals.typing import IntoSeriesT
|
|
from narwhals.typing import _1DArray
|
|
from narwhals.typing import _2DArray
|
|
from narwhals.typing import _NDArray
|
|
from narwhals.typing import _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)."""
|
|
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_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]:
|
|
"""Check whether `index` is a modin Index without importing modin."""
|
|
return (mpd := get_modin()) is not None and isinstance(
|
|
index, mpd.Index
|
|
) # pragma: no cover
|
|
|
|
|
|
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_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 | ArrowChunkedArray,
|
|
) -> TypeIs[ArrowChunkedArray]:
|
|
"""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_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_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_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[np.generic]:
|
|
"""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 np.isscalar(scalar)
|
|
|
|
|
|
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: Any | DataFrame[DataFrameT],
|
|
) -> TypeIs[DataFrame[DataFrameT]]:
|
|
"""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)
|
|
|
|
|
|
__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",
|
|
]
|