feat: start rough draft of narwhals support

Author: machow

shiny/render/_data_frame_utils/_tbl_data.py

@@ -5,6 +5,7 @@
 from functools import singledispatch
 from typing import Any, List, Tuple, cast
 
+import narwhals.stable.v1 as nw
 from htmltools import TagNode
 
 from ..._typing_extensions import TypeIs
@@ -184,7 +185,7 @@ def _(col: PlSeries) -> FrameDtype:
 
     from ._html import col_contains_shiny_html
 
-    if col.dtype.is_(pl.String):
+    if col.dtype == pl.String():
         if col_contains_shiny_html(col):
             type_ = "html"
         else:
@@ -203,6 +204,30 @@ def _(col: PlSeries) -> FrameDtype:
     return {"type": type_}
 
 
+@serialize_dtype.register
+def _(col: nw.Series) -> FrameDtype:
+
+    from ._html import col_contains_shiny_html
+
+    if col.dtype == nw.String():
+        if col_contains_shiny_html(col):
+            type_ = "html"
+        else:
+            type_ = "string"
+    elif col.dtype.is_numeric():
+        type_ = "numeric"
+
+    elif col.dtype == nw.Categorical():
+        categories = col.cat.get_categories().to_list()
+        return {"type": "categorical", "categories": categories}
+    else:
+        type_ = "unknown"
+        if col_contains_shiny_html(col):
+            type_ = "html"
+
+    return {"type": type_}
+
+
 # serialize_frame ----------------------------------------------------------------------
 
 
@@ -218,6 +243,8 @@ def _(data: PdDataFrame) -> FrameJson:
     return serialize_frame_pd(data)
 
 
+# TODO: test this
+@serialize_frame.register(nw.DataFrame)
 @serialize_frame.register
 def _(data: PlDataFrame) -> FrameJson:
     import json
@@ -308,6 +335,7 @@ def _(
         return data.iloc[indx_rows, indx_cols]
 
 
+@subset_frame.register(nw.DataFrame)
 @subset_frame.register
 def _(
     data: PlDataFrame,
@@ -321,7 +349,7 @@ def _(
         else slice(None)
     )
     indx_rows = rows if rows is not None else slice(None)
-    return data[indx_rows, indx_cols]
+    return data[indx_cols][indx_rows]
 
 
 # get_frame_cell -----------------------------------------------------------------------
@@ -341,9 +369,10 @@ def _(data: PdDataFrame, row: int, col: int) -> Any:
     )
 
 
+@get_frame_cell.register(nw.DataFrame)
 @get_frame_cell.register
 def _(data: PlDataFrame, row: int, col: int) -> Any:
-    return data[row, col]
+    return data.item(row, col)
 
 
 # shape --------------------------------------------------------------------------------
@@ -359,6 +388,7 @@ def _(data: PdDataFrame) -> Tuple[int, ...]:
     return data.shape
 
 
+@frame_shape.register(nw.DataFrame)
 @frame_shape.register
 def _(data: PlDataFrame) -> Tuple[int, ...]:
     return data.shape
@@ -377,6 +407,7 @@ def _(data: PdDataFrame) -> PdDataFrame:
     return data.copy()
 
 
+@copy_frame.register(nw.DataFrame)
 @copy_frame.register
 def _(data: PlDataFrame) -> PlDataFrame:
     return data.clone()
@@ -393,6 +424,7 @@ def _(data: PdDataFrame) -> List[str]:
     return data.columns.to_list()
 
 
+@frame_column_names.register(nw.DataFrame)
 @frame_column_names.register
 def _(data: PlDataFrame) -> List[str]:
     return data.columns

tests/pytest/test_render_data_frame.py

@@ -1,3 +1,4 @@
+import narwhals as nw
 import pandas as pd
 import pytest
 

tests/pytest/test_render_data_frame_tbl_data.py

@@ -4,6 +4,7 @@
 from datetime import datetime
 from typing import Any, Union
 
+import narwhals.stable.v1 as nw
 import pandas as pd
 import polars as pl
 import polars.testing as pl_testing
@@ -53,6 +54,9 @@ class D:
 params_frames = [
     pytest.param(pd.DataFrame, id="pandas"),
     pytest.param(pl.DataFrame, id="polars"),
+    pytest.param(
+        lambda d: nw.from_native(pl.DataFrame(d), eager_only=True), id="narwhals"
+    ),
 ]
 
 DataFrameLike: TypeAlias = Union[pd.DataFrame, pl.DataFrame]
@@ -93,6 +97,17 @@ def assert_frame_equal(
         raise NotImplementedError(f"Unsupported data type: {type(src)}")
 
 
+def assert_frame_equal2(
+    src: pd.DataFrame | pl.DataFrame,
+    target_dict: dict,
+    use_index: bool = False,
+):
+    src = nw.to_native(src, strict=False)
+    target = nw.to_native(src, strict=False).__class__(target_dict)
+
+    assert_frame_equal(src, target, use_index)
+
+
 # TODO: explicitly pass dtype= when doing Series construction
 @pytest.mark.parametrize(
     "ser, res_type",
@@ -126,6 +141,13 @@ def test_serialize_dtype(
     ],
     res_type: str,
 ):
+    if isinstance(ser, pl.Series):
+        assert (
+            serialize_dtype(nw.from_native(ser, eager_only=True, allow_series=True))[
+                "type"
+            ]
+            == res_type
+        )
     assert serialize_dtype(ser)["type"] == res_type
 
 
@@ -158,9 +180,9 @@ def test_serialize_frame(df: DataFrameLike):
 def test_subset_frame(df: DataFrameLike):
     # TODO: this assumes subset_frame doesn't reset index
     res = subset_frame(df, rows=[1], cols=["chr", "num"])
-    dst = df.__class__({"chr": ["b"], "num": [2]})
+    dst = {"chr": ["b"], "num": [2]}
 
-    assert_frame_equal(res, dst)
+    assert_frame_equal2(res, dst)
 
 
 def test_get_frame_cell(df: DataFrameLike):
@@ -176,14 +198,20 @@ def test_copy_frame(df: DataFrameLike):
 def test_subset_frame_rows_single(small_df: DataFrameLike):
     res = subset_frame(small_df, rows=[1])
 
-    assert_frame_equal(res, small_df.__class__({"x": [2], "y": [4]}))
+    assert_frame_equal2(
+        res,
+        {"x": [2], "y": [4]},
+    )
 
 
 def test_subset_frame_cols_single(small_df: DataFrameLike):
     # TODO: include test of polars
     res = subset_frame(small_df, cols=["y"])
 
-    assert_frame_equal(res, small_df.__class__({"y": [3, 4]}))
+    assert_frame_equal2(
+        res,
+        {"y": [3, 4]},
+    )
 
 
 def test_shape(small_df: DataFrameLike):