streamlit-nightly 1.43.2.dev20250307__py3-none-any.whl

streamlit/testing/v1/element_tree.py

@@ -0,0 +1,2083 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025)
+#
+# 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 textwrap
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from dataclasses import dataclass, field, fields, is_dataclass
+from datetime import date, datetime, time, timedelta
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generic,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from typing_extensions import TypeAlias
+
+from streamlit import dataframe_util, util
+from streamlit.elements.heading import HeadingProtoTag
+from streamlit.elements.widgets.select_slider import SelectSliderSerde
+from streamlit.elements.widgets.slider import (
+    SliderSerde,
+    SliderStep,
+    SliderValueT,
+)
+from streamlit.elements.widgets.time_widgets import (
+    DateInputSerde,
+    DateWidgetReturn,
+    TimeInputSerde,
+    _parse_date_value,
+)
+from streamlit.proto.Alert_pb2 import Alert as AlertProto
+from streamlit.proto.Checkbox_pb2 import Checkbox as CheckboxProto
+from streamlit.proto.Markdown_pb2 import Markdown as MarkdownProto
+from streamlit.proto.Slider_pb2 import Slider as SliderProto
+from streamlit.proto.WidgetStates_pb2 import WidgetState, WidgetStates
+from streamlit.runtime.state.common import TESTING_KEY, user_key_from_element_id
+
+if TYPE_CHECKING:
+    from pandas import DataFrame as PandasDataframe
+
+    from streamlit.proto.Arrow_pb2 import Arrow as ArrowProto
+    from streamlit.proto.Block_pb2 import Block as BlockProto
+    from streamlit.proto.Button_pb2 import Button as ButtonProto
+    from streamlit.proto.ButtonGroup_pb2 import ButtonGroup as ButtonGroupProto
+    from streamlit.proto.ChatInput_pb2 import ChatInput as ChatInputProto
+    from streamlit.proto.Code_pb2 import Code as CodeProto
+    from streamlit.proto.ColorPicker_pb2 import ColorPicker as ColorPickerProto
+    from streamlit.proto.DateInput_pb2 import DateInput as DateInputProto
+    from streamlit.proto.Element_pb2 import Element as ElementProto
+    from streamlit.proto.Exception_pb2 import Exception as ExceptionProto
+    from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
+    from streamlit.proto.Heading_pb2 import Heading as HeadingProto
+    from streamlit.proto.Json_pb2 import Json as JsonProto
+    from streamlit.proto.Metric_pb2 import Metric as MetricProto
+    from streamlit.proto.MultiSelect_pb2 import MultiSelect as MultiSelectProto
+    from streamlit.proto.NumberInput_pb2 import NumberInput as NumberInputProto
+    from streamlit.proto.Radio_pb2 import Radio as RadioProto
+    from streamlit.proto.Selectbox_pb2 import Selectbox as SelectboxProto
+    from streamlit.proto.Text_pb2 import Text as TextProto
+    from streamlit.proto.TextArea_pb2 import TextArea as TextAreaProto
+    from streamlit.proto.TextInput_pb2 import TextInput as TextInputProto
+    from streamlit.proto.TimeInput_pb2 import TimeInput as TimeInputProto
+    from streamlit.proto.Toast_pb2 import Toast as ToastProto
+    from streamlit.runtime.state.safe_session_state import SafeSessionState
+    from streamlit.testing.v1.app_test import AppTest
+
+T = TypeVar("T")
+
+
+@dataclass
+class InitialValue:
+    """This class is used to represent the initial value of a widget."""
+
+    pass
+
+
+# TODO: This class serves as a fallback option for elements that have not
+# been implemented yet, as well as providing implementations of some
+# trivial methods. It may have significantly reduced scope once all elements
+# have been implemented.
+# This class will not be sufficient implementation for most elements.
+# Widgets need their own classes to translate interactions into the appropriate
+# WidgetState and provide higher level interaction interfaces, and other elements
+# have enough variation in how to get their values that most will need their
+# own classes too.
+@dataclass
+class Element(ABC):
+    """
+    Element base class for testing.
+
+    This class's methods and attributes are universal for all elements
+    implemented in testing. For example, ``Caption``, ``Code``, ``Text``, and
+    ``Title`` inherit from ``Element``. All widget classes also
+    inherit from Element, but have additional methods specific to each
+    widget type. See the ``AppTest`` class for the full list of supported
+    elements.
+
+    For all element classes, parameters of the original element can be obtained
+    as properties. For example, ``Button.label``, ``Caption.help``, and
+    ``Toast.icon``.
+
+    """
+
+    type: str = field(repr=False)
+    proto: Any = field(repr=False)
+    root: ElementTree = field(repr=False)
+    key: str | None
+
+    @abstractmethod
+    def __init__(self, proto: ElementProto, root: ElementTree): ...
+
+    def __iter__(self):
+        yield self
+
+    @property
+    @abstractmethod
+    def value(self) -> Any:
+        """The value or contents of the element."""
+        ...
+
+    def __getattr__(self, name: str) -> Any:
+        """Fallback attempt to get an attribute from the proto"""
+        return getattr(self.proto, name)
+
+    def run(self, *, timeout: float | None = None) -> AppTest:
+        """Run the ``AppTest`` script which contains the element.
+
+        Parameters
+        ----------
+        timeout
+            The maximum number of seconds to run the script. None means
+            use the AppTest's default.
+        """
+        return self.root.run(timeout=timeout)
+
+    def __repr__(self):
+        return util.repr_(self)
+
+
+@dataclass(repr=False)
+class UnknownElement(Element):
+    def __init__(self, proto: ElementProto, root: ElementTree):
+        ty = proto.WhichOneof("type")
+        assert ty is not None
+        self.proto = getattr(proto, ty)
+        self.root = root
+        self.type = ty
+        self.key = None
+
+    @property
+    def value(self) -> Any:
+        try:
+            state = self.root.session_state
+            assert state is not None
+            return state[self.proto.id]
+        except ValueError:
+            # No id field, not a widget
+            return self.proto.value
+
+
+@dataclass(repr=False)
+class Widget(Element, ABC):
+    """Widget base class for testing."""
+
+    id: str = field(repr=False)
+    disabled: bool
+    key: str | None
+    _value: Any
+
+    def __init__(self, proto: Any, root: ElementTree):
+        self.proto = proto
+        self.root = root
+        self.key = user_key_from_element_id(self.id)
+        self._value = None
+
+    def set_value(self, v: Any):
+        """Set the value of the widget."""
+        self._value = v
+        return self
+
+    @property
+    @abstractmethod
+    def _widget_state(self) -> WidgetState: ...
+
+
+El_co = TypeVar("El_co", bound=Element, covariant=True)
+
+
+class ElementList(Generic[El_co]):
+    def __init__(self, els: Sequence[El_co]):
+        self._list: Sequence[El_co] = els
+
+    def __len__(self) -> int:
+        return len(self._list)
+
+    @property
+    def len(self) -> int:
+        return len(self)
+
+    @overload
+    def __getitem__(self, idx: int) -> El_co: ...
+
+    @overload
+    def __getitem__(self, idx: slice) -> ElementList[El_co]: ...
+
+    def __getitem__(self, idx: int | slice) -> El_co | ElementList[El_co]:
+        if isinstance(idx, slice):
+            return ElementList(self._list[idx])
+        else:
+            return self._list[idx]
+
+    def __iter__(self):
+        yield from self._list
+
+    def __repr__(self):
+        return util.repr_(self)
+
+    def __eq__(self, other: ElementList[El_co] | object) -> bool:
+        if isinstance(other, ElementList):
+            return self._list == other._list
+        else:
+            return self._list == other
+
+    @property
+    def values(self) -> Sequence[Any]:
+        return [e.value for e in self]
+
+
+W_co = TypeVar("W_co", bound=Widget, covariant=True)
+
+
+class WidgetList(ElementList[W_co], Generic[W_co]):
+    def __call__(self, key: str) -> W_co:
+        for e in self._list:
+            if e.key == key:
+                return e
+
+        raise KeyError(key)
+
+
+@dataclass(repr=False)
+class AlertBase(Element):
+    proto: AlertProto = field(repr=False)
+    icon: str
+
+    def __init__(self, proto: AlertProto, root: ElementTree):
+        self.proto = proto
+        self.key = None
+        self.root = root
+
+    @property
+    def value(self) -> str:
+        return self.proto.body
+
+
+@dataclass(repr=False)
+class Error(AlertBase):
+    def __init__(self, proto: AlertProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "error"
+
+
+@dataclass(repr=False)
+class Warning(AlertBase):
+    def __init__(self, proto: AlertProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "warning"
+
+
+@dataclass(repr=False)
+class Info(AlertBase):
+    def __init__(self, proto: AlertProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "info"
+
+
+@dataclass(repr=False)
+class Success(AlertBase):
+    def __init__(self, proto: AlertProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "success"
+
+
+@dataclass(repr=False)
+class Button(Widget):
+    """A representation of ``st.button`` and ``st.form_submit_button``."""
+
+    _value: bool
+
+    proto: ButtonProto = field(repr=False)
+    label: str
+    help: str
+    form_id: str
+
+    def __init__(self, proto: ButtonProto, root: ElementTree):
+        super().__init__(proto, root)
+        self._value = False
+        self.type = "button"
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        ws = WidgetState()
+        ws.id = self.id
+        ws.trigger_value = self._value
+        return ws
+
+    @property
+    def value(self) -> bool:
+        """The value of the button. (bool)"""
+        if self._value:
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            return cast(bool, state[TESTING_KEY][self.id])
+
+    def set_value(self, v: bool) -> Button:
+        """Set the value of the button."""
+        self._value = v
+        return self
+
+    def click(self) -> Button:
+        """Set the value of the button to True."""
+        return self.set_value(True)
+
+
+@dataclass(repr=False)
+class ChatInput(Widget):
+    """A representation of ``st.chat_input``."""
+
+    _value: str | None
+    proto: ChatInputProto = field(repr=False)
+    placeholder: str
+
+    def __init__(self, proto: ChatInputProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "chat_input"
+
+    def set_value(self, v: str | None) -> ChatInput:
+        """Set the value of the widget."""
+        self._value = v
+        return self
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        ws = WidgetState()
+        ws.id = self.id
+        if self._value is not None:
+            ws.string_trigger_value.data = self._value
+        return ws
+
+    @property
+    def value(self) -> str | None:
+        """The value of the widget. (str)"""
+        if self._value:
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            return state[TESTING_KEY][self.id]  # type: ignore
+
+
+@dataclass(repr=False)
+class Checkbox(Widget):
+    """A representation of ``st.checkbox``."""
+
+    _value: bool | None
+
+    proto: CheckboxProto = field(repr=False)
+    label: str
+    help: str
+    form_id: str
+
+    def __init__(self, proto: CheckboxProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "checkbox"
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        ws = WidgetState()
+        ws.id = self.id
+        ws.bool_value = self.value
+        return ws
+
+    @property
+    def value(self) -> bool:
+        """The value of the widget. (bool)"""
+        if self._value is not None:
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            return cast(bool, state[self.id])
+
+    def set_value(self, v: bool) -> Checkbox:
+        """Set the value of the widget."""
+        self._value = v
+        return self
+
+    def check(self) -> Checkbox:
+        """Set the value of the widget to True."""
+        return self.set_value(True)
+
+    def uncheck(self) -> Checkbox:
+        """Set the value of the widget to False."""
+        return self.set_value(False)
+
+
+@dataclass(repr=False)
+class Code(Element):
+    """A representation of ``st.code``."""
+
+    proto: CodeProto = field(repr=False)
+
+    language: str
+    show_line_numbers: bool
+    key: None
+
+    def __init__(self, proto: CodeProto, root: ElementTree):
+        self.proto = proto
+        self.key = None
+        self.root = root
+        self.type = "code"
+
+    @property
+    def value(self) -> str:
+        """The value of the element. (str)"""
+        return self.proto.code_text
+
+
+@dataclass(repr=False)
+class ColorPicker(Widget):
+    """A representation of ``st.color_picker``."""
+
+    _value: str | None
+    label: str
+    help: str
+    form_id: str
+
+    proto: ColorPickerProto = field(repr=False)
+
+    def __init__(self, proto: ColorPickerProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "color_picker"
+
+    @property
+    def value(self) -> str:
+        """The currently selected value as a hex string. (str)"""
+        if self._value is not None:
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            return cast(str, state[self.id])
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        """Protobuf message representing the state of the widget, including
+        any interactions that have happened.
+        Should be the same as the frontend would produce for those interactions.
+        """
+        ws = WidgetState()
+        ws.id = self.id
+        ws.string_value = self.value
+        return ws
+
+    def set_value(self, v: str) -> ColorPicker:
+        """Set the value of the widget as a hex string."""
+        self._value = v
+        return self
+
+    def pick(self, v: str) -> ColorPicker:
+        """Set the value of the widget as a hex string. May omit the "#" prefix."""
+        if not v.startswith("#"):
+            v = f"#{v}"
+        return self.set_value(v)
+
+
+@dataclass(repr=False)
+class Dataframe(Element):
+    proto: ArrowProto = field(repr=False)
+
+    def __init__(self, proto: ArrowProto, root: ElementTree):
+        self.key = None
+        self.proto = proto
+        self.root = root
+        self.type = "arrow_data_frame"
+
+    @property
+    def value(self) -> PandasDataframe:
+        return dataframe_util.convert_arrow_bytes_to_pandas_df(self.proto.data)
+
+
+SingleDateValue: TypeAlias = Union[date, datetime]
+DateValue: TypeAlias = Union[SingleDateValue, Sequence[SingleDateValue], None]
+
+
+@dataclass(repr=False)
+class DateInput(Widget):
+    """A representation of ``st.date_input``."""
+
+    _value: DateValue | None | InitialValue
+    proto: DateInputProto = field(repr=False)
+    label: str
+    min: date
+    max: date
+    is_range: bool
+    help: str
+    form_id: str
+
+    def __init__(self, proto: DateInputProto, root: ElementTree):
+        super().__init__(proto, root)
+        self._value = InitialValue()
+        self.type = "date_input"
+        self.min = datetime.strptime(proto.min, "%Y/%m/%d").date()
+        self.max = datetime.strptime(proto.max, "%Y/%m/%d").date()
+
+    def set_value(self, v: DateValue) -> DateInput:
+        """Set the value of the widget."""
+        self._value = v
+        return self
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        ws = WidgetState()
+        ws.id = self.id
+
+        serde = DateInputSerde(None)  # type: ignore
+        ws.string_array_value.data[:] = serde.serialize(self.value)
+        return ws
+
+    @property
+    def value(self) -> DateWidgetReturn:
+        """The value of the widget. (date or Tuple of date)"""
+        if not isinstance(self._value, InitialValue):
+            parsed, _ = _parse_date_value(self._value)
+            return tuple(parsed) if parsed is not None else None  # type: ignore
+        else:
+            state = self.root.session_state
+            assert state
+            return state[self.id]  # type: ignore
+
+
+@dataclass(repr=False)
+class Exception(Element):
+    message: str
+    is_markdown: bool
+    stack_trace: list[str]
+    is_warning: bool
+
+    def __init__(self, proto: ExceptionProto, root: ElementTree):
+        self.key = None
+        self.root = root
+        self.proto = proto
+        self.type = "exception"
+
+        self.is_markdown = proto.message_is_markdown
+        self.stack_trace = list(proto.stack_trace)
+
+    @property
+    def value(self) -> str:
+        return self.message
+
+
+@dataclass(repr=False)
+class HeadingBase(Element, ABC):
+    proto: HeadingProto = field(repr=False)
+
+    tag: str
+    anchor: str | None
+    hide_anchor: bool
+    key: None
+
+    def __init__(self, proto: HeadingProto, root: ElementTree, type_: str):
+        self.proto = proto
+        self.key = None
+        self.root = root
+        self.type = type_
+
+    @property
+    def value(self) -> str:
+        return self.proto.body
+
+
+@dataclass(repr=False)
+class Header(HeadingBase):
+    def __init__(self, proto: HeadingProto, root: ElementTree):
+        super().__init__(proto, root, "header")
+
+
+@dataclass(repr=False)
+class Subheader(HeadingBase):
+    def __init__(self, proto: HeadingProto, root: ElementTree):
+        super().__init__(proto, root, "subheader")
+
+
+@dataclass(repr=False)
+class Title(HeadingBase):
+    def __init__(self, proto: HeadingProto, root: ElementTree):
+        super().__init__(proto, root, "title")
+
+
+@dataclass(repr=False)
+class Json(Element):
+    proto: JsonProto = field(repr=False)
+
+    expanded: bool
+
+    def __init__(self, proto: JsonProto, root: ElementTree):
+        self.proto = proto
+        self.key = None
+        self.root = root
+        self.type = "json"
+
+    @property
+    def value(self) -> str:
+        return self.proto.body
+
+
+@dataclass(repr=False)
+class Markdown(Element):
+    proto: MarkdownProto = field(repr=False)
+
+    is_caption: bool
+    allow_html: bool
+    key: None
+
+    def __init__(self, proto: MarkdownProto, root: ElementTree):
+        self.proto = proto
+        self.key = None
+        self.root = root
+        self.type = "markdown"
+
+    @property
+    def value(self) -> str:
+        return self.proto.body
+
+
+@dataclass(repr=False)
+class Caption(Markdown):
+    def __init__(self, proto: MarkdownProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "caption"
+
+
+@dataclass(repr=False)
+class Divider(Markdown):
+    def __init__(self, proto: MarkdownProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "divider"
+
+
+@dataclass(repr=False)
+class Latex(Markdown):
+    def __init__(self, proto: MarkdownProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "latex"
+
+
+@dataclass(repr=False)
+class Metric(Element):
+    proto: MetricProto
+    label: str
+    delta: str
+    color: str
+    help: str
+
+    def __init__(self, proto: MetricProto, root: ElementTree):
+        self.proto = proto
+        self.key = None
+        self.root = root
+        self.type = "metric"
+
+    @property
+    def value(self) -> str:
+        return self.proto.body
+
+
+@dataclass(repr=False)
+class ButtonGroup(Widget, Generic[T]):
+    """A representation of button_group that is used by ``st.feedback``."""
+
+    _value: list[T] | None
+
+    proto: ButtonGroupProto = field(repr=False)
+    options: list[ButtonGroupProto.Option]
+    form_id: str
+
+    def __init__(self, proto: ButtonGroupProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "button_group"
+        self.options = list(proto.options)
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        """Protobuf message representing the state of the widget, including
+        any interactions that have happened.
+        Should be the same as the frontend would produce for those interactions.
+        """
+        ws = WidgetState()
+        ws.id = self.id
+        ws.int_array_value.data[:] = self.indices
+        return ws
+
+    @property
+    def value(self) -> list[T]:
+        """The currently selected values from the options. (list)"""
+        if self._value is not None:
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            return cast(list[T], state[self.id])
+
+    @property
+    def indices(self) -> Sequence[int]:
+        """The indices of the currently selected values from the options. (list)"""
+        return [self.options.index(self.format_func(v)) for v in self.value]
+
+    @property
+    def format_func(self) -> Callable[[Any], Any]:
+        """The widget's formatting function for displaying options. (callable)"""
+        ss = self.root.session_state
+        return cast(Callable[[Any], Any], ss[TESTING_KEY][self.id])
+
+    def set_value(self, v: list[T]) -> ButtonGroup[T]:
+        """Set the value of the multiselect widget. (list)"""
+
+        self._value = v
+        return self
+
+    def select(self, v: T) -> ButtonGroup[T]:
+        """
+        Add a selection to the widget. Do nothing if the value is already selected.\
+        If testing a multiselect widget with repeated options, use ``set_value``\
+        instead.
+        """
+        current = self.value
+        if v in current:
+            return self
+        else:
+            new = current.copy()
+            new.append(v)
+            self.set_value(new)
+            return self
+
+    def unselect(self, v: T) -> ButtonGroup[T]:
+        """
+        Remove a selection from the widget. Do nothing if the value is not\
+        already selected. If a value is selected multiple times, the first\
+        instance is removed.
+        """
+        current = self.value
+        if v not in current:
+            return self
+        else:
+            new = current.copy()
+            while v in new:
+                new.remove(v)
+            self.set_value(new)
+            return self
+
+
+@dataclass(repr=False)
+class Multiselect(Widget, Generic[T]):
+    """A representation of ``st.multiselect``."""
+
+    _value: list[T] | None
+
+    proto: MultiSelectProto = field(repr=False)
+    label: str
+    options: list[str]
+    max_selections: int
+    help: str
+    form_id: str
+
+    def __init__(self, proto: MultiSelectProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "multiselect"
+        self.options = list(proto.options)
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        """Protobuf message representing the state of the widget, including
+        any interactions that have happened.
+        Should be the same as the frontend would produce for those interactions.
+        """
+        ws = WidgetState()
+        ws.id = self.id
+        ws.int_array_value.data[:] = self.indices
+        return ws
+
+    @property
+    def value(self) -> list[T]:
+        """The currently selected values from the options. (list)"""
+        if self._value is not None:
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            return cast(list[T], state[self.id])
+
+    @property
+    def indices(self) -> Sequence[int]:
+        """The indices of the currently selected values from the options. (list)"""
+        return [self.options.index(self.format_func(v)) for v in self.value]
+
+    @property
+    def format_func(self) -> Callable[[Any], Any]:
+        """The widget's formatting function for displaying options. (callable)"""
+        ss = self.root.session_state
+        return cast(Callable[[Any], Any], ss[TESTING_KEY][self.id])
+
+    def set_value(self, v: list[T]) -> Multiselect[T]:
+        """Set the value of the multiselect widget. (list)"""
+
+        self._value = v
+        return self
+
+    def select(self, v: T) -> Multiselect[T]:
+        """
+        Add a selection to the widget. Do nothing if the value is already selected.\
+        If testing a multiselect widget with repeated options, use ``set_value``\
+        instead.
+        """
+        current = self.value
+        if v in current:
+            return self
+        else:
+            new = current.copy()
+            new.append(v)
+            self.set_value(new)
+            return self
+
+    def unselect(self, v: T) -> Multiselect[T]:
+        """
+        Remove a selection from the widget. Do nothing if the value is not\
+        already selected. If a value is selected multiple times, the first\
+        instance is removed.
+        """
+        current = self.value
+        if v not in current:
+            return self
+        else:
+            new = current.copy()
+            while v in new:
+                new.remove(v)
+            self.set_value(new)
+            return self
+
+
+Number = Union[int, float]
+
+
+@dataclass(repr=False)
+class NumberInput(Widget):
+    """A representation of ``st.number_input``."""
+
+    _value: Number | None | InitialValue
+    proto: NumberInputProto = field(repr=False)
+    label: str
+    min: Number | None
+    max: Number | None
+    step: Number
+    help: str
+    form_id: str
+
+    def __init__(self, proto: NumberInputProto, root: ElementTree):
+        super().__init__(proto, root)
+        self._value = InitialValue()
+        self.type = "number_input"
+        self.min = proto.min if proto.has_min else None
+        self.max = proto.max if proto.has_max else None
+
+    def set_value(self, v: Number | None) -> NumberInput:
+        """Set the value of the ``st.number_input`` widget."""
+        self._value = v
+        return self
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        ws = WidgetState()
+        ws.id = self.id
+        if self.value is not None:
+            ws.double_value = self.value
+        return ws
+
+    @property
+    def value(self) -> Number | None:
+        """Get the current value of the ``st.number_input`` widget."""
+        if not isinstance(self._value, InitialValue):
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+
+            # Awkward to do this with `cast`
+            return state[self.id]  # type: ignore
+
+    def increment(self) -> NumberInput:
+        """Increment the ``st.number_input`` widget as if the user clicked "+"."""
+        if self.value is None:
+            return self
+
+        v = min(self.value + self.step, self.max or float("inf"))
+        return self.set_value(v)
+
+    def decrement(self) -> NumberInput:
+        """Decrement the ``st.number_input`` widget as if the user clicked "-"."""
+        if self.value is None:
+            return self
+
+        v = max(self.value - self.step, self.min or float("-inf"))
+        return self.set_value(v)
+
+
+@dataclass(repr=False)
+class Radio(Widget, Generic[T]):
+    """A representation of ``st.radio``."""
+
+    _value: T | None | InitialValue
+
+    proto: RadioProto = field(repr=False)
+    label: str
+    options: list[str]
+    horizontal: bool
+    help: str
+    form_id: str
+
+    def __init__(self, proto: RadioProto, root: ElementTree):
+        super().__init__(proto, root)
+        self._value = InitialValue()
+        self.type = "radio"
+        self.options = list(proto.options)
+
+    @property
+    def index(self) -> int | None:
+        """The index of the current selection. (int)"""
+        if self.value is None:
+            return None
+        return self.options.index(self.format_func(self.value))
+
+    @property
+    def value(self) -> T | None:
+        """The currently selected value from the options. (Any)"""
+        if not isinstance(self._value, InitialValue):
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            return cast(T, state[self.id])
+
+    @property
+    def format_func(self) -> Callable[[Any], Any]:
+        """The widget's formatting function for displaying options. (callable)"""
+        ss = self.root.session_state
+        return cast(Callable[[Any], Any], ss[TESTING_KEY][self.id])
+
+    def set_value(self, v: T | None) -> Radio[T]:
+        """Set the selection by value."""
+        self._value = v
+        return self
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        """Protobuf message representing the state of the widget, including
+        any interactions that have happened.
+        Should be the same as the frontend would produce for those interactions.
+        """
+        ws = WidgetState()
+        ws.id = self.id
+        if self.index is not None:
+            ws.int_value = self.index
+        return ws
+
+
+@dataclass(repr=False)
+class Selectbox(Widget, Generic[T]):
+    """A representation of ``st.selectbox``."""
+
+    _value: T | None | InitialValue
+
+    proto: SelectboxProto = field(repr=False)
+    label: str
+    options: list[str]
+    help: str
+    form_id: str
+
+    def __init__(self, proto: SelectboxProto, root: ElementTree):
+        super().__init__(proto, root)
+        self._value = InitialValue()
+        self.type = "selectbox"
+        self.options = list(proto.options)
+
+    @property
+    def index(self) -> int | None:
+        """The index of the current selection. (int)"""
+        if self.value is None:
+            return None
+
+        if len(self.options) == 0:
+            return 0
+        return self.options.index(self.format_func(self.value))
+
+    @property
+    def value(self) -> T | None:
+        """The currently selected value from the options. (Any)"""
+        if not isinstance(self._value, InitialValue):
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            return cast(T, state[self.id])
+
+    @property
+    def format_func(self) -> Callable[[Any], Any]:
+        """The widget's formatting function for displaying options. (callable)"""
+        ss = self.root.session_state
+        return cast(Callable[[Any], Any], ss[TESTING_KEY][self.id])
+
+    def set_value(self, v: T | None) -> Selectbox[T]:
+        """Set the selection by value."""
+        self._value = v
+        return self
+
+    def select(self, v: T | None) -> Selectbox[T]:
+        """Set the selection by value."""
+        return self.set_value(v)
+
+    def select_index(self, index: int | None) -> Selectbox[T]:
+        """Set the selection by index."""
+        if index is None:
+            return self.set_value(None)
+        return self.set_value(cast(T, self.options[index]))
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        """Protobuf message representing the state of the widget, including
+        any interactions that have happened.
+        Should be the same as the frontend would produce for those interactions.
+        """
+        ws = WidgetState()
+        ws.id = self.id
+        if self.index is not None:
+            ws.int_value = self.index
+        return ws
+
+
+@dataclass(repr=False)
+class SelectSlider(Widget, Generic[T]):
+    """A representation of ``st.select_slider``."""
+
+    _value: T | Sequence[T] | None
+
+    proto: SliderProto = field(repr=False)
+    label: str
+    data_type: SliderProto.DataType.ValueType
+    options: list[str]
+    help: str
+    form_id: str
+
+    def __init__(self, proto: SliderProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "select_slider"
+        self.options = list(proto.options)
+
+    def set_value(self, v: T | Sequence[T]) -> SelectSlider[T]:
+        """Set the (single) selection by value."""
+        self._value = v
+        return self
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        serde = SelectSliderSerde(self.options, [], False)
+        try:
+            v = serde.serialize(self.format_func(self.value))
+        except (ValueError, TypeError):
+            try:
+                v = serde.serialize([self.format_func(val) for val in self.value])  # type: ignore
+            except:  # noqa: E722
+                raise ValueError(f"Could not find index for {self.value}")
+
+        ws = WidgetState()
+        ws.id = self.id
+        ws.double_array_value.data[:] = v
+        return ws
+
+    @property
+    def value(self) -> T | Sequence[T]:
+        """The currently selected value or range. (Any or Sequence of Any)"""
+        if self._value is not None:
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            # Awkward to do this with `cast`
+            return state[self.id]  # type: ignore
+
+    @property
+    def format_func(self) -> Callable[[Any], Any]:
+        """The widget's formatting function for displaying options. (callable)"""
+        ss = self.root.session_state
+        return cast(Callable[[Any], Any], ss[TESTING_KEY][self.id])
+
+    def set_range(self, lower: T, upper: T) -> SelectSlider[T]:
+        """Set the ranged selection by values."""
+        return self.set_value([lower, upper])
+
+
+@dataclass(repr=False)
+class Slider(Widget, Generic[SliderValueT]):
+    """A representation of ``st.slider``."""
+
+    _value: SliderValueT | Sequence[SliderValueT] | None
+
+    proto: SliderProto = field(repr=False)
+    label: str
+    data_type: SliderProto.DataType.ValueType
+    min: SliderValueT
+    max: SliderValueT
+    step: SliderStep
+    help: str
+    form_id: str
+
+    def __init__(self, proto: SliderProto, root: ElementTree):
+        super().__init__(proto, root)
+        self.type = "slider"
+
+    def set_value(
+        self, v: SliderValueT | Sequence[SliderValueT]
+    ) -> Slider[SliderValueT]:
+        """Set the (single) value of the slider."""
+        self._value = v
+        return self
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        data_type = self.proto.data_type
+        serde = SliderSerde([], data_type, True, None)
+        v = serde.serialize(self.value)
+
+        ws = WidgetState()
+        ws.id = self.id
+        ws.double_array_value.data[:] = v
+        return ws
+
+    @property
+    def value(self) -> SliderValueT | Sequence[SliderValueT]:
+        """The currently selected value or range. (Any or Sequence of Any)"""
+        if self._value is not None:
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            # Awkward to do this with `cast`
+            return state[self.id]  # type: ignore
+
+    def set_range(
+        self, lower: SliderValueT, upper: SliderValueT
+    ) -> Slider[SliderValueT]:
+        """Set the ranged value of the slider."""
+        return self.set_value([lower, upper])
+
+
+@dataclass(repr=False)
+class Table(Element):
+    proto: ArrowProto = field(repr=False)
+
+    def __init__(self, proto: ArrowProto, root: ElementTree):
+        self.key = None
+        self.proto = proto
+        self.root = root
+        self.type = "arrow_table"
+
+    @property
+    def value(self) -> PandasDataframe:
+        return dataframe_util.convert_arrow_bytes_to_pandas_df(self.proto.data)
+
+
+@dataclass(repr=False)
+class Text(Element):
+    proto: TextProto = field(repr=False)
+
+    key: None = None
+
+    def __init__(self, proto: TextProto, root: ElementTree):
+        self.proto = proto
+        self.root = root
+        self.type = "text"
+
+    @property
+    def value(self) -> str:
+        """The value of the element. (str)"""
+        return self.proto.body
+
+
+@dataclass(repr=False)
+class TextArea(Widget):
+    """A representation of ``st.text_area``."""
+
+    _value: str | None | InitialValue
+
+    proto: TextAreaProto = field(repr=False)
+    label: str
+    max_chars: int
+    placeholder: str
+    help: str
+    form_id: str
+
+    def __init__(self, proto: TextAreaProto, root: ElementTree):
+        super().__init__(proto, root)
+        self._value = InitialValue()
+        self.type = "text_area"
+
+    def set_value(self, v: str | None) -> TextArea:
+        """Set the value of the widget."""
+        self._value = v
+        return self
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        ws = WidgetState()
+        ws.id = self.id
+        if self.value is not None:
+            ws.string_value = self.value
+        return ws
+
+    @property
+    def value(self) -> str | None:
+        """The current value of the widget. (str)"""
+        if not isinstance(self._value, InitialValue):
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            # Awkward to do this with `cast`
+            return state[self.id]  # type: ignore
+
+    def input(self, v: str) -> TextArea:
+        """
+        Set the value of the widget only if the value does not exceed the\
+        maximum allowed characters.
+        """
+        # TODO: should input be setting or appending?
+        if self.max_chars and len(v) > self.max_chars:
+            return self
+        return self.set_value(v)
+
+
+@dataclass(repr=False)
+class TextInput(Widget):
+    """A representation of ``st.text_input``."""
+
+    _value: str | None | InitialValue
+    proto: TextInputProto = field(repr=False)
+    label: str
+    max_chars: int
+    autocomplete: str
+    placeholder: str
+    help: str
+    form_id: str
+
+    def __init__(self, proto: TextInputProto, root: ElementTree):
+        super().__init__(proto, root)
+        self._value = InitialValue()
+        self.type = "text_input"
+
+    def set_value(self, v: str | None) -> TextInput:
+        """Set the value of the widget."""
+        self._value = v
+        return self
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        ws = WidgetState()
+        ws.id = self.id
+        if self.value is not None:
+            ws.string_value = self.value
+        return ws
+
+    @property
+    def value(self) -> str | None:
+        """The current value of the widget. (str)"""
+        if not isinstance(self._value, InitialValue):
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            # Awkward to do this with `cast`
+            return state[self.id]  # type: ignore
+
+    def input(self, v: str) -> TextInput:
+        """
+        Set the value of the widget only if the value does not exceed the\
+        maximum allowed characters.
+        """
+        # TODO: should input be setting or appending?
+        if self.max_chars and len(v) > self.max_chars:
+            return self
+        return self.set_value(v)
+
+
+TimeValue: TypeAlias = Union[time, datetime]
+
+
+@dataclass(repr=False)
+class TimeInput(Widget):
+    """A representation of ``st.time_input``."""
+
+    _value: TimeValue | None | InitialValue
+    proto: TimeInputProto = field(repr=False)
+    label: str
+    step: int
+    help: str
+    form_id: str
+
+    def __init__(self, proto: TimeInputProto, root: ElementTree):
+        super().__init__(proto, root)
+        self._value = InitialValue()
+        self.type = "time_input"
+
+    def set_value(self, v: TimeValue | None) -> TimeInput:
+        """Set the value of the widget."""
+        self._value = v
+        return self
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        ws = WidgetState()
+        ws.id = self.id
+
+        serde = TimeInputSerde(None)
+        serialized_value = serde.serialize(self.value)
+        if serialized_value is not None:
+            ws.string_value = serialized_value
+        return ws
+
+    @property
+    def value(self) -> time | None:
+        """The current value of the widget. (time)"""
+        if not isinstance(self._value, InitialValue):
+            v = self._value
+            v = v.time() if isinstance(v, datetime) else v
+            return v
+        else:
+            state = self.root.session_state
+            assert state
+            return state[self.id]  # type: ignore
+
+    def increment(self) -> TimeInput:
+        """Select the next available time."""
+        if self.value is None:
+            return self
+        dt = datetime.combine(date.today(), self.value) + timedelta(seconds=self.step)
+        return self.set_value(dt.time())
+
+    def decrement(self) -> TimeInput:
+        """Select the previous available time."""
+        if self.value is None:
+            return self
+        dt = datetime.combine(date.today(), self.value) - timedelta(seconds=self.step)
+        return self.set_value(dt.time())
+
+
+@dataclass(repr=False)
+class Toast(Element):
+    proto: ToastProto = field(repr=False)
+    icon: str
+
+    def __init__(self, proto: ToastProto, root: ElementTree):
+        self.proto = proto
+        self.key = None
+        self.root = root
+        self.type = "toast"
+
+    @property
+    def value(self) -> str:
+        return self.proto.body
+
+
+@dataclass(repr=False)
+class Toggle(Widget):
+    """A representation of ``st.toggle``."""
+
+    _value: bool | None
+
+    proto: CheckboxProto = field(repr=False)
+    label: str
+    help: str
+    form_id: str
+
+    def __init__(self, proto: CheckboxProto, root: ElementTree):
+        super().__init__(proto, root)
+        self._value = None
+        self.type = "toggle"
+
+    @property
+    def _widget_state(self) -> WidgetState:
+        ws = WidgetState()
+        ws.id = self.id
+        ws.bool_value = self.value
+        return ws
+
+    @property
+    def value(self) -> bool:
+        """The current value of the widget. (bool)"""
+        if self._value is not None:
+            return self._value
+        else:
+            state = self.root.session_state
+            assert state
+            return cast(bool, state[self.id])
+
+    def set_value(self, v: bool) -> Toggle:
+        """Set the value of the widget."""
+        self._value = v
+        return self
+
+
+@dataclass(repr=False)
+class Block:
+    """A container of other elements.
+
+    Elements within a Block can be inspected and interacted with. This follows
+    the same syntax as inspecting and interacting within an ``AppTest`` object.
+
+    For all container classes, parameters of the original element can be
+    obtained as properties. For example, ``ChatMessage.avatar`` and
+    ``Tab.label``.
+    """
+
+    type: str
+    children: dict[int, Node]
+    proto: Any = field(repr=False)
+    root: ElementTree = field(repr=False)
+
+    def __init__(
+        self,
+        proto: BlockProto | None,
+        root: ElementTree,
+    ):
+        self.children = {}
+        self.proto = proto
+        if proto:
+            ty = proto.WhichOneof("type")
+            if ty is not None:
+                self.type = ty
+            else:
+                # `st.container` has no sub-message
+                self.type = "container"
+        else:
+            self.type = "unknown"
+        self.root = root
+
+    def __len__(self) -> int:
+        return len(self.children)
+
+    def __iter__(self):
+        yield self
+        for child_idx in self.children:
+            yield from self.children[child_idx]
+
+    def __getitem__(self, k: int) -> Node:
+        return self.children[k]
+
+    @property
+    def key(self) -> str | None:
+        return None
+
+    # We could implement these using __getattr__ but that would have
+    # much worse type information.
+    @property
+    def button(self) -> WidgetList[Button]:
+        return WidgetList(self.get("button"))  # type: ignore
+
+    @property
+    def button_group(self) -> WidgetList[ButtonGroup[Any]]:
+        return WidgetList(self.get("button_group"))  # type: ignore
+
+    @property
+    def caption(self) -> ElementList[Caption]:
+        return ElementList(self.get("caption"))  # type: ignore
+
+    @property
+    def chat_input(self) -> WidgetList[ChatInput]:
+        return WidgetList(self.get("chat_input"))  # type: ignore
+
+    @property
+    def chat_message(self) -> Sequence[ChatMessage]:
+        return self.get("chat_message")  # type: ignore
+
+    @property
+    def checkbox(self) -> WidgetList[Checkbox]:
+        return WidgetList(self.get("checkbox"))  # type: ignore
+
+    @property
+    def code(self) -> ElementList[Code]:
+        return ElementList(self.get("code"))  # type: ignore
+
+    @property
+    def color_picker(self) -> WidgetList[ColorPicker]:
+        return WidgetList(self.get("color_picker"))  # type: ignore
+
+    @property
+    def columns(self) -> Sequence[Column]:
+        return self.get("column")  # type: ignore
+
+    @property
+    def dataframe(self) -> ElementList[Dataframe]:
+        return ElementList(self.get("arrow_data_frame"))  # type: ignore
+
+    @property
+    def date_input(self) -> WidgetList[DateInput]:
+        return WidgetList(self.get("date_input"))  # type: ignore
+
+    @property
+    def divider(self) -> ElementList[Divider]:
+        return ElementList(self.get("divider"))  # type: ignore
+
+    @property
+    def error(self) -> ElementList[Error]:
+        return ElementList(self.get("error"))  # type: ignore
+
+    @property
+    def exception(self) -> ElementList[Exception]:
+        return ElementList(self.get("exception"))  # type: ignore
+
+    @property
+    def expander(self) -> Sequence[Expander]:
+        return self.get("expander")  # type: ignore
+
+    @property
+    def header(self) -> ElementList[Header]:
+        return ElementList(self.get("header"))  # type: ignore
+
+    @property
+    def info(self) -> ElementList[Info]:
+        return ElementList(self.get("info"))  # type: ignore
+
+    @property
+    def json(self) -> ElementList[Json]:
+        return ElementList(self.get("json"))  # type: ignore
+
+    @property
+    def latex(self) -> ElementList[Latex]:
+        return ElementList(self.get("latex"))  # type: ignore
+
+    @property
+    def markdown(self) -> ElementList[Markdown]:
+        return ElementList(self.get("markdown"))  # type: ignore
+
+    @property
+    def metric(self) -> ElementList[Metric]:
+        return ElementList(self.get("metric"))  # type: ignore
+
+    @property
+    def multiselect(self) -> WidgetList[Multiselect[Any]]:
+        return WidgetList(self.get("multiselect"))  # type: ignore
+
+    @property
+    def number_input(self) -> WidgetList[NumberInput]:
+        return WidgetList(self.get("number_input"))  # type: ignore
+
+    @property
+    def radio(self) -> WidgetList[Radio[Any]]:
+        return WidgetList(self.get("radio"))  # type: ignore
+
+    @property
+    def select_slider(self) -> WidgetList[SelectSlider[Any]]:
+        return WidgetList(self.get("select_slider"))  # type: ignore
+
+    @property
+    def selectbox(self) -> WidgetList[Selectbox[Any]]:
+        return WidgetList(self.get("selectbox"))  # type: ignore
+
+    @property
+    def slider(self) -> WidgetList[Slider[Any]]:
+        return WidgetList(self.get("slider"))  # type: ignore
+
+    @property
+    def status(self) -> Sequence[Status]:
+        return self.get("status")  # type: ignore
+
+    @property
+    def subheader(self) -> ElementList[Subheader]:
+        return ElementList(self.get("subheader"))  # type: ignore
+
+    @property
+    def success(self) -> ElementList[Success]:
+        return ElementList(self.get("success"))  # type: ignore
+
+    @property
+    def table(self) -> ElementList[Table]:
+        return ElementList(self.get("arrow_table"))  # type: ignore
+
+    @property
+    def tabs(self) -> Sequence[Tab]:
+        return self.get("tab")  # type: ignore
+
+    @property
+    def text(self) -> ElementList[Text]:
+        return ElementList(self.get("text"))  # type: ignore
+
+    @property
+    def text_area(self) -> WidgetList[TextArea]:
+        return WidgetList(self.get("text_area"))  # type: ignore
+
+    @property
+    def text_input(self) -> WidgetList[TextInput]:
+        return WidgetList(self.get("text_input"))  # type: ignore
+
+    @property
+    def time_input(self) -> WidgetList[TimeInput]:
+        return WidgetList(self.get("time_input"))  # type: ignore
+
+    @property
+    def title(self) -> ElementList[Title]:
+        return ElementList(self.get("title"))  # type: ignore
+
+    @property
+    def toast(self) -> ElementList[Toast]:
+        return ElementList(self.get("toast"))  # type: ignore
+
+    @property
+    def toggle(self) -> WidgetList[Toggle]:
+        return WidgetList(self.get("toggle"))  # type: ignore
+
+    @property
+    def warning(self) -> ElementList[Warning]:
+        return ElementList(self.get("warning"))  # type: ignore
+
+    def get(self, element_type: str) -> Sequence[Node]:
+        return [e for e in self if e.type == element_type]
+
+    def run(self, *, timeout: float | None = None) -> AppTest:
+        """Run the script with updated widget values.
+
+        Parameters
+        ----------
+        timeout
+            The maximum number of seconds to run the script. None means
+            use the AppTest's default.
+        """
+        return self.root.run(timeout=timeout)
+
+    def __repr__(self):
+        return repr_(self)
+
+
+def repr_(self) -> str:
+    """A custom repr similar to `streamlit.util.repr_` but that shows tree
+    structure using indentation.
+    """
+    classname = self.__class__.__name__
+
+    defaults: list[Any] = [None, "", False, [], set(), {}]
+
+    if is_dataclass(self):
+        fields_vals = (
+            (f.name, getattr(self, f.name))
+            for f in fields(self)
+            if f.repr
+            and getattr(self, f.name) != f.default
+            and getattr(self, f.name) not in defaults
+        )
+    else:
+        fields_vals = ((f, v) for (f, v) in self.__dict__.items() if v not in defaults)
+
+    reprs = []
+    for field_name, value in fields_vals:
+        if isinstance(value, dict):
+            line = f"{field_name}={format_dict(value)}"
+        else:
+            line = f"{field_name}={value!r}"
+        reprs.append(line)
+
+    reprs[0] = "\n" + reprs[0]
+    field_reprs = ",\n".join(reprs)
+
+    field_reprs = textwrap.indent(field_reprs, " " * 4)
+    return f"{classname}({field_reprs}\n)"
+
+
+def format_dict(d: dict[Any, Any]):
+    lines = []
+    for k, v in d.items():
+        line = f"{k}: {v!r}"
+        lines.append(line)
+    r = ",\n".join(lines)
+    r = textwrap.indent(r, " " * 4)
+    r = f"{{\n{r}\n}}"
+    return r
+
+
+@dataclass(repr=False)
+class SpecialBlock(Block):
+    """Base class for the sidebar and main body containers."""
+
+    def __init__(
+        self,
+        proto: BlockProto | None,
+        root: ElementTree,
+        type: str | None = None,
+    ):
+        self.children = {}
+        self.proto = proto
+        if type:
+            self.type = type
+        elif proto and proto.WhichOneof("type"):
+            ty = proto.WhichOneof("type")
+            assert ty is not None
+            self.type = ty
+        else:
+            self.type = "unknown"
+        self.root = root
+
+
+@dataclass(repr=False)
+class ChatMessage(Block):
+    """A representation of ``st.chat_message``."""
+
+    type: str = field(repr=False)
+    proto: BlockProto.ChatMessage = field(repr=False)
+    name: str
+    avatar: str
+
+    def __init__(
+        self,
+        proto: BlockProto.ChatMessage,
+        root: ElementTree,
+    ):
+        self.children = {}
+        self.proto = proto
+        self.root = root
+        self.type = "chat_message"
+        self.name = proto.name
+        self.avatar = proto.avatar
+
+
+@dataclass(repr=False)
+class Column(Block):
+    """A representation of a column within ``st.columns``."""
+
+    type: str = field(repr=False)
+    proto: BlockProto.Column = field(repr=False)
+    weight: float
+    gap: str
+
+    def __init__(
+        self,
+        proto: BlockProto.Column,
+        root: ElementTree,
+    ):
+        self.children = {}
+        self.proto = proto
+        self.root = root
+        self.type = "column"
+        self.weight = proto.weight
+        self.gap = proto.gap
+
+
+@dataclass(repr=False)
+class Expander(Block):
+    type: str = field(repr=False)
+    proto: BlockProto.Expandable = field(repr=False)
+    icon: str
+    label: str
+
+    def __init__(self, proto: BlockProto.Expandable, root: ElementTree):
+        self.children = {}
+        self.proto = proto
+        self.root = root
+        # The internal name is "expandable" but the public API uses "expander"
+        # so the naming of the class and type follows the public name.
+        self.type = "expander"
+        self.icon = proto.icon
+        self.label = proto.label
+
+
+@dataclass(repr=False)
+class Status(Block):
+    type: str = field(repr=False)
+    proto: BlockProto.Expandable = field(repr=False)
+    icon: str
+    label: str
+
+    def __init__(self, proto: BlockProto.Expandable, root: ElementTree):
+        self.children = {}
+        self.proto = proto
+        self.root = root
+        self.type = "status"
+        self.icon = proto.icon
+        self.label = proto.label
+
+    @property
+    def state(self):
+        if self.icon == "spinner":
+            return "running"
+        elif self.icon == ":material/check:":
+            return "complete"
+        elif self.icon == ":material/error:":
+            return "error"
+        else:
+            raise ValueError("Unknown Status state")
+
+
+@dataclass(repr=False)
+class Tab(Block):
+    """A representation of tab within ``st.tabs``."""
+
+    type: str = field(repr=False)
+    proto: BlockProto.Tab = field(repr=False)
+    label: str
+
+    def __init__(
+        self,
+        proto: BlockProto.Tab,
+        root: ElementTree,
+    ):
+        self.children = {}
+        self.proto = proto
+        self.root = root
+        self.type = "tab"
+        self.label = proto.label
+
+
+Node: TypeAlias = Union[Element, Block]
+
+
+def get_widget_state(node: Node) -> WidgetState | None:
+    if isinstance(node, Widget):
+        return node._widget_state
+    else:
+        return None
+
+
+@dataclass(repr=False)
+class ElementTree(Block):
+    """A tree of the elements produced by running a streamlit script.
+
+    Elements can be queried in three ways:
+    - By element type, using `.foo` properties to get a list of all of that element,
+    in the order they appear in the app
+    - By user key, for widgets, by calling the above list with a key: `.foo(key='bar')`
+    - Positionally, using list indexing syntax (`[...]`) to access a child of a
+    block element. Not recommended because the exact tree structure can be surprising.
+
+    Element queries made on a block container will return only the elements
+    descending from that block.
+
+    Returned elements have methods for accessing whatever attributes are relevant.
+    For very simple elements this may be only its value, while complex elements
+    like widgets have many.
+
+    Widgets provide a fluent API for faking frontend interaction and rerunning
+    the script with the new widget values. All widgets provide a low level `set_value`
+    method, along with higher level methods specific to that type of widget.
+    After an interaction, calling `.run()` will update the AppTest with the
+    results of that script run.
+    """
+
+    _runner: AppTest | None = field(repr=False, default=None)
+
+    def __init__(self):
+        self.children = {}
+        self.root = self
+        self.type = "root"
+
+    @property
+    def main(self) -> Block:
+        m = self[0]
+        assert isinstance(m, Block)
+        return m
+
+    @property
+    def sidebar(self) -> Block:
+        s = self[1]
+        assert isinstance(s, Block)
+        return s
+
+    @property
+    def session_state(self) -> SafeSessionState:
+        assert self._runner is not None
+        return self._runner.session_state
+
+    def get_widget_states(self) -> WidgetStates:
+        ws = WidgetStates()
+        for node in self:
+            w = get_widget_state(node)
+            if w is not None:
+                ws.widgets.append(w)
+
+        return ws
+
+    def run(self, *, timeout: float | None = None) -> AppTest:
+        """Run the script with updated widget values.
+
+        Parameters
+        ----------
+        timeout
+            The maximum number of seconds to run the script. None means
+            use the AppTest's default.
+        """
+        assert self._runner is not None
+
+        widget_states = self.get_widget_states()
+        return self._runner._run(widget_states, timeout=timeout)
+
+    def __repr__(self):
+        return format_dict(self.children)
+
+
+def parse_tree_from_messages(messages: list[ForwardMsg]) -> ElementTree:
+    """Transform a list of `ForwardMsg` into a tree matching the implicit
+    tree structure of blocks and elements in a streamlit app.
+
+    Returns the root of the tree, which acts as the entrypoint for the query
+    and interaction API.
+    """
+    root = ElementTree()
+    root.children = {
+        0: SpecialBlock(type="main", root=root, proto=None),
+        1: SpecialBlock(type="sidebar", root=root, proto=None),
+        2: SpecialBlock(type="event", root=root, proto=None),
+    }
+
+    for msg in messages:
+        if not msg.HasField("delta"):
+            continue
+        delta_path = msg.metadata.delta_path
+        delta = msg.delta
+        if delta.WhichOneof("type") == "new_element":
+            elt = delta.new_element
+            ty = elt.WhichOneof("type")
+            new_node: Node
+            if ty == "alert":
+                format = elt.alert.format
+                if format == AlertProto.Format.ERROR:
+                    new_node = Error(elt.alert, root=root)
+                elif format == AlertProto.Format.INFO:
+                    new_node = Info(elt.alert, root=root)
+                elif format == AlertProto.Format.SUCCESS:
+                    new_node = Success(elt.alert, root=root)
+                elif format == AlertProto.Format.WARNING:
+                    new_node = Warning(elt.alert, root=root)
+                else:
+                    raise ValueError(
+                        f"Unknown alert type with format {elt.alert.format}"
+                    )
+            elif ty == "arrow_data_frame":
+                new_node = Dataframe(elt.arrow_data_frame, root=root)
+            elif ty == "arrow_table":
+                new_node = Table(elt.arrow_table, root=root)
+            elif ty == "button":
+                new_node = Button(elt.button, root=root)
+            elif ty == "button_group":
+                new_node = ButtonGroup(elt.button_group, root=root)
+            elif ty == "chat_input":
+                new_node = ChatInput(elt.chat_input, root=root)
+            elif ty == "checkbox":
+                style = elt.checkbox.type
+                if style == CheckboxProto.StyleType.TOGGLE:
+                    new_node = Toggle(elt.checkbox, root=root)
+                else:
+                    new_node = Checkbox(elt.checkbox, root=root)
+            elif ty == "code":
+                new_node = Code(elt.code, root=root)
+            elif ty == "color_picker":
+                new_node = ColorPicker(elt.color_picker, root=root)
+            elif ty == "date_input":
+                new_node = DateInput(elt.date_input, root=root)
+            elif ty == "exception":
+                new_node = Exception(elt.exception, root=root)
+            elif ty == "heading":
+                if elt.heading.tag == HeadingProtoTag.TITLE_TAG.value:
+                    new_node = Title(elt.heading, root=root)
+                elif elt.heading.tag == HeadingProtoTag.HEADER_TAG.value:
+                    new_node = Header(elt.heading, root=root)
+                elif elt.heading.tag == HeadingProtoTag.SUBHEADER_TAG.value:
+                    new_node = Subheader(elt.heading, root=root)
+                else:
+                    raise ValueError(f"Unknown heading type with tag {elt.heading.tag}")
+            elif ty == "json":
+                new_node = Json(elt.json, root=root)
+            elif ty == "markdown":
+                if elt.markdown.element_type == MarkdownProto.Type.NATIVE:
+                    new_node = Markdown(elt.markdown, root=root)
+                elif elt.markdown.element_type == MarkdownProto.Type.CAPTION:
+                    new_node = Caption(elt.markdown, root=root)
+                elif elt.markdown.element_type == MarkdownProto.Type.LATEX:
+                    new_node = Latex(elt.markdown, root=root)
+                elif elt.markdown.element_type == MarkdownProto.Type.DIVIDER:
+                    new_node = Divider(elt.markdown, root=root)
+                else:
+                    raise ValueError(
+                        f"Unknown markdown type {elt.markdown.element_type}"
+                    )
+            elif ty == "metric":
+                new_node = Metric(elt.metric, root=root)
+            elif ty == "multiselect":
+                new_node = Multiselect(elt.multiselect, root=root)
+            elif ty == "number_input":
+                new_node = NumberInput(elt.number_input, root=root)
+            elif ty == "radio":
+                new_node = Radio(elt.radio, root=root)
+            elif ty == "selectbox":
+                new_node = Selectbox(elt.selectbox, root=root)
+            elif ty == "slider":
+                if elt.slider.type == SliderProto.Type.SLIDER:
+                    new_node = Slider(elt.slider, root=root)
+                elif elt.slider.type == SliderProto.Type.SELECT_SLIDER:
+                    new_node = SelectSlider(elt.slider, root=root)
+                else:
+                    raise ValueError(f"Slider with unknown type {elt.slider}")
+            elif ty == "text":
+                new_node = Text(elt.text, root=root)
+            elif ty == "text_area":
+                new_node = TextArea(elt.text_area, root=root)
+            elif ty == "text_input":
+                new_node = TextInput(elt.text_input, root=root)
+            elif ty == "time_input":
+                new_node = TimeInput(elt.time_input, root=root)
+            elif ty == "toast":
+                new_node = Toast(elt.toast, root=root)
+            else:
+                new_node = UnknownElement(elt, root=root)
+        elif delta.WhichOneof("type") == "add_block":
+            block = delta.add_block
+            bty = block.WhichOneof("type")
+            if bty == "chat_message":
+                new_node = ChatMessage(block.chat_message, root=root)
+            elif bty == "column":
+                new_node = Column(block.column, root=root)
+            elif bty == "expandable":
+                if block.expandable.icon:
+                    new_node = Status(block.expandable, root=root)
+                else:
+                    new_node = Expander(block.expandable, root=root)
+            elif bty == "tab":
+                new_node = Tab(block.tab, root=root)
+            else:
+                new_node = Block(proto=block, root=root)
+        else:
+            # add_rows
+            continue
+
+        current_node: Block = root
+        # Every node up to the end is a Block
+        for idx in delta_path[:-1]:
+            children = current_node.children
+            child = children.get(idx)
+            if child is None:
+                child = Block(proto=None, root=root)
+                children[idx] = child
+            assert isinstance(child, Block)
+            current_node = child
+
+        # Handle a block when we already have a placeholder for that location
+        if isinstance(new_node, Block):
+            placeholder_block = current_node.children.get(delta_path[-1])
+            if placeholder_block is not None:
+                new_node.children = placeholder_block.children
+
+        current_node.children[delta_path[-1]] = new_node
+
+    return root