Support dataframes as state vars (#324)

pynecone/components/datadisplay/datatable.py

@@ -5,7 +5,7 @@ from typing import Any, List
 from pynecone import utils
 from pynecone.components.component import Component, ImportDict
 from pynecone.components.tags import Tag
-from pynecone.var import Var
+from pynecone.var import BaseVar, Var
 
 
 class Gridjs(Component):
@@ -69,9 +69,23 @@ class DataTable(Gridjs):
         )
 
     def _render(self) -> Tag:
+        # If given a pandas df break up the data and columns
         if utils.is_dataframe(type(self.data)):
-            # If given a pandas df break up the data and columns
             self.columns = Var.create(list(self.data.columns.values.tolist()))  # type: ignore
             self.data = Var.create(list(self.data.values.tolist()))  # type: ignore
 
+        # If given a var dataframe, get the data and columns
+        if isinstance(self.data, Var):
+            self.columns = BaseVar(
+                name=f"{self.data.name}.columns",
+                type_=List[Any],
+                state=self.data.state,
+            )
+            self.data = BaseVar(
+                name=f"{self.data.name}.data",
+                type_=List[List[Any]],
+                state=self.data.state,
+            )
+
+        # Render the table.
         return super()._render()

pynecone/state.py

@@ -112,7 +112,7 @@ class State(Base, ABC):
 
         # Setup the base vars at the class level.
         for prop in cls.base_vars.values():
-            if not utils._issubclass(prop.type_, utils.StateVar):
+            if not utils.is_valid_var_type(prop.type_):
                 raise TypeError(
                     "State vars must be primitive Python types, "
                     "Plotly figures, Pandas dataframes, "

pynecone/utils.py

@@ -976,6 +976,18 @@ def is_dataframe(value: Type) -> bool:
     return value.__name__ == "DataFrame"
 
 
+def is_valid_var_type(var: Type) -> bool:
+    """Check if the given value is a valid prop type.
+
+    Args:
+        var: The value to check.
+
+    Returns:
+        Whether the value is a valid prop type.
+    """
+    return _issubclass(var, StateVar) or is_dataframe(var)
+
+
 def format_state(value: Any) -> Dict:
     """Recursively format values in the given state.
 

tests/components/graphing/test_victory_data.py

@@ -1,15 +1,8 @@
-from typing import List, Set, Type
-
 import pytest
 
 from pynecone import data
-from pynecone.components.component import Component, CustomComponent, ImportDict
-from pynecone.components.layout.box import Box
-from pynecone.event import EVENT_TRIGGERS, EventHandler
-from pynecone.state import State
-from pynecone.style import Style
-from pynecone.var import Var
 
+# Test data.
 x_num = [1, 2, 3, 4, 5]
 x_str = ["Cats", "Dogs", "Birds", "Fish", "Reptiles"]
 y = [1, 2, 3, 4, 10]

tests/test_state.py

@@ -4,10 +4,10 @@ import pytest
 
 from pynecone import utils
 from pynecone.base import Base
+from pynecone.constants import RouteVar
 from pynecone.event import Event
 from pynecone.state import State
 from pynecone.var import BaseVar, ComputedVar
-from pynecone.constants import RouteVar
 
 
 class Object(Base):