streamlit-nightly 1.36.1.dev20240714__py2.py3-none-any.whl → 1.36.1.dev20240716__py2.py3-none-any.whl

streamlit/runtime/context.py

@@ -0,0 +1,157 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2024)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from functools import lru_cache
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, Mapping, cast
+
+from streamlit import runtime
+from streamlit.runtime.metrics_util import gather_metrics
+from streamlit.runtime.scriptrunner import get_script_run_ctx
+from streamlit.type_util import is_type
+
+if TYPE_CHECKING:
+    from http.cookies import Morsel
+
+    from tornado.httputil import HTTPHeaders, HTTPServerRequest
+    from tornado.web import RequestHandler
+
+
+def _get_request() -> HTTPServerRequest | None:
+    ctx = get_script_run_ctx()
+    if ctx is None:
+        return None
+
+    session_client = runtime.get_instance().get_client(ctx.session_id)
+    if session_client is None:
+        return None
+
+    # We return websocket request only if session_client is an instance of
+    # BrowserWebSocketHandler (which is True for the Streamlit open-source
+    # implementation). For any other implementation, we return None.
+    if not is_type(
+        session_client,
+        "streamlit.web.server.browser_websocket_handler.BrowserWebSocketHandler",
+    ):
+        return None
+    return cast("RequestHandler", session_client).request
+
+
+@lru_cache
+def _normalize_header(name: str) -> str:
+    """Map a header name to Http-Header-Case.
+
+    >>> _normalize_header("coNtent-TYPE")
+    'Content-Type'
+    """
+    return "-".join(w.capitalize() for w in name.split("-"))
+
+
+class StreamlitHeaders(Mapping[str, str]):
+    def __init__(self, headers: Iterable[tuple[str, str]]):
+        dict_like_headers: dict[str, list[str]] = {}
+
+        for key, value in headers:
+            header_value = dict_like_headers.setdefault(_normalize_header(key), [])
+            header_value.append(value)
+
+        self._headers = dict_like_headers
+
+    @classmethod
+    def from_tornado_headers(cls, tornado_headers: HTTPHeaders) -> StreamlitHeaders:
+        return cls(tornado_headers.get_all())
+
+    def get_all(self, key: str) -> list[str]:
+        return list(self._headers.get(_normalize_header(key), []))
+
+    def __getitem__(self, key: str) -> str:
+        try:
+            return self._headers[_normalize_header(key)][0]
+        except LookupError:
+            raise KeyError(key) from None
+
+    def __len__(self) -> int:
+        """Number of unique headers present in request."""
+        return len(self._headers)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._headers)
+
+    def to_dict(self) -> dict[str, str]:
+        return {key: self[key] for key in self}
+
+
+class StreamlitCookies(Mapping[str, str]):
+    def __init__(self, cookies: Mapping[str, str]):
+        self._cookies = MappingProxyType(cookies)
+
+    @classmethod
+    def from_tornado_cookies(
+        cls, tornado_cookies: dict[str, Morsel[Any]]
+    ) -> StreamlitCookies:
+        dict_like_cookies = {}
+        for key, morsel in tornado_cookies.items():
+            dict_like_cookies[key] = morsel.value
+        return cls(dict_like_cookies)
+
+    def __getitem__(self, key: str) -> str:
+        return self._cookies[key]
+
+    def __len__(self) -> int:
+        """Number of unique headers present in request."""
+        return len(self._cookies)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._cookies)
+
+    def to_dict(self) -> dict[str, str]:
+        return dict(self._cookies)
+
+
+class ContextProxy:
+    """An interface to context about the current user session. Context is exposed as
+    properties on `st.context`, such as `st.context.headers` or `st.context.cookies`.
+    Each property description explains more about the contents."""
+
+    @property
+    @gather_metrics("context.headers")
+    def headers(self) -> StreamlitHeaders:
+        """A read-only, dict-like access to headers sent in the initial request.
+        Keys are case-insensitive. Use get_all() to see all values if the same header
+        is set multiple times.
+        """
+        # We have a docstring in line above as one-liner, to have a correct docstring
+        # in the st.write(st,context) call.
+        session_client_request = _get_request()
+
+        if session_client_request is None:
+            return StreamlitHeaders({})
+
+        return StreamlitHeaders.from_tornado_headers(session_client_request.headers)
+
+    @property
+    @gather_metrics("context.cookies")
+    def cookies(self) -> StreamlitCookies:
+        """A read-only, dict-like access to cookies sent in the initial request."""
+        # We have a docstring in line above as one-liner, to have a correct docstring
+        # in the st.write(st,context) call.
+        session_client_request = _get_request()
+
+        if session_client_request is None:
+            return StreamlitCookies({})
+
+        cookies = session_client_request.cookies
+        return StreamlitCookies.from_tornado_cookies(cookies)

streamlit/runtime/forward_msg_queue.py

@@ -98,6 +98,7 @@ class ForwardMsgQueue:
                 for msg in self._queue
                 if msg.WhichOneof("type")
                 in {
+                    "new_session",
                     "script_finished",
                     "session_status_changed",
                     "parent_message",

streamlit/runtime/fragment.py

@@ -22,22 +22,24 @@ from copy import deepcopy
 from functools import wraps
 from typing import TYPE_CHECKING, Any, Callable, Protocol, TypeVar, overload
 
+from streamlit.error_util import handle_uncaught_app_exception
+from streamlit.errors import FragmentHandledException, FragmentStorageKeyError
 from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
 from streamlit.runtime.metrics_util import gather_metrics
 from streamlit.runtime.scriptrunner import get_script_run_ctx
-from streamlit.runtime.scriptrunner.exec_code import exec_func_with_error_handling
+from streamlit.runtime.scriptrunner.exceptions import RerunException, StopException
 from streamlit.time_util import time_to_seconds
 
 if TYPE_CHECKING:
     from datetime import timedelta
 
+
 F = TypeVar("F", bound=Callable[..., Any])
 Fragment = Callable[[], Any]
 
 
 class FragmentStorage(Protocol):
-    """A key-value store for Fragments. Used to implement the @st.experimental_fragment
-    decorator.
+    """A key-value store for Fragments. Used to implement the @st.fragment decorator.
 
     We intentionally define this as its own protocol despite how generic it appears to
     be at first glance. The reason why is that, in any case where fragments aren't just
@@ -47,6 +49,14 @@ class FragmentStorage(Protocol):
     protocols.
     """
 
+    # Weirdly, we have to define this above the `set` method, or mypy gets it confused
+    # with the `set` type of `new_fragments_ids`.
+    @abstractmethod
+    def clear(self, new_fragment_ids: set[str] | None = None) -> None:
+        """Remove all fragments saved in this FragmentStorage unless listed in
+        new_fragment_ids."""
+        raise NotImplementedError
+
     @abstractmethod
     def get(self, key: str) -> Fragment:
         """Returns the stored fragment for the given key."""
@@ -63,8 +73,8 @@ class FragmentStorage(Protocol):
         raise NotImplementedError
 
     @abstractmethod
-    def clear(self) -> None:
-        """Remove all fragments saved in this FragmentStorage."""
+    def contains(self, key: str) -> bool:
+        """Return whether the given key is present in this FragmentStorage."""
         raise NotImplementedError
 
 
@@ -83,21 +93,42 @@ class MemoryFragmentStorage(FragmentStorage):
     def __init__(self):
         self._fragments: dict[str, Fragment] = {}
 
+    # Weirdly, we have to define this above the `set` method, or mypy gets it confused
+    # with the `set` type of `new_fragments_ids`.
+    def clear(self, new_fragment_ids: set[str] | None = None) -> None:
+        if new_fragment_ids is None:
+            new_fragment_ids = set()
+
+        fragment_ids = list(self._fragments.keys())
+
+        for fid in fragment_ids:
+            if fid not in new_fragment_ids:
+                del self._fragments[fid]
+
     def get(self, key: str) -> Fragment:
-        return self._fragments[key]
+        try:
+            return self._fragments[key]
+        except KeyError as e:
+            raise FragmentStorageKeyError(str(e))
 
     def set(self, key: str, value: Fragment) -> None:
         self._fragments[key] = value
 
     def delete(self, key: str) -> None:
-        del self._fragments[key]
+        try:
+            del self._fragments[key]
+        except KeyError as e:
+            raise FragmentStorageKeyError(str(e))
 
-    def clear(self) -> None:
-        self._fragments.clear()
+    def contains(self, key: str) -> bool:
+        return key in self._fragments
 
 
 def _fragment(
-    func: F | None = None, *, run_every: int | float | timedelta | str | None = None
+    func: F | None = None,
+    *,
+    run_every: int | float | timedelta | str | None = None,
+    additional_hash_info: str = "",
 ) -> Callable[[F], F] | F:
     """Contains the actual fragment logic.
 
@@ -128,10 +159,9 @@ def _fragment(
 
         cursors_snapshot = deepcopy(ctx.cursors)
         dg_stack_snapshot = deepcopy(dg_stack.get())
-        active_dg = dg_stack_snapshot[-1]
         h = hashlib.new("md5")
         h.update(
-            f"{non_optional_func.__module__}.{non_optional_func.__qualname__}{active_dg._get_delta_path_str()}".encode()
+            f"{non_optional_func.__module__}.{non_optional_func.__qualname__}{dg_stack_snapshot[-1]._get_delta_path_str()}{additional_hash_info}".encode()
         )
         fragment_id = h.hexdigest()
 
@@ -149,18 +179,22 @@ def _fragment(
             ctx = get_script_run_ctx(suppress_warning=True)
             assert ctx is not None
 
-            if ctx.fragment_ids_this_run:
+            if ctx.script_requests and ctx.script_requests.fragment_id_queue:
                 # This script run is a run of one or more fragments. We restore the
                 # state of ctx.cursors and dg_stack to the snapshots we took when this
                 # fragment was declared.
                 ctx.cursors = deepcopy(cursors_snapshot)
                 dg_stack.set(deepcopy(dg_stack_snapshot))
-            else:
-                # Otherwise, we must be in a full script run. We need to temporarily set
-                # ctx.current_fragment_id so that elements corresponding to this
-                # fragment get tagged with the appropriate ID. ctx.current_fragment_id
-                # gets reset after the fragment function finishes running.
-                ctx.current_fragment_id = fragment_id
+
+            # Always add the fragment id to new_fragment_ids. For full app runs
+            # we need to add them anyways and for fragment runs we add them
+            # in case the to-be-executed fragment id was cleared from the storage
+            # by the full app run.
+            ctx.new_fragment_ids.add(fragment_id)
+            # Set ctx.current_fragment_id so that elements corresponding to this
+            # fragment get tagged with the appropriate ID. ctx.current_fragment_id gets
+            # reset after the fragment function finishes running.
+            ctx.current_fragment_id = fragment_id
 
             try:
                 # Make sure we set the active script hash to the same value
@@ -174,18 +208,48 @@ def _fragment(
                     if initialized_active_script_hash != ctx.active_script_hash
                     else contextlib.nullcontext()
                 )
+                result = None
                 with active_hash_context:
                     with st.container():
-                        ctx.current_fragment_delta_path = (
-                            active_dg._cursor.delta_path if active_dg._cursor else []
-                        )
-                        result = non_optional_func(*args, **kwargs)
+                        try:
+                            # use dg_stack instead of active_dg to have correct copy
+                            # during execution (otherwise we can run into concurrency
+                            # issues with multiple fragments). Use dg_stack because we
+                            # just entered a container and [:-1] of the delta path
+                            # because thats the prefix of the fragment,
+                            # e.g. [0, 3, 0] -> [0, 3].
+                            # All fragment elements start with [0, 3].
+                            active_dg = dg_stack.get()[-1]
+                            ctx.current_fragment_delta_path = (
+                                active_dg._cursor.delta_path
+                                if active_dg._cursor
+                                else []
+                            )[:-1]
+                            result = non_optional_func(*args, **kwargs)
+                        except (
+                            RerunException,
+                            StopException,
+                        ) as e:
+                            # The wrapped_fragment function is executed
+                            # inside of a exec_func_with_error_handling call, so
+                            # there is a correct handler for these exceptions.
+                            raise e
+                        except Exception as e:
+                            # render error here so that the delta path is correct
+                            # for full app runs, the error will be displayed by the
+                            # main code handler
+                            # if not is_full_app_run:
+                            handle_uncaught_app_exception(e)
+                            # raise here again in case we are in full app execution
+                            # and some flags have to be set
+                            raise FragmentHandledException(e)
+                    return result
             finally:
                 ctx.current_fragment_id = None
+                ctx.current_fragment_delta_path = []
 
-            return result
-
-        ctx.fragment_storage.set(fragment_id, wrapped_fragment)
+        if not ctx.fragment_storage.contains(fragment_id):
+            ctx.fragment_storage.set(fragment_id, wrapped_fragment)
 
         if run_every:
             msg = ForwardMsg()
@@ -193,15 +257,8 @@ def _fragment(
             msg.auto_rerun.fragment_id = fragment_id
             ctx.enqueue(msg)
 
-        # Wrap the fragment function in the same try-except block as in a normal
-        # script_run so that for a main-app run (this execution) and a fragment-rerun
-        # the same execution and error-handling logic is used. This makes errors in the
-        # fragment appear in the fragment path also for the first execution here in
-        # context of a full app run.
-        result, _, _, _ = exec_func_with_error_handling(
-            wrapped_fragment, ctx, reraise_rerun_exception=True
-        )
-        return result
+        # Immediate execute the wrapped fragment since we are in a full app run
+        return wrapped_fragment()
 
     with contextlib.suppress(AttributeError):
         # Make this a well-behaved decorator by preserving important function
@@ -230,7 +287,7 @@ def fragment(
 ) -> Callable[[F], F]: ...
 
 
-@gather_metrics("experimental_fragment")
+@gather_metrics("fragment")
 def fragment(
     func: F | None = None,
     *,
@@ -297,14 +354,14 @@ def fragment(
     Examples
     --------
     The following example demonstrates basic usage of
-    ``@st.experimental_fragment``. As an anology, "inflating balloons" is a
-    slow process that happens outside of the fragment. "Releasing balloons" is
-    a quick process that happens inside of the fragment.
+    ``@st.fragment``. As an analogy, "inflating balloons" is a slow process that happens
+    outside of the fragment. "Releasing balloons" is a quick process that happens inside
+    of the fragment.
 
     >>> import streamlit as st
     >>> import time
     >>>
-    >>> @st.experimental_fragment
+    >>> @st.fragment
     >>> def release_the_balloons():
     >>>     st.button("Release the balloons", help="Fragment rerun")
     >>>     st.balloons()
@@ -332,14 +389,14 @@ def fragment(
     >>>     st.session_state.app_runs = 0
     >>>     st.session_state.fragment_runs = 0
     >>>
-    >>> @st.experimental_fragment
-    >>> def fragment():
+    >>> @st.fragment
+    >>> def my_fragment():
     >>>     st.session_state.fragment_runs += 1
     >>>     st.button("Rerun fragment")
     >>>     st.write(f"Fragment says it ran {st.session_state.fragment_runs} times.")
     >>>
     >>> st.session_state.app_runs += 1
-    >>> fragment()
+    >>> my_fragment()
     >>> st.button("Rerun full app")
     >>> st.write(f"Full app says it ran {st.session_state.app_runs} times.")
     >>> st.write(f"Full app sees that fragment ran {st.session_state.fragment_runs} times.")
@@ -356,7 +413,7 @@ def fragment(
     >>> if "clicks" not in st.session_state:
     >>>     st.session_state.clicks = 0
     >>>
-    >>> @st.experimental_fragment
+    >>> @st.fragment
     >>> def count_to_five():
     >>>     if st.button("Plus one!"):
     >>>         st.session_state.clicks += 1
@@ -376,3 +433,31 @@ def fragment(
 
     """
     return _fragment(func, run_every=run_every)
+
+
+@overload
+def experimental_fragment(
+    func: F,
+    *,
+    run_every: int | float | timedelta | str | None = None,
+) -> F: ...
+
+
+# Support being able to pass parameters to this decorator (that is, being able to write
+# `@fragment(run_every=5.0)`).
+@overload
+def experimental_fragment(
+    func: None = None,
+    *,
+    run_every: int | float | timedelta | str | None = None,
+) -> Callable[[F], F]: ...
+
+
+@gather_metrics("experimental_fragment")
+def experimental_fragment(
+    func: F | None = None,
+    *,
+    run_every: int | float | timedelta | str | None = None,
+) -> Callable[[F], F] | F:
+    """Deprecated alias for @st.fragment. See the docstring for the decorator's new name."""
+    return _fragment(func, run_every=run_every)

streamlit/runtime/metrics_util.py

@@ -482,6 +482,8 @@ def create_page_profile_message(
         page_profile.uncaught_exception = uncaught_exception
 
     if ctx := get_script_run_ctx():
-        page_profile.is_fragment_run = bool(ctx.fragment_ids_this_run)
+        page_profile.is_fragment_run = bool(
+            ctx.script_requests and ctx.script_requests.fragment_id_queue
+        )
 
     return msg

streamlit/runtime/scriptrunner/exec_code.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, Callable
 
 from streamlit.error_util import handle_uncaught_app_exception
+from streamlit.errors import FragmentHandledException
 from streamlit.runtime.scriptrunner.exceptions import RerunException, StopException
 
 if TYPE_CHECKING:
@@ -25,10 +26,7 @@ if TYPE_CHECKING:
 
 
 def exec_func_with_error_handling(
-    func: Callable[[], None],
-    ctx: ScriptRunContext,
-    *,
-    reraise_rerun_exception: bool = False,
+    func: Callable[[], None], ctx: ScriptRunContext
 ) -> tuple[Any | None, bool, RerunData | None, bool]:
     """Execute the passed function wrapped in a try/except block.
 
@@ -43,10 +41,6 @@ def exec_func_with_error_handling(
         The function to execute wrapped in the try/except block.
     ctx : ScriptRunContext
         The context in which the script is being run.
-    reraise_rerun_exception : bool, default False
-        If True, an occuring RerunException will be raised instead of handled. This can
-        be used if this function is called outside of the script_run context and we want
-        the script_runner to react on the rerun exception.
 
     Returns
     -------
@@ -70,16 +64,6 @@ def exec_func_with_error_handling(
     # is interrupted by a RerunException.
     rerun_exception_data: RerunData | None = None
 
-    # Saving and restoring our original cursors/dg_stack is needed
-    # specifically to handle the case where a RerunException is raised while
-    # running a fragment. In this case, we need to restore both to their states
-    # at the start of the script run to ensure that we write to the correct
-    # places in the app during the rerun (without this, ctx.cursors and dg_stack
-    # will still be set to the snapshots they were restored from when running
-    # the fragment).
-    original_cursors = ctx.cursors
-    original_dg_stack = dg_stack.get()
-
     # If the script stops early, we don't want to remove unseen widgets,
     # so we track this to potentially skip session state cleanup later.
     premature_stop: bool = False
@@ -90,22 +74,17 @@ def exec_func_with_error_handling(
     try:
         result = func()
     except RerunException as e:
-        if reraise_rerun_exception:
-            raise e
-
         rerun_exception_data = e.rerun_data
-        if rerun_exception_data.fragment_id_queue:
-            # This is a fragment-specific rerun, so we need to restore the stack
-            ctx.cursors = original_cursors
-            dg_stack.set(original_dg_stack)
-        else:
-            # If it is a full-app rerun, the stack needs to be refreshed.
-            # We should land here when `st.rerun` is called from within a
-            # fragment. Since we re-use the same thread, we have to clear the
-            # stack or otherwise we might render the main app in the old
-            # fragment's dg_stack.
-            ctx.cursors.clear()
-            dg_stack.set(get_default_dg_stack())
+
+        # Since the script is about to rerun, we may need to reset our cursors/dg_stack
+        # so that we write to the right place in the app. For full script runs, this
+        # needs to happen in case the same thread reruns our script (a different thread
+        # would automatically come with fresh cursors/dg_stack values). For fragments,
+        # it doesn't matter either way since the fragment resets these values from its
+        # snapshot before execution.
+        ctx.cursors.clear()
+        dg_stack.set(get_default_dg_stack())
+
         # Interruption due to a rerun is usually from `st.rerun()`, which
         # we want to count as a script completion so triggers reset.
         # It is also possible for this to happen if fast reruns is off,
@@ -116,11 +95,12 @@ def exec_func_with_error_handling(
         # This is thrown when the script executes `st.stop()`.
         # We don't have to do anything here.
         premature_stop = True
-
+    except FragmentHandledException:
+        run_without_errors = False
+        premature_stop = True
     except Exception as ex:
         run_without_errors = False
-        uncaught_exception = ex
-        handle_uncaught_app_exception(uncaught_exception)
         premature_stop = True
+        handle_uncaught_app_exception(ex)
 
     return result, run_without_errors, rerun_exception_data, premature_stop