From d376d2972b8365a46d2105a490920c6179dfc944 Mon Sep 17 00:00:00 2001
From: Nikhil Rao <nikhil@pynecone.io>
Date: Mon, 23 Jan 2023 18:47:21 -0800
Subject: [PATCH] Support dataframes as state vars (#324)

---
 pynecone/components/datadisplay/datatable.py   | 18 ++++++++++++++++--
 pynecone/state.py                              |  2 +-
 pynecone/utils.py                              | 12 ++++++++++++
 tests/components/graphing/test_victory_data.py |  9 +--------
 tests/test_state.py                            |  2 +-
 5 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/pynecone/components/datadisplay/datatable.py b/pynecone/components/datadisplay/datatable.py
index bc2028dfb..e7b4f6bd0 100644
--- a/pynecone/components/datadisplay/datatable.py
+++ b/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()
diff --git a/pynecone/state.py b/pynecone/state.py
index be99263ad..18b1341bf 100644
--- a/pynecone/state.py
+++ b/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, "
diff --git a/pynecone/utils.py b/pynecone/utils.py
index bc2bcb559..6b1b8bcf6 100644
--- a/pynecone/utils.py
+++ b/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.
 
diff --git a/tests/components/graphing/test_victory_data.py b/tests/components/graphing/test_victory_data.py
index f966347e3..e50db1cb6 100644
--- a/tests/components/graphing/test_victory_data.py
+++ b/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]
diff --git a/tests/test_state.py b/tests/test_state.py
index 9be355f08..39112014e 100644
--- a/tests/test_state.py
+++ b/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):