EventContext

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

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import threading
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Optional,
+    Generator,
+)
+from contextlib import contextmanager
+from .NoContextException import NoContextException
+
+OnEventCallable = Callable[[str, Dict[str, Any]], None]
+"""
+Callable that is called when an event is queued up.
+"""
+
+_local_data = threading.local()
+
+
+def get_event_context() -> EventContext:
+    """
+    Gets the currently active context, or throws NoContextException if none is set.
+
+    Returns:
+        The active EventContext, or throws if none is present.
+    """
+    try:
+        return _local_data.context
+    except AttributeError:
+        raise NoContextException("No context set")
+
+
+def _set_event_context(context: Optional[EventContext]):
+    """
+    Set the current context for the thread. Can be set to None to unset the context for a thread.
+    """
+    if context is None:
+        del _local_data.context
+    else:
+        _local_data.context = context
+
+
+class EventContext:
+    _on_send_event: OnEventCallable
+    """
+    The callback to call when sending an event.
+    """
+
+    def __init__(
+        self,
+        on_send_event: OnEventCallable,
+    ):
+        """
+        Create a new event context.
+
+        Args:
+            on_send_event: The callback to call when sending an event.
+        """
+
+        self._on_send_event = on_send_event
+
+    @contextmanager
+    def open(self) -> Generator[EventContext, None, None]:
+        """
+        Opens this context.
+
+        Returns:
+            A context manager to manage EventContext resources.
+        """
+        old_context = get_event_context()
+        _set_event_context(self)
+        yield self
+        _set_event_context(old_context)
+
+    def send_event(self, name: str, params: dict) -> None:
+        """
+        Send an event to the client.
+
+        Args:
+            name: The name of the event.
+            params: The params of the event.
+        """
+        self._on_send_event(name, params)

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

@@ -0,0 +1,2 @@
+class NoContextException(Exception):
+    pass

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

@@ -20,6 +20,7 @@
 from deephaven.liveness_scope import LivenessScope
 from contextlib import contextmanager
 from dataclasses import dataclass
+from .NoContextException import NoContextException
 
 logger = logging.getLogger(__name__)
 
@@ -33,10 +34,6 @@
 Callable that is called when there is a change in the context (setting the state).
 """
 
-OnEventCallable = Callable[[str, dict], None]
-"""
-Callable that is called when an event is queued up.
-"""
 
 StateKey = int
 """
@@ -131,10 +128,6 @@ def _should_retain_value(value: ValueWithLiveness[T | None]) -> bool:
 _local_data = threading.local()
 
 
-class NoContextException(Exception):
-    pass
-
-
 def get_context() -> RenderContext:
     """
     Gets the currently active context, or throws NoContextException if none is set.
@@ -228,12 +221,7 @@ class RenderContext:
     Flag to indicate if this context is mounted. It is unusable after being unmounted.
     """
 
-    def __init__(
-        self,
-        on_change: OnChangeCallable,
-        on_queue_render: OnChangeCallable,
-        on_queue_event: OnEventCallable,
-    ):
+    def __init__(self, on_change: OnChangeCallable, on_queue_render: OnChangeCallable):
         """
         Create a new render context.
 
@@ -248,7 +236,6 @@ def __init__(
         self._children_context = {}
         self._on_change = on_change
         self._on_queue_render = on_queue_render
-        self._on_queue_event = on_queue_event
         self._collected_scopes = set()
         self._collected_effects = []
         self._collected_unmount_listeners = []
@@ -461,9 +448,7 @@ def get_child_context(self, key: ContextKey) -> "RenderContext":
         """
         logger.debug("Getting child context for key %s", key)
         if key not in self._children_context:
-            child_context = RenderContext(
-                self._on_change, self._on_queue_render, self._on_queue_event
-            )
+            child_context = RenderContext(self._on_change, self._on_queue_render)
             logger.debug(
                 "Created new child context %s for key %s in %s",
                 child_context,
@@ -497,16 +482,6 @@ def queue_render(self, update: Callable[[], None]) -> None:
         """
         self._on_queue_render(update)
 
-    def queue_event(self, name: str, params: dict) -> None:
-        """
-        Queue up an event to be sent to the client.
-
-        Args:
-            name: The name of the event.
-            params: The params of the event.
-        """
-        self._on_queue_event(name, params)
-
     def manage(self, liveness_scope: LivenessScope) -> None:
         """
         Indicates that the given LivenessScope must live until the end of the next

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

@@ -1,9 +1,13 @@
+from .EventContext import (
+    EventContext,
+    OnEventCallable,
+    get_event_context,
+)
 from .RenderContext import (
     RenderContext,
     StateKey,
     StateUpdateCallable,
     OnChangeCallable,
-    OnEventCallable,
     InitializerFunction,
     UpdaterFunction,
     get_context,