event encoder

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

@@ -26,6 +26,7 @@
     ExportedRenderState,
     EventContext,
 )
+from .EventEncoder import EventEncoder
 from .ErrorCode import ErrorCode
 
 logger = logging.getLogger(__name__)
@@ -68,6 +69,11 @@ class ElementMessageStream(MessageStream):
     Encoder to use to encode the document.
     """
 
+    _event_encoder: EventEncoder
+    """
+    Encoder to use to encode events.
+    """
+
     _message_id: int
     """
     The next message ID to use.
@@ -172,6 +178,9 @@ def __init__(self, element: Element, connection: MessageStream):
         self._manager = JSONRPCResponseManager()
         self._dispatcher = self._make_dispatcher()
         self._encoder = NodeEncoder(separators=(",", ":"))
+        self._event_encoder = EventEncoder(
+            self._serialize_callables, separators=(",", ":")
+        )
         self._context = RenderContext(self._queue_state_update, self._queue_callable)
         self._event_context = EventContext(self._send_event)
         self._renderer = Renderer(self._context)
@@ -383,6 +392,24 @@ def _set_state(self, state: ExportedRenderState) -> None:
         self._context.import_state(state)
         self._mark_dirty()
 
+    def _serialize_callables(self, node: Any) -> Any:
+        """
+        Serialize a callable.
+
+        Args:
+            node: The node to serialize
+        """
+        if callable(node):
+            new_id = f"tempCb{self._next_temp_callable_id}"
+            self._next_temp_callable_id += 1
+            self._temp_callable_dict[new_id] = node
+            return {
+                CALLABLE_KEY: new_id,
+            }
+        raise TypeError(
+            f"A Deephaven UI callback returned a non-serializable value. Object of type {type(node).__name__} is not JSON serializable"
+        )
+
     def _call_callable(self, callable_id: str, args: Any) -> Any:
         """
         Call a callable by its ID.
@@ -401,20 +428,8 @@ def _call_callable(self, callable_id: str, args: Any) -> Any:
             return
         result = fn(*args)
 
-        def serialize_callables(node: Any) -> Any:
-            if callable(node):
-                new_id = f"tempCb{self._next_temp_callable_id}"
-                self._next_temp_callable_id += 1
-                self._temp_callable_dict[new_id] = node
-                return {
-                    CALLABLE_KEY: new_id,
-                }
-            raise TypeError(
-                f"A Deephaven UI callback returned a non-serializable value. Object of type {type(node).__name__} is not JSON serializable"
-            )
-
         try:
-            return json.dumps(result, default=serialize_callables)
+            return json.dumps(result, default=self._serialize_callables)
         except Exception as e:
             # This is shown to the user in the Python console
             # The stack trace from logger.exception is useless to the user
@@ -501,11 +516,7 @@ def _send_event(self, name: str, params: dict) -> None:
             name: The name of the event
             params: The params of the event
         """
-        encoded_params, callable_id_dict = self._encoder.encode_event_params(params)
-        # Register any new callables
-        for callable, callable_id in callable_id_dict.items():
-            logger.debug("Registering callable %s", callable_id)
-            self._temp_callable_dict[callable_id] = wrap_callable(callable)
+        encoded_params = self._event_encoder.encode(params)
         request = self._make_notification("event", name, encoded_params)
         payload = json.dumps(request)
         self._connection.on_data(payload.encode(), [])

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

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import json
+from typing import Any, Callable
+
+
+class EventEncoder(json.JSONEncoder):
+    """
+    Encode an event in JSON.
+    """
+
+    _convert_callable: Callable[[Any], Any]
+    """
+    Function that will be called to serialize callables.
+    """
+
+    def __init__(
+        self,
+        convert_callable: Callable[[Any], Any],
+        *args: Any,
+        **kwargs: Any,
+    ):
+        """
+        Create a new EventEncoder.
+
+        Args:
+            convert_callable: A function that will be called to serialize callables
+            *args: Arguments to pass to the JSONEncoder constructor
+            **kwargs: Args to pass to the JSONEncoder constructor
+        """
+        super().__init__(*args, **kwargs)
+        self._convert_callable = convert_callable
+
+    def default(self, o: Any):
+        if callable(o):
+            return self._convert_callable(o)
+        else:
+            return super().default(o)

plugins/ui/src/deephaven/ui/renderer/NodeEncoder.py

@@ -2,7 +2,7 @@
 
 import json
 import logging
-from typing import Any, Callable, TypedDict, Tuple
+from typing import Any, Callable, TypedDict
 from weakref import WeakKeyDictionary
 from .RenderedNode import RenderedNode
 
@@ -152,18 +152,6 @@ def encode_node(self, node: RenderedNode) -> NodeEncoderResult:
             "callable_id_dict": self._callable_dict,
         }
 
-    def encode_event_params(
-        self, params: dict[str, Any]
-    ) -> Tuple[str, WeakKeyDictionary[Callable[..., Any], CallableId]]:
-        """
-        Encode the event parameters.
-
-        Args:
-            params: The parameters to encode
-        """
-        encoded_params = super().encode(params)
-        return encoded_params, self._callable_dict
-
     def _convert_rendered_node(self, node: RenderedNode):
         result: dict[str, Any] = {ELEMENT_KEY: node.name}
         if node.props is not None: