Cleanup based on review

plugins/ui/src/deephaven/ui/_internal/RenderContext.py

@@ -5,13 +5,13 @@
 from typing import (
     Any,
     Callable,
+    Dict,
     Optional,
     TypeVar,
     Union,
     Generator,
     Generic,
     cast,
-    Set,
 )
 from functools import partial
 from deephaven import DHError
@@ -53,6 +53,11 @@
 The key for a child context.
 """
 
+ExportedRenderState = Dict[str, Any]
+"""
+The serializable state of a RenderContext. Used to serialize the state for the client.
+"""
+
 
 @dataclass
 class ValueWithLiveness(Generic[T]):
@@ -368,11 +373,11 @@ def manage(self, liveness_scope: LivenessScope) -> None:
         assert self is get_context()
         self._collected_scopes.add(cast(LivenessScope, liveness_scope.j_scope))
 
-    def export_state(self) -> dict[str, Any]:
+    def export_state(self) -> ExportedRenderState:
         """
         Export the state of this context. This is used to serialize the state for the client.
         """
-        exported_state: dict[str, Any] = {}
+        exported_state: ExportedRenderState = {}
         # We need to iterate through all of our state and export anything that doesn't have a LivenessScope right now (anything serializable)
         state = {
             key: value.value

plugins/ui/src/deephaven/ui/_internal/__init__.py

@@ -8,6 +8,7 @@
     get_context,
     NoContextException,
     ValueWithLiveness,
+    ExportedRenderState,
 )
 from .utils import (
     get_component_name,

plugins/ui/src/deephaven/ui/object_types/ElementMessageStream.py

@@ -16,7 +16,7 @@
 from .._internal import wrap_callable
 from ..elements import Element
 from ..renderer import NodeEncoder, Renderer, RenderedNode
-from .._internal import RenderContext, StateUpdateCallable
+from .._internal import RenderContext, StateUpdateCallable, ExportedRenderState
 
 logger = logging.getLogger(__name__)
 
@@ -155,12 +155,12 @@ def _render(self) -> None:
 
         try:
             node = self._renderer.render(self._element)
+            state = self._context.export_state()
+            self._send_document_update(node, state)
         except Exception as e:
             logger.exception("Error rendering %s", self._element.name)
             raise e
 
-        self._send_document_update(node, self._context)
-
     def _process_callable_queue(self) -> None:
         """
         Process any queued callables, then re-renders the element if it is dirty.
@@ -310,7 +310,7 @@ def _make_dispatcher(self) -> Dispatcher:
         dispatcher["setState"] = self._set_state
         return dispatcher
 
-    def _set_state(self, state: dict[str, Any]) -> None:
+    def _set_state(self, state: ExportedRenderState) -> None:
         """
         Set the state of the element. This is called by the client on initial load.
 
@@ -321,7 +321,9 @@ def _set_state(self, state: dict[str, Any]) -> None:
         self._context.import_state(state)
         self._mark_dirty()
 
-    def _send_document_update(self, root: RenderedNode, context: RenderContext) -> None:
+    def _send_document_update(
+        self, root: RenderedNode, state: ExportedRenderState
+    ) -> None:
         """
         Send a document update to the client. Currently just sends the entire document for each update.
 
@@ -336,7 +338,6 @@ def _send_document_update(self, root: RenderedNode, context: RenderContext) -> N
         new_objects = encoder_result["new_objects"]
         callable_id_dict = encoder_result["callable_id_dict"]
 
-        state = context.export_state()
         logger.debug("Exported state: %s", state)
         encoded_state = json.dumps(state)
 
@@ -352,15 +353,3 @@ def _send_document_update(self, root: RenderedNode, context: RenderContext) -> N
             dispatcher[callable_id] = wrap_callable(callable)
         self._dispatcher = dispatcher
         self._connection.on_data(payload.encode(), new_objects)
-
-    def send_state_update(self, state: RenderContext) -> None:
-        """
-        Send a state update to the client. Currently just sends the entire state for each update.
-
-        Args:
-            state: The state to send
-        """
-        request = self._make_notification("stateUpdated", json.dumps(state))
-        payload = json.dumps(request)
-        logger.debug(f"Sending payload: {payload}")
-        self._connection.on_data(payload.encode(), [])