Move dg_stack into a ContextVar to support with blocks in separate threads

lib/streamlit/delta_generator.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import contextvars
 import sys
 from typing import (
     TYPE_CHECKING,
@@ -25,6 +26,7 @@
     Iterable,
     NoReturn,
     Optional,
+    Tuple,
     Type,
     TypeVar,
     cast,
@@ -48,7 +50,7 @@
 from streamlit.elements.doc_string import HelpMixin
 from streamlit.elements.empty import EmptyMixin
 from streamlit.elements.exception import ExceptionMixin
 from streamlit.elements.form import FormData, FormMixin, current_form_id
 from streamlit.elements.graphviz_chart import GraphvizMixin
 from streamlit.elements.heading import HeadingMixin
 from streamlit.elements.iframe import IframeMixin
@@ -147,6 +149,14 @@
             )
 
 
+# The dg_stack tracks the currently active DeltaGenerator, and is pushed to when
+# a DeltaGenerator is entered via a `with` block. This is implemented as a ContextVar
+# so that different threads or async tasks can have their own stacks.
+dg_stack: contextvars.ContextVar[Tuple[DeltaGenerator, ...]] = contextvars.ContextVar(
+    "dg_stack", default=tuple()
+)
+
+
 class DeltaGenerator(
     AlertMixin,
     BalloonsMixin,
@@ -282,9 +292,7 @@
 
     def __enter__(self) -> None:
         # with block started
-        ctx = get_script_run_ctx()
-        if ctx:
-            ctx.dg_stack.append(self)
+        dg_stack.set(dg_stack.get() + (self,))
 
     def __exit__(
         self,
@@ -293,9 +301,8 @@
         traceback: Any,
     ) -> Literal[False]:
         # with block ended
-        ctx = get_script_run_ctx()
-        if ctx is not None:
-            ctx.dg_stack.pop()
+
+        dg_stack.set(dg_stack.get()[:-1])
 
         # Re-raise any exceptions
         return False
@@ -310,9 +317,9 @@
         if self == self._main_dg:
             # We're being invoked via an `st.foo` pattern - use the current
             # `with` dg (aka the top of the stack).
-            ctx = get_script_run_ctx()
-            if ctx and len(ctx.dg_stack) > 0:
-                return ctx.dg_stack[-1]
+            current_stack = dg_stack.get()
+            if len(current_stack) > 0:
+                return current_stack[-1]
 
         # We're being invoked via an `st.sidebar.foo` pattern - ignore the
         # current `with` dg.

lib/streamlit/elements/form.py

@@ -14,20 +14,18 @@
 from __future__ import annotations
 
 import textwrap
-from typing import TYPE_CHECKING, NamedTuple, cast
+from typing import NamedTuple, cast
 
 from typing_extensions import Literal
 
 from streamlit import runtime
+from streamlit.delta_generator import DeltaGenerator, dg_stack
 from streamlit.errors import StreamlitAPIException
 from streamlit.proto import Block_pb2
 from streamlit.runtime.metrics_util import gather_metrics
 from streamlit.runtime.scriptrunner import ScriptRunContext, get_script_run_ctx
 from streamlit.runtime.state import WidgetArgs, WidgetCallback, WidgetKwargs
 
-if TYPE_CHECKING:
-    from streamlit.delta_generator import DeltaGenerator
-
 
 class FormData(NamedTuple):
     """Form data stored on a DeltaGenerator."""
@@ -52,11 +50,7 @@
     if this_dg == this_dg._main_dg:
         # We were created via an `st.foo` call.
         # Walk up the dg_stack to see if we're nested inside a `with st.form` statement.
-        ctx = get_script_run_ctx()
-        if ctx is None or len(ctx.dg_stack) == 0:
-            return None
-
-        for dg in reversed(ctx.dg_stack):
+        for dg in reversed(dg_stack.get()):
             if dg._form_data is not None:
                 return dg._form_data
     else:

lib/streamlit/runtime/scriptrunner/script_run_context.py

@@ -40,7 +40,9 @@ class ScriptRunContext:
     scoped to a single connected "session").
 
     ScriptRunContext is used internally by virtually every `st.foo()` function.
-    It is accessed only from the script thread that's created by ScriptRunner.
+    It is accessed only from the script thread that's created by ScriptRunner,
+    or from app-created helper threads that have been "attached" to the
+    ScriptRunContext via `add_script_run_ctx`.
 
     Streamlit code typically retrieves the active ScriptRunContext via the
     `get_script_run_ctx` function.
@@ -64,9 +66,6 @@ class ScriptRunContext:
     widget_user_keys_this_run: Set[str] = field(default_factory=set)
     form_ids_this_run: Set[str] = field(default_factory=set)
     cursors: Dict[int, "streamlit.cursor.RunningCursor"] = field(default_factory=dict)
-    dg_stack: List["streamlit.delta_generator.DeltaGenerator"] = field(
-        default_factory=list
-    )
     script_requests: Optional[ScriptRequests] = None
 
     def reset(self, query_string: str = "", page_script_hash: str = "") -> None: