improve datatable hacks, also default to ObjectVar for Unions

reflex/components/gridjs/datatable.py

@@ -116,14 +116,8 @@ class DataTable(Gridjs):
         if isinstance(self.data, ImmutableVar) and types.is_dataframe(
             self.data._var_type
         ):
-            self.columns = self.data._replace(
-                _var_name=f"{self.data._var_name}.columns",
-                _var_type=List[Any],
-            )
-            self.data = self.data._replace(
-                _var_name=f"{self.data._var_name}.data",
-                _var_type=List[List[Any]],
-            )
+            self.columns = self.data.columns
+            self.data = self.data.data
         if types.is_dataframe(type(self.data)):
             # If given a pandas df break up the data and columns
             data = serialize(self.data)

reflex/ivars/base.py

@@ -482,7 +482,7 @@ class ImmutableVar(Var, Generic[VAR_TYPE]):
             if all(inspect.isclass(t) and issubclass(t, Base) for t in inner_types):
                 return self.to(ObjectVar, self._var_type)
 
-            return self
+            return self.to(ObjectVar, self._var_type)
 
         if not inspect.isclass(fixed_type):
             raise TypeError(f"Unsupported type {var_type} for guess_type.")

reflex/utils/types.py

@@ -338,6 +338,12 @@ def get_attribute_access_type(cls: GenericType, name: str) -> GenericType | None
                 return hints[name]
         except exceptions as e:
             console.warn(f"Failed to resolve ForwardRefs for {cls}.{name} due to {e}")
+            # hardcoded fallbacks for forward ref issues
+            if is_dataframe(cls):
+                if name == "columns":
+                    return List[Any]
+                if name == "data":
+                    return List[List[Any]]
             pass
     return None  # Attribute is not accessible.
 

tests/components/datadisplay/conftest.py

@@ -10,7 +10,7 @@ from reflex.state import BaseState
 
 
 @pytest.fixture
-def data_table_state(request):
+def data_table_state(request: pytest.FixtureRequest):
     """Get a data table state.
 
     Args:

tests/components/datadisplay/test_datatable.py

@@ -23,7 +23,7 @@ from reflex.utils.serializers import serialize, serialize_dataframe
     ],
     indirect=["data_table_state"],
 )
-def test_validate_data_table(data_table_state: rx.State, expected):
+def test_validate_data_table(data_table_state: rx.State, expected: str) -> None:
     """Test the str/render function.
 
     Args:
@@ -40,13 +40,13 @@ def test_validate_data_table(data_table_state: rx.State, expected):
 
     data_table_dict = data_table_component.render()
 
-    # prefix expected with state name
-    state_name = data_table_state.get_name()
-    expected = f"{state_name}.{expected}" if expected else state_name
+    var = data_table_state
+    if expected:
+        var = getattr(var, expected)
 
     assert data_table_dict["props"] == [
-        f"columns={{{expected}.columns}}",
-        f"data={{{expected}.data}}",
+        f"columns={{{var.columns._var_name}}}",
+        f"data={{{var.data._var_name}}}",
     ]