From 76cd7abf65140191c649c36a6b296232702c3e78 Mon Sep 17 00:00:00 2001
From: Joe <josephnumainville@deephaven.io>
Date: Tue, 23 Jan 2024 09:39:29 -0600
Subject: [PATCH] feat: use_execution_context hook (#205)

The hook can be used to call code within an execution context

It is a fairly niche hook, but here as an example of it being used to do
a table update on a separate thread

```
import deephaven.ui as ui
import threading
from deephaven.execution_context import get_exec_ctx
from deephaven import agg as agg
import deephaven.pandas as dhpd
import deephaven.plot.express as dx

stocks = dx.data.stocks()

def print_sym(source):
    print(dhpd.to_pandas(source.update(["Symbol = sym"]) \
        .agg_by([agg.avg(cols=["AVG = size"])], by=["Symbol"])))

@ui.component
def exec_ctx_hook(source):

    with_exec_ctx = ui.use_execution_context()

    def thread_func():
        #print_sym(source)
        with_exec_ctx(lambda: print_sym(source))

    def start_thread():
        thread = threading.Thread(target=thread_func)
        thread.start()

    ui.use_memo(start_thread, source)

    return None

exec_hook = exec_ctx_hook(stocks)
```

Fixes #184
---
 plugins/ui/src/deephaven/ui/hooks/__init__.py |  2 +
 .../ui/hooks/use_execution_context.py         | 36 ++++++++++++++
 plugins/ui/test/deephaven/ui/test_hooks.py    | 48 +++++++++++++++++++
 3 files changed, 86 insertions(+)
 create mode 100644 plugins/ui/src/deephaven/ui/hooks/use_execution_context.py

diff --git a/plugins/ui/src/deephaven/ui/hooks/__init__.py b/plugins/ui/src/deephaven/ui/hooks/__init__.py
index 6c9eb7a60..55312852a 100644
--- a/plugins/ui/src/deephaven/ui/hooks/__init__.py
+++ b/plugins/ui/src/deephaven/ui/hooks/__init__.py
@@ -10,6 +10,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__ = [
@@ -25,4 +26,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 fc2133479..da05b9b73 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()