diff --git a/docs/api-reference/narwhals.md b/docs/api-reference/narwhals.md index d808263aa..9d7ac384e 100644 --- a/docs/api-reference/narwhals.md +++ b/docs/api-reference/narwhals.md @@ -21,6 +21,7 @@ Here are the top-level functions available in Narwhals. - max - maybe_align_index - maybe_convert_dtypes + - maybe_get_index - maybe_set_index - mean - mean_horizontal diff --git a/narwhals/__init__.py b/narwhals/__init__.py index 5180ef095..0977e716a 100644 --- a/narwhals/__init__.py +++ b/narwhals/__init__.py @@ -50,6 +50,7 @@ from narwhals.utils import is_ordered_categorical from narwhals.utils import maybe_align_index from narwhals.utils import maybe_convert_dtypes +from narwhals.utils import maybe_get_index from narwhals.utils import maybe_set_index __version__ = "1.5.5" @@ -66,6 +67,7 @@ "is_ordered_categorical", "maybe_align_index", "maybe_convert_dtypes", + "maybe_get_index", "maybe_set_index", "get_native_namespace", "all", diff --git a/narwhals/stable/v1.py b/narwhals/stable/v1.py index 8d79d6c73..d166a814f 100644 --- a/narwhals/stable/v1.py +++ b/narwhals/stable/v1.py @@ -49,6 +49,7 @@ from narwhals.utils import is_ordered_categorical as nw_is_ordered_categorical from narwhals.utils import maybe_align_index as nw_maybe_align_index from narwhals.utils import maybe_convert_dtypes as nw_maybe_convert_dtypes +from narwhals.utils import maybe_get_index as nw_maybe_get_index from narwhals.utils import maybe_set_index as nw_maybe_set_index if TYPE_CHECKING: @@ -1438,6 +1439,33 @@ def maybe_convert_dtypes(df: T, *args: bool, **kwargs: bool | str) -> T: return nw_maybe_convert_dtypes(df, *args, **kwargs) +def maybe_get_index(obj: T) -> Any | None: + """ + Get the index of a DataFrame or a Series, if it's pandas-like. + + Notes: + This is only really intended for backwards-compatibility purposes, + for example if your library already aligns indices for users. + If you're designing a new library, we highly encourage you to not + rely on the Index. + For non-pandas-like inputs, this returns `None`. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals.stable.v1 as nw + >>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]}) + >>> df = nw.from_native(df_pd) + >>> nw.maybe_get_index(df) + RangeIndex(start=0, stop=2, step=1) + >>> series_pd = pd.Series([1, 2]) + >>> series = nw.from_native(series_pd, series_only=True) + >>> nw.maybe_get_index(series) + RangeIndex(start=0, stop=2, step=1) + """ + return nw_maybe_get_index(obj) + + def maybe_set_index(df: T, column_names: str | list[str]) -> T: """ Set columns `columns` to be the index of `df`, if `df` is pandas-like. @@ -1678,6 +1706,7 @@ def from_dict( "is_ordered_categorical", "maybe_align_index", "maybe_convert_dtypes", + "maybe_get_index", "maybe_set_index", "get_native_namespace", "get_level", diff --git a/narwhals/utils.py b/narwhals/utils.py index f0e5c1cf6..6c1b5c1b4 100644 --- a/narwhals/utils.py +++ b/narwhals/utils.py @@ -240,6 +240,37 @@ def _validate_index(index: Any) -> None: return lhs +def maybe_get_index(obj: T) -> Any | None: + """ + Get the index of a DataFrame or a Series, if it's pandas-like. + + Notes: + This is only really intended for backwards-compatibility purposes, + for example if your library already aligns indices for users. + If you're designing a new library, we highly encourage you to not + rely on the Index. + For non-pandas-like inputs, this returns `None`. + + Examples: + >>> import pandas as pd + >>> import polars as pl + >>> import narwhals as nw + >>> df_pd = pd.DataFrame({"a": [1, 2], "b": [4, 5]}) + >>> df = nw.from_native(df_pd) + >>> nw.maybe_get_index(df) + RangeIndex(start=0, stop=2, step=1) + >>> series_pd = pd.Series([1, 2]) + >>> series = nw.from_native(series_pd, series_only=True) + >>> nw.maybe_get_index(series) + RangeIndex(start=0, stop=2, step=1) + """ + obj_any = cast(Any, obj) + native_obj = to_native(obj_any) + if is_pandas_like_dataframe(native_obj) or is_pandas_like_series(native_obj): + return native_obj.index + return None + + def maybe_set_index(df: T, column_names: str | list[str]) -> T: """ Set columns `columns` to be the index of `df`, if `df` is pandas-like. diff --git a/tests/test_utils.py b/tests/test_utils.py index b94c8371e..68dc90ed7 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -2,6 +2,7 @@ import polars as pl import pytest from pandas.testing import assert_frame_equal +from pandas.testing import assert_index_equal from pandas.testing import assert_series_equal import narwhals.stable.v1 as nw @@ -65,6 +66,24 @@ def test_maybe_set_index_polars() -> None: assert result is df +def test_maybe_get_index_pandas() -> None: + pandas_df = pd.DataFrame({"a": [1, 2, 3]}, index=[1, 2, 0]) + result = nw.maybe_get_index(nw.from_native(pandas_df)) + assert_index_equal(result, pandas_df.index) + pandas_series = pd.Series([1, 2, 3], index=[1, 2, 0]) + result_s = nw.maybe_get_index(nw.from_native(pandas_series, series_only=True)) + assert_index_equal(result_s, pandas_series.index) + + +def test_maybe_get_index_polars() -> None: + df = nw.from_native(pl.DataFrame({"a": [1, 2, 3]})) + result = nw.maybe_get_index(df) + assert result is None + series = nw.from_native(pl.Series([1, 2, 3]), series_only=True) + result = nw.maybe_get_index(series) + assert result is None + + @pytest.mark.skipif( parse_version(pd.__version__) < parse_version("1.0.0"), reason="too old for convert_dtypes",