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 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()