From 25d670fb63aa05874868e05c4a281c4fb294ed03 Mon Sep 17 00:00:00 2001
From: Chung Min Kim <chungmin99@gmail.com>
Date: Sun, 21 Apr 2024 15:30:06 -0700
Subject: [PATCH] ruff, type checks?

---
 src/viser/_gui_api.py     |  5 +++--
 src/viser/_gui_handles.py | 14 ++++++++------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/viser/_gui_api.py b/src/viser/_gui_api.py
index f30e3e64c..93110ad24 100644
--- a/src/viser/_gui_api.py
+++ b/src/viser/_gui_api.py
@@ -34,7 +34,6 @@
     get_origin,
     get_type_hints,
 )
-import plotly.graph_objects as go
 
 from . import _messages
 from ._gui_handles import (
@@ -46,8 +45,8 @@
     GuiFolderHandle,
     GuiInputHandle,
     GuiMarkdownHandle,
-    GuiPlotlyHandle,
     GuiModalHandle,
+    GuiPlotlyHandle,
     GuiTabGroupHandle,
     GuiUploadButtonHandle,
     SupportsRemoveProtocol,
@@ -62,6 +61,8 @@
 from ._messages import FileTransferPartAck
 
 if TYPE_CHECKING:
+    import plotly.graph_objects as go
+
     from .infra import ClientId
 
 IntOrFloat = TypeVar("IntOrFloat", int, float)
diff --git a/src/viser/_gui_handles.py b/src/viser/_gui_handles.py
index 11018f66a..fc57fb6be 100644
--- a/src/viser/_gui_handles.py
+++ b/src/viser/_gui_handles.py
@@ -24,7 +24,6 @@
 import imageio.v3 as iio
 import numpy as onp
 from typing_extensions import Protocol
-import plotly.graph_objects as go
 
 from ._icons import svg_from_icon
 from ._icons_enum import IconName
@@ -33,6 +32,8 @@
 from .infra import ClientId
 
 if TYPE_CHECKING:
+    import plotly.graph_objects as go
+
     from ._gui_api import GuiApi
     from ._viser import ClientHandle
 
@@ -624,13 +625,12 @@ class GuiPlotlyHandle:
     _visible: bool
     _container_id: str  # Parent.
     _order: float
-    _figure: Optional[go.Figure]
-    _aspect_ratio: Optional[float]
+    _figure: go.Figure
+    _aspect_ratio: float
 
     @property
-    def figure(self) -> str:
+    def figure(self) -> go.Figure:
         """Current content of this markdown element. Synchronized automatically when assigned."""
-        assert self._figure is not None
         return self._figure
 
     @figure.setter
@@ -660,7 +660,9 @@ def aspect_ratio(self, aspect_ratio: float) -> None:
 
     def plot_to_json(self) -> str:
         """Convert the plotly figure to an HTML string."""
-        return self._figure.to_json()
+        json_str = self._figure.to_json()
+        assert isinstance(json_str, str)
+        return json_str
 
     @property
     def order(self) -> float: