feat: add as_series flag to maybe_get_index to return the index as a narwhals Series

src/narwhals/_utils.py

@@ -868,11 +868,17 @@ def _validate_index(index: Any) -> None:
     return lhs
 
 
-def maybe_get_index(obj: DataFrame[Any] | LazyFrame[Any] | Series[Any]) -> Any | None:
+def maybe_get_index(
+    obj: DataFrame[Any] | LazyFrame[Any] | Series[Any], *, as_series: bool = False
+) -> Any | None:
     """Get the index of a DataFrame or a Series, if it's pandas-like.
 
     Arguments:
         obj: Dataframe or Series.
+        as_series: If `True`, return the index wrapped as a Narwhals `Series`
+            (matching the implementation of `obj`) instead of the raw native
+            index. This makes the result directly usable as the `index`
+            argument of [`maybe_set_index`][narwhals.maybe_set_index].
 
     Notes:
         This is only really intended for backwards-compatibility purposes,
@@ -893,11 +899,37 @@ def maybe_get_index(obj: DataFrame[Any] | LazyFrame[Any] | Series[Any]) -> Any |
         >>> series = nw.from_native(series_pd, series_only=True)
         >>> nw.maybe_get_index(series)
         RangeIndex(start=0, stop=2, step=1)
+        >>> nw.maybe_get_index(df, as_series=True)
+        ┌───────────────┐
+        |Narwhals Series|
+        |---------------|
+        | 0    0        |
+        | 1    1        |
+        | dtype: int64  |
+        └───────────────┘
     """
+    from narwhals.translate import _from_native_impl
+
     obj_any = cast("Any", obj)
     native_obj = obj_any.to_native()
     if is_pandas_like_dataframe(native_obj) or is_pandas_like_series(native_obj):
-        return native_obj.index
+        native_index = native_obj.index
+        if not as_series:
+            return native_index
+        compliant = getattr(obj_any, "_compliant_frame", None)
+        if compliant is None:
+            compliant = obj_any._compliant_series
+        ns = obj_any.implementation.to_native_namespace()
+        native_index_series = ns.Series(native_index, name=native_index.name)
+        return _from_native_impl(
+            native_index_series,
+            pass_through=False,
+            eager_only=False,
+            eager_or_interchange_only=False,
+            series_only=True,
+            allow_series=None,
+            version=compliant._version,
+        )
     return None
 
 

tests/utils_test.py

@@ -197,6 +197,41 @@ def test_maybe_get_index_pandas() -> None:
     assert_index_equal(result_s, pandas_series.index)
 
 
+def test_maybe_get_index_pandas_as_series() -> None:
+    pandas_df = pd.DataFrame({"a": [1, 2, 3]}, index=[1, 2, 0])
+    result = nw.maybe_get_index(nw.from_native(pandas_df), as_series=True)
+    assert isinstance(result, nw.Series)
+    assert result.implementation.is_pandas()
+    assert_index_equal(pd.Index(result.to_native()), 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), as_series=True
+    )
+    assert isinstance(result_s, nw.Series)
+    assert_index_equal(pd.Index(result_s.to_native()), pandas_series.index)
+
+
+def test_maybe_get_index_pandas_as_series_empty() -> None:
+    df = nw.from_native(pd.DataFrame({"a": []}))
+    result = nw.maybe_get_index(df, as_series=True)
+    assert isinstance(result, nw.Series)
+    assert len(result) == 0
+    series = nw.from_native(pd.Series([], dtype="float64"), series_only=True)
+    result_s = nw.maybe_get_index(series, as_series=True)
+    assert isinstance(result_s, nw.Series)
+    assert len(result_s) == 0
+
+
+def test_maybe_get_index_pandas_as_series_round_trip() -> None:
+    like = nw.from_native(pd.DataFrame({"a": [1, 2, 3]}, index=[7, 8, 9]))
+    new_series = nw.from_native(pd.Series([4, 5, 6]), series_only=True)
+    result = nw.maybe_set_index(
+        new_series, index=nw.maybe_get_index(like, as_series=True)
+    )
+    assert_index_equal(result.to_native().index, like.to_native().index)
+
+
 def test_maybe_get_index_polars() -> None:
     pytest.importorskip("polars")
     import polars as pl
@@ -209,6 +244,16 @@ def test_maybe_get_index_polars() -> None:
     assert result is None
 
 
+def test_maybe_get_index_polars_as_series() -> None:
+    pytest.importorskip("polars")
+    import polars as pl
+
+    df = nw.from_native(pl.DataFrame({"a": [1, 2, 3]}))
+    assert nw.maybe_get_index(df, as_series=True) is None
+    series = nw.from_native(pl.Series([1, 2, 3]), series_only=True)
+    assert nw.maybe_get_index(series, as_series=True) is None
+
+
 def test_maybe_reset_index_pandas() -> None:
     pandas_df = nw.from_native(
         pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[7, 8, 9])