Repository: narwhals-dev/narwhals
Pull request #148 (2024-05-12 13:37:59)

Maybe set index


Patch:
diff --git a/docs/api-reference/narwhals.md b/docs/api-reference/narwhals.md
index 8425ed17..863d3ad9 100644
--- a/docs/api-reference/narwhals.md
+++ b/docs/api-reference/narwhals.md
@@ -13,6 +13,7 @@ Here are the top-level functions available in Narwhals.
         - get_native_namespace
         - len
         - maybe_align_index
+        - maybe_set_index
         - max
         - mean
         - min
diff --git a/narwhals/__init__.py b/narwhals/__init__.py
index d3038b49..2bd270b0 100644
--- a/narwhals/__init__.py
+++ b/narwhals/__init__.py
@@ -28,6 +28,7 @@ from narwhals.translate import from_native
 from narwhals.translate import get_native_namespace
 from narwhals.translate import to_native
 from narwhals.utils import maybe_align_index
+from narwhals.utils import maybe_set_index
 
 __version__ = "0.8.11"
 
@@ -36,6 +37,7 @@ __all__ = [
     "to_native",
     "from_native",
     "maybe_align_index",
+    "maybe_set_index",
     "get_native_namespace",
     "all",
     "col",
diff --git a/narwhals/utils.py b/narwhals/utils.py
index 81b02279..0ee71ef6 100644
--- a/narwhals/utils.py
+++ b/narwhals/utils.py
@@ -161,3 +161,39 @@ def maybe_align_index(lhs: T, rhs: Series | BaseFrame) -> T:
         msg = f"Expected `lhs` and `rhs` to have the same length, got {len(lhs_any)} and {len(rhs_any)}"
         raise ValueError(msg)
     return lhs
+
+
+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.
+
+    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 is a no-op.
+
+    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.to_native(nw.maybe_set_index(df, 'b'))  # doctest: +NORMALIZE_WHITESPACE
+           a
+        b
+        4  1
+        5  2
+    """
+    from narwhals._pandas_like.dataframe import PandasDataFrame
+    from narwhals.dataframe import DataFrame
+
+    df_any = cast(Any, df)
+    if isinstance(getattr(df_any, "_dataframe", None), PandasDataFrame):
+        return DataFrame(  # type: ignore[return-value]
+            df_any._dataframe._from_dataframe(
+                df_any._dataframe._dataframe.set_index(column_names)
+            )
+        )
+    return df


Test patch:
diff --git a/tests/utils_test.py b/tests/utils_test.py
index bdd22796..228e7ec7 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -45,3 +45,18 @@ def test_native_namespace() -> None:
     assert nw.get_native_namespace(df) is pl
     df = nw.from_native(pd.DataFrame({"a": [1, 2, 3]}))
     assert nw.get_native_namespace(df) is pd
+
+
+def test_maybe_set_index_pandas() -> None:
+    df = nw.from_native(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[1, 2, 0]))
+    result = nw.maybe_set_index(df, "b")
+    expected = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[1, 2, 0]).set_index(
+        "b"
+    )
+    assert_frame_equal(nw.to_native(result), expected)
+
+
+def test_maybe_set_index_polars() -> None:
+    df = nw.from_native(pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    result = nw.maybe_set_index(df, "b")
+    assert result is df