# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026) # # 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 import time from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast from typing_extensions import Self from streamlit.delta_generator import DeltaGenerator from streamlit.elements.lib.layout_utils import ( EXPANDABLE_TYPE_TO_PROTO_MAPPING, ExpandableType, WidthWithoutContent, get_width_config, validate_width, ) from streamlit.errors import StreamlitValueError from streamlit.proto.Block_pb2 import Block as BlockProto from streamlit.proto.ForwardMsg_pb2 import ForwardMsg from streamlit.runtime.scriptrunner_utils.script_run_context import enqueue_message if TYPE_CHECKING: from types import TracebackType from streamlit.cursor import Cursor States: TypeAlias = Literal["running", "complete", "error"] # Maps each state to the icon the default and compact styles render. Also # defines the valid `state` values, which both validation sites check against. _STATE_ICONS: Final[dict[States, str]] = { "running": "spinner", "complete": ":material/check:", "error": ":material/error:", } # Semantic progress state sent with every status container. The step style uses # it for the icon and the screen-reader announcement; AppTest reads it instead # of inferring the state from the icon. _STATE_PROTO_VALUES: Final[dict[States, BlockProto.Expandable.State.ValueType]] = { "running": BlockProto.Expandable.State.RUNNING, "complete": BlockProto.Expandable.State.COMPLETE, "error": BlockProto.Expandable.State.ERROR, } class StatusContainer(DeltaGenerator): @staticmethod def _create( parent: DeltaGenerator, label: str, expanded: bool = False, state: States = "running", type: ExpandableType = "default", width: WidthWithoutContent = "stretch", ) -> StatusContainer: if type not in EXPANDABLE_TYPE_TO_PROTO_MAPPING: raise StreamlitValueError( "type", [repr(name) for name in EXPANDABLE_TYPE_TO_PROTO_MAPPING] ) if state not in _STATE_ICONS: raise StreamlitValueError("state", [repr(name) for name in _STATE_ICONS]) expandable_proto = BlockProto.Expandable() expandable_proto.expanded = expanded expandable_proto.label = label or "" expandable_proto.type = EXPANDABLE_TYPE_TO_PROTO_MAPPING[type] expandable_proto.icon = _STATE_ICONS[state] expandable_proto.state = _STATE_PROTO_VALUES[state] block_proto = BlockProto() block_proto.allow_empty = True block_proto.expandable.CopyFrom(expandable_proto) validate_width(width=width) block_proto.width_config.CopyFrom(get_width_config(width)) status_container = cast( "StatusContainer", parent._block(block_proto=block_proto, dg_type=StatusContainer), ) # `update()` re-sends the block proto at this path. Use the path `_block()` wrote # to, which can be deeper than the parent cursor points to. See issue #16281. status_container._delta_path = status_container._block_delta_path status_container._current_proto = block_proto status_container._current_state = state # We need to sleep here for a very short time to prevent issues when # the status is updated too quickly. If an .update() directly follows the # the initialization, sometimes only the latest update is applied. # Adding a short timeout here allows the frontend to render the update before. time.sleep(0.05) return status_container def __init__( self, root_container: int | None, cursor: Cursor | None, parent: DeltaGenerator | None, block_type: str | None, ) -> None: super().__init__(root_container, cursor, parent, block_type) # Initialized in `_create()`: self._current_proto: BlockProto | None = None self._current_state: States | None = None self._delta_path: list[int] | None = None def update( self, *, label: str | None = None, expanded: bool | None = None, state: States | None = None, ) -> None: """Update the status container. Only specified arguments are updated. Container contents and unspecified arguments remain unchanged. Parameters ---------- label : str or None A new label of the status container. If None, the label is not changed. expanded : bool or None The new expanded state of the status container. If None, the expanded state is not changed. state : "running", "complete", "error", or None The new state of the status container. This mainly changes the icon. If None, the state is not changed. """ if ( self._current_proto is None or self._delta_path is None ): # pragma: no cover - defensive raise RuntimeError( "StatusContainer is not correctly initialized. This should never happen." ) msg = ForwardMsg() msg.metadata.delta_path[:] = self._delta_path msg.delta.add_block.CopyFrom(self._current_proto) if expanded is not None: msg.delta.add_block.expandable.expanded = expanded else: msg.delta.add_block.expandable.ClearField("expanded") if label is not None: msg.delta.add_block.expandable.label = label if state is not None: if state not in _STATE_ICONS: raise StreamlitValueError( "state", [repr(name) for name in _STATE_ICONS] ) msg.delta.add_block.expandable.icon = _STATE_ICONS[state] msg.delta.add_block.expandable.state = _STATE_PROTO_VALUES[state] self._current_state = state self._current_proto = msg.delta.add_block enqueue_message(msg) def __enter__(self) -> Self: # type: ignore[override] # ty: ignore[invalid-method-override] # This is a little dubious: we're returning a different type than # our superclass' `__enter__` function. Maybe DeltaGenerator.__enter__ # should always return `self`? super().__enter__() return self def __exit__( self, typ: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None, ) -> Literal[False]: # Only update if the current state is running if self._current_state == "running": # We need to sleep here for a very short time to prevent issues when # the status is updated too quickly. If an .update() is directly followed # by the exit of the context manager, sometimes only the last update # (to complete) is applied. Adding a short timeout here allows the frontend # to render the update before. time.sleep(0.05) if typ is not None: # If an exception was raised in the context, # we want to update the status to error. self.update(state="error") else: self.update(state="complete") return super().__exit__(typ, exc, tb)