Repository: deephaven/deephaven-plugins (language: python)

Issue:
Hook to use execution context
from Bender on #168:
We probably should make a hook just for submitting tasks with the execution context, similar to how I made the use_render_queue hook for submitting tasks on the render thread.

Could be named use_execution_context


Discussion / hints:

Fix patch:
diff --git a/plugins/ui/src/deephaven/ui/hooks/__init__.py b/plugins/ui/src/deephaven/ui/hooks/__init__.py
index a3cc11bfc..b06b500be 100644
--- a/plugins/ui/src/deephaven/ui/hooks/__init__.py
+++ b/plugins/ui/src/deephaven/ui/hooks/__init__.py
@@ -9,6 +9,7 @@
 from .use_row_data import use_row_data
 from .use_row_list import use_row_list
 from .use_cell_data import use_cell_data
+from .use_execution_context import use_execution_context
 
 
 __all__ = [
@@ -23,4 +24,5 @@
     "use_row_data",
     "use_row_list",
     "use_cell_data",
+    "use_execution_context",
 ]
diff --git a/plugins/ui/src/deephaven/ui/hooks/use_execution_context.py b/plugins/ui/src/deephaven/ui/hooks/use_execution_context.py
new file mode 100644
index 000000000..b136a4048
--- /dev/null
+++ b/plugins/ui/src/deephaven/ui/hooks/use_execution_context.py
@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable
+
+from deephaven.execution_context import get_exec_ctx, ExecutionContext
+
+from . import use_memo
+
+
+def func_with_ctx(
+    exec_ctx: ExecutionContext,
+    func: Callable,
+) -> None:
+    """
+    Call the function within an execution context.
+
+    Args:
+        exec_ctx: ExecutionContext: The execution context to use.
+        func: Callable: The function to call.
+    """
+    with exec_ctx:
+        func()
+
+
+def use_execution_context(exec_ctx: ExecutionContext = None) -> None:
+    """
+    Create an execution context wrapper for a function.
+
+    Args:
+        exec_ctx: ExecutionContext: The execution context to use. Defaults to
+            the current execution context if not provided.
+    """
+    exec_ctx = use_memo(lambda: exec_ctx if exec_ctx else get_exec_ctx(), [exec_ctx])
+    exec_fn = use_memo(lambda: partial(func_with_ctx, exec_ctx), [exec_ctx])
+    return exec_fn


Test patch:
diff --git a/plugins/ui/test/deephaven/ui/test_hooks.py b/plugins/ui/test/deephaven/ui/test_hooks.py
index fe76fb5e2..41621d26d 100644
--- a/plugins/ui/test/deephaven/ui/test_hooks.py
+++ b/plugins/ui/test/deephaven/ui/test_hooks.py
@@ -490,6 +490,54 @@ def _test_cell_data(t=table):
 
         self.assertEqual(result, expected)
 
+    def test_execution_context(self):
+        from deephaven.ui.hooks import use_execution_context, use_state, use_memo
+        from deephaven import empty_table
+
+        def _test_execution_context():
+            with_exec_ctx = use_execution_context()
+
+            def table_func():
+                # This would fail if not using an execution context
+                empty_table(0).update("X=1")
+
+            def thread_func():
+                with_exec_ctx(table_func)
+
+            def start_thread():
+                thread = threading.Thread(target=thread_func)
+                thread.start()
+                thread.join()
+
+            use_memo(start_thread, [])
+
+        render_hook(_test_execution_context)
+
+    def test_execution_context_custom(self):
+        from deephaven.ui.hooks import use_execution_context, use_state, use_memo
+        from deephaven import empty_table
+        from deephaven.execution_context import make_user_exec_ctx
+
+        table = None
+
+        def _test_execution_context():
+            with_exec_ctx = use_execution_context(make_user_exec_ctx())
+
+            def table_func():
+                # This would fail if not using an execution context
+                empty_table(0).update("X=1")
+
+            def thread_func():
+                with_exec_ctx(table_func)
+
+            def start_thread():
+                thread = threading.Thread(target=thread_func)
+                thread.start()
+
+            use_memo(start_thread, [])
+
+        render_hook(_test_execution_context)
+
 
 if __name__ == "__main__":
     unittest.main()