From 68bc8629c53b2ee7945ea624781414f1a8d0d0a2 Mon Sep 17 00:00:00 2001 From: Marvin Lopez Date: Tue, 10 Dec 2024 07:27:33 -0500 Subject: [PATCH] enh: nw.dependencies.is_narwhals_dataframe / nw.dependencies.is_narwhals_lazyframe / nw.dependencies.is_narwhals_series (#1550) --------- Co-authored-by: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> --- narwhals/dependencies.py | 42 +++++++++++++++++++ narwhals/stable/v1/dependencies.py | 6 +++ .../is_narwhals_dataframe_test.py | 18 ++++++++ .../is_narwhals_lazyframe_test.py | 19 +++++++++ tests/dependencies/is_narwhals_series_test.py | 18 ++++++++ 5 files changed, 103 insertions(+) create mode 100644 tests/dependencies/is_narwhals_dataframe_test.py create mode 100644 tests/dependencies/is_narwhals_lazyframe_test.py create mode 100644 tests/dependencies/is_narwhals_series_test.py diff --git a/narwhals/dependencies.py b/narwhals/dependencies.py index 463a64a67..61937f08d 100644 --- a/narwhals/dependencies.py +++ b/narwhals/dependencies.py @@ -24,6 +24,9 @@ import pyarrow as pa import pyspark.sql as pyspark_sql + from narwhals.dataframe import DataFrame + from narwhals.dataframe import LazyFrame + from narwhals.series import Series from narwhals.typing import IntoSeries # We silently allow these but - given that they claim @@ -320,6 +323,42 @@ def is_into_dataframe(native_dataframe: Any) -> bool: ) +def is_narwhals_dataframe(df: Any) -> TypeGuard[DataFrame[Any]]: + """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) -> TypeGuard[LazyFrame[Any]]: + """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) -> TypeGuard[Series[Any]]: + """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", @@ -336,6 +375,9 @@ def is_into_dataframe(native_dataframe: Any) -> bool: "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", diff --git a/narwhals/stable/v1/dependencies.py b/narwhals/stable/v1/dependencies.py index 6a020622e..8ad1ad201 100644 --- a/narwhals/stable/v1/dependencies.py +++ b/narwhals/stable/v1/dependencies.py @@ -15,6 +15,9 @@ from narwhals.dependencies import is_into_series from narwhals.dependencies import is_modin_dataframe from narwhals.dependencies import is_modin_series +from narwhals.dependencies import is_narwhals_dataframe +from narwhals.dependencies import is_narwhals_lazyframe +from narwhals.dependencies import is_narwhals_series from narwhals.dependencies import is_numpy_array from narwhals.dependencies import is_pandas_dataframe from narwhals.dependencies import is_pandas_index @@ -43,6 +46,9 @@ "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", diff --git a/tests/dependencies/is_narwhals_dataframe_test.py b/tests/dependencies/is_narwhals_dataframe_test.py new file mode 100644 index 000000000..0a25bc8b6 --- /dev/null +++ b/tests/dependencies/is_narwhals_dataframe_test.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import narwhals as nw +import narwhals.stable.v1 as nws +from narwhals.stable.v1.dependencies import is_narwhals_dataframe + +if TYPE_CHECKING: + from tests.utils import ConstructorEager + + +def test_is_narwhals_dataframe(constructor_eager: ConstructorEager) -> None: + df = constructor_eager({"col1": [1, 2], "col2": [3, 4]}) + + assert is_narwhals_dataframe(nw.from_native(df)) + assert is_narwhals_dataframe(nws.from_native(df)) + assert not is_narwhals_dataframe(df) diff --git a/tests/dependencies/is_narwhals_lazyframe_test.py b/tests/dependencies/is_narwhals_lazyframe_test.py new file mode 100644 index 000000000..1dcbe5fcc --- /dev/null +++ b/tests/dependencies/is_narwhals_lazyframe_test.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import narwhals as nw +import narwhals.stable.v1 as nws +from narwhals.stable.v1.dependencies import is_narwhals_lazyframe +from tests.utils import Constructor + +if TYPE_CHECKING: + from tests.utils import Constructor + + +def test_is_narwhals_lazyframe(constructor: Constructor) -> None: + lf = constructor({"a": [1, 2, 3]}) + + assert is_narwhals_lazyframe(nw.from_native(lf).lazy()) + assert is_narwhals_lazyframe(nws.from_native(lf).lazy()) + assert not is_narwhals_lazyframe(lf) diff --git a/tests/dependencies/is_narwhals_series_test.py b/tests/dependencies/is_narwhals_series_test.py new file mode 100644 index 000000000..0beb3fc15 --- /dev/null +++ b/tests/dependencies/is_narwhals_series_test.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import narwhals as nw +import narwhals.stable.v1 as nws +from narwhals.stable.v1.dependencies import is_narwhals_series + +if TYPE_CHECKING: + from tests.utils import ConstructorEager + + +def test_is_narwhals_series(constructor_eager: ConstructorEager) -> None: + df = constructor_eager({"col1": [1, 2], "col2": [3, 4]}) + + assert is_narwhals_series(nw.from_native(df, eager_only=True)["col1"]) + assert is_narwhals_series(nws.from_native(df, eager_only=True)["col1"]) + assert not is_narwhals_series(nw.from_native(df, eager_only=True)["col1"].to_native())