fix(datatable): ensure columns work with DataFrame, list, and Var inputs

Author: zain2983

pyi_hashes.json

@@ -1,3 +1,3 @@
 {
-  "reflex/components/gridjs/datatable.pyi": "98a7e1b3f3b60cafcdfcd8879750ee42"
+  "reflex/components/gridjs/datatable.pyi": "e1f34ade3873a931770da4a35586f298"
 }

reflex/components/gridjs/datatable.py

@@ -1,4 +1,4 @@
-"""Table components."""
+"""Table components for Reflex using Gridjs."""
 
 from __future__ import annotations
 
@@ -14,58 +14,52 @@
 
 
 class Gridjs(NoSSRComponent):
-    """A component that wraps a nivo bar component."""
+    """A base component that wraps Gridjs (JS library) for tables."""
 
     library = "gridjs-react@6.1.1"
-
     lib_dependencies: list[str] = ["gridjs@6.2.0"]
 
 
 class DataTable(Gridjs):
-    """A data table component."""
+    """A flexible data table component for Reflex.
 
-    tag = "Grid"
+    Supports:
+    - Pandas DataFrames
+    - Python lists
+    - Reflex Vars (state variables)
+    """
 
+    tag = "Grid"
     alias = "DataTableGrid"
 
-    # The data to display. Either a list of lists or a pandas dataframe.
-    data: Any
-
-    # The list of columns to display. Required if data is a list and should not be provided
-    # if the data field is a dataframe
-    columns: Var[Sequence]
-
-    # Enable a search bar.
-    search: Var[bool]
-
-    # Enable sorting on columns.
-    sort: Var[bool]
-
-    # Enable resizable columns.
-    resizable: Var[bool]
-
-    # Enable pagination.
-    pagination: Var[bool | dict]
-
+    # -----------------------------
+    # Component Props
+    # -----------------------------
+    data: Any  # The data to display (list of lists or DataFrame)
+    columns: Var[Sequence]  # Columns to display (optional if using DataFrame)
+    search: Var[bool]  # Enable search
+    sort: Var[bool]  # Enable sorting
+    resizable: Var[bool]  # Enable column resizing
+    pagination: Var[bool | dict]  # Enable pagination
+
+    # -----------------------------
+    # Component creation
+    # -----------------------------
     @classmethod
     def create(cls, *children, **props):
-        """Create a datatable component.
+        """Create a DataTable component with proper validation.
 
-        Args:
-            *children: The children of the component.
-            **props: The props to pass to the component.
+        Raises:
+            ValueError: If both DataFrame and columns are provided, or
+                        if columns are missing for a list-type data field.
 
         Returns:
-            The datatable component.
-
-        Raises:
-            ValueError: If a pandas dataframe is passed in and columns are also provided.
+            DataTable: The created DataTable component.
         """
         data = props.get("data")
         columns = props.get("columns")
 
-        # The annotation should be provided if data is a computed var. We need this to know how to
-        # render pandas dataframes.
+        # 1️⃣ Ensure computed Vars have type annotations
         if is_computed_var(data) and data._var_type == Any:
             msg = "Annotation of the computed var assigned to the data field should be provided."
             raise ValueError(msg)
@@ -78,38 +72,50 @@ def create(cls, *children, **props):
             msg = "Annotation of the computed var assigned to the column field should be provided."
             raise ValueError(msg)
 
-        # If data is a pandas dataframe and columns are provided throw an error.
+        # 2️⃣ Disallow DataFrame + columns (columns auto-detected from DataFrame)
         if (
             types.is_dataframe(type(data))
             or (isinstance(data, Var) and types.is_dataframe(data._var_type))
         ) and columns is not None:
             msg = "Cannot pass in both a pandas dataframe and columns to the data_table component."
             raise ValueError(msg)
 
-        # If data is a list and columns are not provided, throw an error
+        # 3️⃣ Require columns if data is a list
         if (
             (isinstance(data, Var) and types.typehint_issubclass(data._var_type, list))
             or isinstance(data, list)
         ) and columns is None:
-            msg = "column field should be specified when the data field is a list type"
+            msg = "Column field should be specified when the data field is a list type"
             raise ValueError(msg)
 
-        # Create the component.
-        return super().create(
-            *children,
-            **props,
-        )
+        # 4️⃣ Call parent create method
+        return super().create(*children, **props)
 
+    # -----------------------------
+    # Add external imports (CSS)
+    # -----------------------------
     def add_imports(self) -> ImportDict:
-        """Add the imports for the datatable component.
+        """Add CSS for Gridjs.
 
         Returns:
-            The import dict for the component.
+            ImportDict: The import dictionary required for the component.
         """
         return {"": "gridjs/dist/theme/mermaid.css"}
 
+    # -----------------------------
+    # Render component
+    # -----------------------------
     def _render(self) -> Tag:
+        """Normalize columns and prepare data for front-end rendering.
+
+        Returns:
+            Tag: The rendered table component.
+        """
+        # -----------------------------
+        # Case 1: DataFrame coming from State (Var)
+        # -----------------------------
         if isinstance(self.data, Var) and types.is_dataframe(self.data._var_type):
+            # Convert DataFrame to front-end-safe Vars
             self.columns = self.data._replace(
                 _js_expr=f"{self.data._js_expr}.columns",
                 _var_type=list[Any],
@@ -118,23 +124,53 @@ def _render(self) -> Tag:
                 _js_expr=f"{self.data._js_expr}.data",
                 _var_type=list[list[Any]],
             )
+
+        # -----------------------------
+        # Case 2: DataFrame passed directly from Python
+        # -----------------------------
         if types.is_dataframe(type(self.data)):
-            # If given a pandas df break up the data and columns
             data = serialize(self.data)
             if not isinstance(data, dict):
                 msg = "Serialized dataframe should be a dict."
                 raise ValueError(msg)
+
+            # Convert Python lists to LiteralVars for front-end rendering
             self.columns = LiteralVar.create(data["columns"])
             self.data = LiteralVar.create(data["data"])
 
-        # If columns is a list of strings convert to list of dicts with id and name keys
-        if isinstance(self.columns, LiteralVar) and isinstance(
-            self.columns._var_value, list
-        ):
-            self.columns = LiteralVar.create([
-                {"id": col, "name": col} if isinstance(col, str) else col
-                for col in self.columns._var_value
-            ])
-
-        # Render the table.
+        # -----------------------------
+        # Case 3: Normalize columns for all other scenarios
+        # -----------------------------
+        if self.columns is not None:
+            # Python list → LiteralVar
+            if isinstance(self.columns, list):
+                self.columns = LiteralVar.create([
+                    {"id": col, "name": col} if isinstance(col, str) else col
+                    for col in self.columns
+                ])
+
+            # LiteralVar[list] → normalized LiteralVar
+            elif isinstance(self.columns, LiteralVar) and isinstance(
+                self.columns._var_value, list
+            ):
+                self.columns = LiteralVar.create([
+                    {"id": col, "name": col} if isinstance(col, str) else col
+                    for col in self.columns._var_value
+                ])
+
+            # Var[list] → frontend-safe JS mapping (compile-time + runtime safe)
+            elif isinstance(self.columns, Var):
+                self.columns = self.columns._replace(
+                    _js_expr=(
+                        f"{self.columns._js_expr}.map("
+                        "(col) => typeof col === 'string' "
+                        "? ({ id: col, name: col }) "
+                        ": col)"
+                    ),
+                    _var_type=list[Any],
+                )
+
+        # -----------------------------
+        # Case 4: Render component
+        # -----------------------------
         return super()._render()

tests/units/components/gridjs/__init__.py

@@ -0,0 +1 @@
+"""gridjs component tests."""

tests/units/components/gridjs/test_datatable.py

@@ -0,0 +1,90 @@
+import pandas as pd
+import pytest
+
+from reflex.components.gridjs import DataTable
+from reflex.state import State
+from reflex.vars.base import LiteralVar, Var
+
+
+class TableState(State):
+    """TableState used by tests."""
+
+    df: pd.DataFrame = pd.DataFrame({"A": [1], "B": [2]})
+    data: list[list[int]] = [[1, 2]]
+    columns: list[str] = ["A", "B"]
+
+
+def test_dataframe_python_columns_normalized():
+    """DataFrame passed as a Python value should produce normalized LiteralVar columns."""
+    df = pd.DataFrame({"A": [1], "B": [2]})
+
+    table = DataTable.create(data=df)
+    table._render()
+
+    columns = table.columns  # type: ignore[attr-defined]
+
+    assert isinstance(columns, LiteralVar)
+    assert columns._var_value == [
+        {"id": "A", "name": "A"},
+        {"id": "B", "name": "B"},
+    ]
+
+
+def test_dataframe_var_columns_preserved():
+    """DataFrame coming from State is a runtime value.
+    Columns must remain a Var (frontend expression),
+    not be eagerly converted to LiteralVar.
+    """
+    table = DataTable.create(data=TableState.df)
+    table._render()
+
+    columns = table.columns  # type: ignore[attr-defined]
+
+    assert isinstance(columns, Var)
+    assert not isinstance(columns, LiteralVar)
+
+
+def test_list_columns_python_normalized():
+    """Python list of column names should be normalized eagerly."""
+    table = DataTable.create(
+        data=[[1, 2]],
+        columns=["A", "B"],
+    )
+    table._render()
+
+    columns = table.columns  # type: ignore[attr-defined]
+
+    assert isinstance(columns, LiteralVar)
+    assert columns._var_value == [
+        {"id": "A", "name": "A"},
+        {"id": "B", "name": "B"},
+    ]
+
+
+def test_list_columns_var_preserved():
+    """Columns coming from State must remain a Var so they
+    can be transformed on the frontend.
+    """
+    table = DataTable.create(
+        data=TableState.data,
+        columns=TableState.columns,
+    )
+    table._render()
+
+    columns = table.columns  # type: ignore[attr-defined]
+
+    assert isinstance(columns, Var)
+    assert not isinstance(columns, LiteralVar)
+
+
+def test_dataframe_with_columns_raises():
+    """DataFrame already defines columns.
+    Passing explicit columns is ambiguous and must error.
+    """
+    df = pd.DataFrame({"A": [1]})
+
+    with pytest.raises(ValueError):
+        DataTable.create(
+            data=df,
+            columns=["A"],
+        )