From 3b3b41550fa6286ddfdb4ec3e0bcbb14eb4a7c33 Mon Sep 17 00:00:00 2001
From: Joe Numainville <josephnumainville@deephaven.io>
Date: Wed, 13 Mar 2024 12:42:20 -0500
Subject: [PATCH] better check

---
 .../src/deephaven/plot/express/__init__.py               | 2 +-
 .../plot/express/deephaven_figure/DeephavenFigure.py     | 9 +++++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/plugins/plotly-express/src/deephaven/plot/express/__init__.py b/plugins/plotly-express/src/deephaven/plot/express/__init__.py
index 7ab97095f..cca071669 100644
--- a/plugins/plotly-express/src/deephaven/plot/express/__init__.py
+++ b/plugins/plotly-express/src/deephaven/plot/express/__init__.py
@@ -94,7 +94,7 @@ def create_client_connection(
         """
         if isinstance(obj, Figure):
             # this is a plotly figure, it will never be updated, so wrap once and send
-            obj = DeephavenFigure(obj)
+            obj = DeephavenFigure(obj, is_plotly_fig=True)
 
         figure_connection = DeephavenFigureConnection(obj, connection)
         initial_message = json.dumps(
diff --git a/plugins/plotly-express/src/deephaven/plot/express/deephaven_figure/DeephavenFigure.py b/plugins/plotly-express/src/deephaven/plot/express/deephaven_figure/DeephavenFigure.py
index 058f841f3..d70828519 100644
--- a/plugins/plotly-express/src/deephaven/plot/express/deephaven_figure/DeephavenFigure.py
+++ b/plugins/plotly-express/src/deephaven/plot/express/deephaven_figure/DeephavenFigure.py
@@ -394,6 +394,7 @@ def __init__(
         has_color: bool = False,
         trace_generator: Generator[dict[str, Any], None, None] | None = None,
         has_subplots: bool = False,
+        is_plotly_fig: bool = False,
     ):
         """
         Create a new DeephavenFigure
@@ -407,6 +408,7 @@ def __init__(
             has_color: If this figure has color
             trace_generator: The trace generator
             has_subplots: If this figure has subplots
+            is_plotly_fig: If this is a plotly figure
         """
         # keep track of function that called this, and it's args
         self._head_node = DeephavenHeadNode()
@@ -428,6 +430,8 @@ def __init__(
 
         self._has_subplots = has_subplots
 
+        self._is_plotly_fig = is_plotly_fig
+
         self._liveness_scope = LivenessScope()
 
     def copy_mappings(self: DeephavenFigure, offset: int = 0) -> list[DataMapping]:
@@ -584,8 +588,8 @@ def get_figure(self) -> DeephavenFigure | None:
         Returns:
             The figure
         """
-        if not self._head_node.get_figure():
-            # if there is no figure stored in the node, a plotly figure was passed directly
+        if self._is_plotly_fig:
+            # a plotly figure was passed directly
             # just return this figure since it will never be updated
             return self
         return self._head_node.get_figure()
@@ -680,6 +684,7 @@ def copy(self) -> DeephavenFigure:
             self._has_color,
             self._trace_generator,
             self._has_subplots,
+            self._is_plotly_fig,
         )
         new_figure._head_node = self._head_node.copy_graph()
         return new_figure