From d1ff7d481f12c0d2fff69bc3847de7a6ebcae209 Mon Sep 17 00:00:00 2001
From: lawrence-axb <112918402+lawrence-axb@users.noreply.github.com>
Date: Sun, 5 Feb 2023 09:13:18 +0900
Subject: [PATCH] Allow plotly figures as state vars (#440)

---
 CONTRIBUTING.md     |  2 +-
 pynecone/utils.py   | 15 +++++++++++++--
 tests/test_state.py |  4 +++-
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index cf7c1dadc..3c28d0d50 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -93,6 +93,6 @@ find pynecone tests -name "*.py" -not -path pynecone/pc.py | xargs poetry run da
 ```
 Finally run `black` to format your code.
 ``` bash
-poetry run black pynecone
+poetry run black pynecone tests
 ```
 That's it you can now submit your pr. Thanks for contributing to Pynecone!
diff --git a/pynecone/utils.py b/pynecone/utils.py
index 1a7653056..6d2a26392 100644
--- a/pynecone/utils.py
+++ b/pynecone/utils.py
@@ -1100,6 +1100,18 @@ def is_dataframe(value: Type) -> bool:
     return value.__name__ == "DataFrame"
 
 
+def is_figure(value: Type) -> bool:
+    """Check if the given value is a figure.
+
+    Args:
+        value: The value to check.
+
+    Returns:
+        Whether the value is a figure.
+    """
+    return value.__name__ == "Figure"
+
+
 def is_valid_var_type(var: Type) -> bool:
     """Check if the given value is a valid prop type.
 
@@ -1109,7 +1121,7 @@ def is_valid_var_type(var: Type) -> bool:
     Returns:
         Whether the value is a valid prop type.
     """
-    return _issubclass(var, StateVar) or is_dataframe(var)
+    return _issubclass(var, StateVar) or is_dataframe(var) or is_figure(var)
 
 
 def format_state(value: Any) -> Dict:
@@ -1279,7 +1291,6 @@ def fix_events(events: Optional[List[Event]], token: str) -> List[Event]:
     # Fix the events created by the handler.
     out = []
     for e in events:
-
         # If it is already an event, don't modify it.
         if isinstance(e, Event):
             name = e.name
diff --git a/tests/test_state.py b/tests/test_state.py
index ea4930872..3d773e526 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -8,6 +8,7 @@ from pynecone.constants import RouteVar
 from pynecone.event import Event
 from pynecone.state import State
 from pynecone.var import BaseVar, ComputedVar
+from plotly.graph_objects import Figure
 
 
 class Object(Base):
@@ -27,6 +28,7 @@ class TestState(State):
     mapping: Dict[str, List[int]] = {"a": [1, 2, 3], "b": [4, 5, 6]}
     obj: Object = Object()
     complex: Dict[int, Object] = {1: Object(), 2: Object()}
+    fig: Figure = Figure()
 
     @ComputedVar
     def sum(self) -> float:
@@ -195,6 +197,7 @@ def test_class_vars(test_state):
         "complex",
         "sum",
         "upper",
+        "fig",
     }
 
 
@@ -622,7 +625,6 @@ def test_get_token(test_state):
 
 
 def test_get_current_page(test_state):
-
     assert test_state.get_current_page() == ""
 
     route = "mypage/subpage"