Write test for the codeajitesh123/Perf-Review-AI#9

Sign in

1# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2024)
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15from streamlit.testing.v1.app_test import AppTest
16

1# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2024)
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14from __future__ import annotations
15
16import textwrap
17from abc import ABC, abstractmethod
18from dataclasses import dataclass, field, fields, is_dataclass
19from datetime import date, datetime, time, timedelta
20from typing import (
21    TYPE_CHECKING,
22    Any,
23    Generic,
24    List,
25    Sequence,
26    TypeVar,
27    Union,
28    cast,
29    overload,
30)
31
32from pandas import DataFrame
33from typing_extensions import TypeAlias
34
35from streamlit import type_util, util
36from streamlit.elements.heading import HeadingProtoTag
37from streamlit.elements.widgets.select_slider import SelectSliderSerde
38from streamlit.elements.widgets.slider import (
39    SliderScalar,
40    SliderScalarT,
41    SliderSerde,
42    Step,
43)
44from streamlit.elements.widgets.time_widgets import (
45    DateInputSerde,
46    DateWidgetReturn,
47    TimeInputSerde,
48    _parse_date_value,
49)
50from streamlit.proto.Alert_pb2 import Alert as AlertProto
51from streamlit.proto.Arrow_pb2 import Arrow as ArrowProto
52from streamlit.proto.Block_pb2 import Block as BlockProto
53from streamlit.proto.Button_pb2 import Button as ButtonProto
54from streamlit.proto.ChatInput_pb2 import ChatInput as ChatInputProto
55from streamlit.proto.Checkbox_pb2 import Checkbox as CheckboxProto
56from streamlit.proto.Code_pb2 import Code as CodeProto
57from streamlit.proto.ColorPicker_pb2 import ColorPicker as ColorPickerProto
58from streamlit.proto.DateInput_pb2 import DateInput as DateInputProto
59from streamlit.proto.Element_pb2 import Element as ElementProto
60from streamlit.proto.Exception_pb2 import Exception as ExceptionProto
61from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
62from streamlit.proto.Heading_pb2 import Heading as HeadingProto
63from streamlit.proto.Json_pb2 import Json as JsonProto
64from streamlit.proto.Markdown_pb2 import Markdown as MarkdownProto
65from streamlit.proto.Metric_pb2 import Metric as MetricProto
66from streamlit.proto.MultiSelect_pb2 import MultiSelect as MultiSelectProto
67from streamlit.proto.NumberInput_pb2 import NumberInput as NumberInputProto
68from streamlit.proto.Radio_pb2 import Radio as RadioProto
69from streamlit.proto.Selectbox_pb2 import Selectbox as SelectboxProto
70from streamlit.proto.Slider_pb2 import Slider as SliderProto
71from streamlit.proto.Text_pb2 import Text as TextProto
72from streamlit.proto.TextArea_pb2 import TextArea as TextAreaProto
73from streamlit.proto.TextInput_pb2 import TextInput as TextInputProto
74from streamlit.proto.TimeInput_pb2 import TimeInput as TimeInputProto
75from streamlit.proto.Toast_pb2 import Toast as ToastProto
76from streamlit.proto.WidgetStates_pb2 import WidgetState, WidgetStates
77from streamlit.runtime.state.common import user_key_from_widget_id
78from streamlit.runtime.state.safe_session_state import SafeSessionState
79
80if TYPE_CHECKING:
81    from streamlit.testing.v1.app_test import AppTest
82
83T = TypeVar("T")
84
85
86@dataclass
87class InitialValue:
88    """This class is used to represent the initial value of a widget."""
89
90    pass
91
92
93# TODO This class serves as a fallback option for elements that have not
94# been implemented yet, as well as providing implementations of some
95# trivial methods. It may have significantly reduced scope once all elements
96# have been implemented.
97# This class will not be sufficient implementation for most elements.
98# Widgets need their own classes to translate interactions into the appropriate
99# WidgetState and provide higher level interaction interfaces, and other elements
100# have enough variation in how to get their values that most will need their
101# own classes too.
102@dataclass
103class Element(ABC):
104    """
105    Element base class for testing.
106
107    This class's methods and attributes are universal for all elements
108    implemented in testing. For example, ``Caption``, ``Code``, ``Text``, and
109    ``Title`` inherit from ``Element``. All widget classes also
110    inherit from Element, but have additional methods specific to each
111    widget type. See the ``AppTest`` class for the full list of supported
112    elements.
113
114    For all element classes, parameters of the original element can be obtained
115    as properties. For example, ``Button.label``, ``Caption.help``, and
116    ``Toast.icon``.
117
118    """
119
120    type: str = field(repr=False)
121    proto: Any = field(repr=False)
122    root: ElementTree = field(repr=False)
123    key: str | None
124
125    @abstractmethod
126    def __init__(self, proto: ElementProto, root: ElementTree):
127        ...
128
129    def __iter__(self):
130        yield self
131
132    @property
133    @abstractmethod
134    def value(self) -> Any:
135        """The value or contents of the element."""
136        ...
137
138    def __getattr__(self, name: str) -> Any:
139        """Fallback attempt to get an attribute from the proto"""
140        return getattr(self.proto, name)
141
142    def run(self, *, timeout: float | None = None) -> AppTest:
143        """Run the ``AppTest`` script which contains the element.
144
145        Parameters
146        ----------
147        timeout
148            The maximum number of seconds to run the script. None means
149            use the AppTest's default.
150        """
151        return self.root.run(timeout=timeout)
152
153    def __repr__(self):
154        return util.repr_(self)
155
156
157@dataclass(repr=False)
158class UnknownElement(Element):
159    def __init__(self, proto: ElementProto, root: ElementTree):
160        ty = proto.WhichOneof("type")
161        assert ty is not None
162        self.proto = getattr(proto, ty)
163        self.root = root
164        self.type = ty
165        self.key = None
166
167    @property
168    def value(self) -> Any:
169        try:
170            state = self.root.session_state
171            assert state is not None
172            return state[self.proto.id]
173        except ValueError:
174            # No id field, not a widget
175            return self.proto.value
176
177
178@dataclass(repr=False)
179class Widget(Element, ABC):
180    """Widget base class for testing."""
181
182    id: str = field(repr=False)
183    disabled: bool
184    key: str | None
185    _value: Any
186
187    def __init__(self, proto: Any, root: ElementTree):
188        self.proto = proto
189        self.root = root
190        self.key = user_key_from_widget_id(self.id)
191        self._value = None
192
193    def set_value(self, v: Any):
194        """Set the value of the widget."""
195        self._value = v
196        return self
197
198    @property
199    @abstractmethod
200    def _widget_state(self) -> WidgetState:
201        ...
202
203
204El = TypeVar("El", bound=Element, covariant=True)
205
206
207class ElementList(Generic[El]):
208    def __init__(self, els: Sequence[El]):
209        self._list: Sequence[El] = els
210
211    def __len__(self) -> int:
212        return len(self._list)
213
214    @property
215    def len(self) -> int:
216        return len(self)
217
218    @overload
219    def __getitem__(self, idx: int) -> El:
220        ...
221
222    @overload
223    def __getitem__(self, idx: slice) -> ElementList[El]:
224        ...
225
226    def __getitem__(self, idx: int | slice) -> El | ElementList[El]:
227        if isinstance(idx, slice):
228            return ElementList(self._list[idx])
229        else:
230            return self._list[idx]
231
232    def __iter__(self):
233        yield from self._list
234
235    def __repr__(self):
236        return util.repr_(self)
237
238    def __eq__(self, other: ElementList[El] | object) -> bool:
239        if isinstance(other, ElementList):
240            return self._list == other._list
241        else:
242            return self._list == other
243
244    @property
245    def values(self) -> Sequence[Any]:
246        return [e.value for e in self]
247
248
249W = TypeVar("W", bound=Widget, covariant=True)
250
251
252class WidgetList(Generic[W], ElementList[W]):
253    def __call__(self, key: str) -> W:
254        for e in self._list:
255            if e.key == key:
256                return e
257
258        raise KeyError(key)
259
260
261@dataclass(repr=False)
262class AlertBase(Element):
263    proto: AlertProto = field(repr=False)
264    icon: str
265
266    def __init__(self, proto: AlertProto, root: ElementTree):
267        self.proto = proto
268        self.key = None
269        self.root = root
270
271    @property
272    def value(self) -> str:
273        return self.proto.body
274
275
276@dataclass(repr=False)
277class Error(AlertBase):
278    def __init__(self, proto: AlertProto, root: ElementTree):
279        super().__init__(proto, root)
280        self.type = "error"
281
282
283@dataclass(repr=False)
284class Warning(AlertBase):
285    def __init__(self, proto: AlertProto, root: ElementTree):
286        super().__init__(proto, root)
287        self.type = "warning"
288
289
290@dataclass(repr=False)
291class Info(AlertBase):
292    def __init__(self, proto: AlertProto, root: ElementTree):
293        super().__init__(proto, root)
294        self.type = "info"
295
296
297@dataclass(repr=False)
298class Success(AlertBase):
299    def __init__(self, proto: AlertProto, root: ElementTree):
300        super().__init__(proto, root)
301        self.type = "success"
302
303
304@dataclass(repr=False)
305class Button(Widget):
306    """A representation of ``st.button`` and ``st.form_submit_button``."""
307
308    _value: bool
309
310    proto: ButtonProto = field(repr=False)
311    label: str
312    help: str
313    form_id: str
314
315    def __init__(self, proto: ButtonProto, root: ElementTree):
316        super().__init__(proto, root)
317        self._value = False
318        self.type = "button"
319
320    @property
321    def _widget_state(self) -> WidgetState:
322        ws = WidgetState()
323        ws.id = self.id
324        ws.trigger_value = self._value
325        return ws
326
327    @property
328    def value(self) -> bool:
329        """The value of the button. (bool)"""
330        if self._value:
331            return self._value
332        else:
333            state = self.root.session_state
334            assert state
335            return cast(bool, state[self.id])
336
337    def set_value(self, v: bool) -> Button:
338        """Set the value of the button."""
339        self._value = v
340        return self
341
342    def click(self) -> Button:
343        """Set the value of the button to True."""
344        return self.set_value(True)
345
346
347@dataclass(repr=False)
348class ChatInput(Widget):
349    """A representation of ``st.chat_input``."""
350
351    _value: str | None
352    proto: ChatInputProto = field(repr=False)
353    placeholder: str
354
355    def __init__(self, proto: ChatInputProto, root: ElementTree):
356        super().__init__(proto, root)
357        self.type = "chat_input"
358
359    def set_value(self, v: str | None) -> ChatInput:
360        """Set the value of the widget."""
361        self._value = v
362        return self
363
364    @property
365    def _widget_state(self) -> WidgetState:
366        ws = WidgetState()
367        ws.id = self.id
368        if self._value is not None:
369            ws.string_trigger_value.data = self._value
370        return ws
371
372    @property
373    def value(self) -> str | None:
374        """The value of the widget. (str)"""
375        if self._value:
376            return self._value
377        else:
378            state = self.root.session_state
379            assert state
380            return state[self.id]  # type: ignore
381
382
383@dataclass(repr=False)
384class Checkbox(Widget):
385    """A representation of ``st.checkbox``."""
386
387    _value: bool | None
388
389    proto: CheckboxProto = field(repr=False)
390    label: str
391    help: str
392    form_id: str
393
394    def __init__(self, proto: CheckboxProto, root: ElementTree):
395        super().__init__(proto, root)
396        self.type = "checkbox"
397
398    @property
399    def _widget_state(self) -> WidgetState:
400        ws = WidgetState()

401        ws.id = self.id
402        ws.bool_value = self.value
403        return ws
404
405    @property
406    def value(self) -> bool:
407        """The value of the widget. (bool)"""
408        if self._value is not None:
409            return self._value
410        else:
411            state = self.root.session_state
412            assert state
413            return cast(bool, state[self.id])
414
415    def set_value(self, v: bool) -> Checkbox:
416        """Set the value of the widget."""
417        self._value = v
418        return self
419
420    def check(self) -> Checkbox:
421        """Set the value of the widget to True."""
422        return self.set_value(True)
423
424    def uncheck(self) -> Checkbox:
425        """Set the value of the widget to False."""
426        return self.set_value(False)
427
428
429@dataclass(repr=False)
430class Code(Element):
431    """A representation of ``st.code``."""
432
433    proto: CodeProto = field(repr=False)
434
435    language: str
436    show_line_numbers: bool
437    key: None
438
439    def __init__(self, proto: CodeProto, root: ElementTree):
440        self.proto = proto
441        self.key = None
442        self.root = root
443        self.type = "code"
444
445    @property
446    def value(self) -> str:
447        """The value of the element. (str)"""
448        return self.proto.code_text
449
450
451@dataclass(repr=False)
452class ColorPicker(Widget):
453    """A representation of ``st.color_picker``."""
454
455    _value: str | None
456    label: str
457    help: str
458    form_id: str
459
460    proto: ColorPickerProto = field(repr=False)
461
462    def __init__(self, proto: ColorPickerProto, root: ElementTree):
463        super().__init__(proto, root)
464        self.type = "color_picker"
465
466    @property
467    def value(self) -> str:
468        """The currently selected value as a hex string. (str)"""
469        if self._value is not None:
470            return self._value
471        else:
472            state = self.root.session_state
473            assert state
474            return cast(str, state[self.id])
475
476    @property
477    def _widget_state(self) -> WidgetState:
478        """Protobuf message representing the state of the widget, including
479        any interactions that have happened.
480        Should be the same as the frontend would produce for those interactions.
481        """
482        ws = WidgetState()
483        ws.id = self.id
484        ws.string_value = self.value
485        return ws
486
487    def set_value(self, v: str) -> ColorPicker:
488        """Set the value of the widget as a hex string."""
489        self._value = v
490        return self
491
492    def pick(self, v: str) -> ColorPicker:
493        """Set the value of the widget as a hex string. May omit the "#" prefix."""
494        if not v.startswith("#"):
495            v = f"#{v}"
496        return self.set_value(v)
497
498
499@dataclass(repr=False)
500class Dataframe(Element):
501    proto: ArrowProto = field(repr=False)
502
503    def __init__(self, proto: ArrowProto, root: ElementTree):
504        self.key = None
505        self.proto = proto
506        self.root = root
507        self.type = "arrow_data_frame"
508
509    @property
510    def value(self) -> DataFrame:
511        return type_util.bytes_to_data_frame(self.proto.data)
512
513
514SingleDateValue: TypeAlias = Union[date, datetime]
515DateValue: TypeAlias = Union[SingleDateValue, Sequence[SingleDateValue], None]
516
517
518@dataclass(repr=False)
519class DateInput(Widget):
520    """A representation of ``st.date_input``."""
521
522    _value: DateValue | None | InitialValue
523    proto: DateInputProto = field(repr=False)
524    label: str
525    min: date
526    max: date
527    is_range: bool
528    help: str
529    form_id: str
530
531    def __init__(self, proto: DateInputProto, root: ElementTree):
532        super().__init__(proto, root)
533        self._value = InitialValue()
534        self.type = "date_input"
535        self.min = datetime.strptime(proto.min, "%Y/%m/%d").date()
536        self.max = datetime.strptime(proto.max, "%Y/%m/%d").date()
537
538    def set_value(self, v: DateValue) -> DateInput:
539        """Set the value of the widget."""
540        self._value = v
541        return self
542
543    @property
544    def _widget_state(self) -> WidgetState:
545        ws = WidgetState()
546        ws.id = self.id
547
548        serde = DateInputSerde(None)  # type: ignore
549        ws.string_array_value.data[:] = serde.serialize(self.value)
550        return ws
551
552    @property
553    def value(self) -> DateWidgetReturn:
554        """The value of the widget. (date or Tuple of date)"""
555        if not isinstance(self._value, InitialValue):
556            parsed, _ = _parse_date_value(self._value)
557            return tuple(parsed) if parsed is not None else None  # type: ignore
558        else:
559            state = self.root.session_state
560            assert state
561            return state[self.id]  # type: ignore
562
563
564@dataclass(repr=False)
565class Exception(Element):
566    message: str
567    is_markdown: bool
568    stack_trace: list[str]
569    is_warning: bool
570
571    def __init__(self, proto: ExceptionProto, root: ElementTree):
572        self.key = None
573        self.root = root
574        self.proto = proto
575        self.type = "exception"
576
577        self.is_markdown = proto.message_is_markdown
578        self.stack_trace = list(proto.stack_trace)
579
580    @property
581    def value(self) -> str:
582        return self.message
583
584
585@dataclass(repr=False)
586class HeadingBase(Element, ABC):
587    proto: HeadingProto = field(repr=False)
588
589    tag: str
590    anchor: str | None
591    hide_anchor: bool
592    key: None
593
594    def __init__(self, proto: HeadingProto, root: ElementTree, type_: str):
595        self.proto = proto
596        self.key = None
597        self.root = root
598        self.type = type_
599
600    @property
601    def value(self) -> str:
602        return self.proto.body
603
604
605@dataclass(repr=False)
606class Header(HeadingBase):
607    def __init__(self, proto: HeadingProto, root: ElementTree):
608        super().__init__(proto, root, "header")
609
610
611@dataclass(repr=False)
612class Subheader(HeadingBase):
613    def __init__(self, proto: HeadingProto, root: ElementTree):
614        super().__init__(proto, root, "subheader")
615
616
617@dataclass(repr=False)
618class Title(HeadingBase):
619    def __init__(self, proto: HeadingProto, root: ElementTree):
620        super().__init__(proto, root, "title")
621
622
623@dataclass(repr=False)
624class Json(Element):
625    proto: JsonProto = field(repr=False)
626
627    expanded: bool
628
629    def __init__(self, proto: JsonProto, root: ElementTree):
630        self.proto = proto
631        self.key = None
632        self.root = root
633        self.type = "json"
634
635    @property
636    def value(self) -> str:
637        return self.proto.body
638
639
640@dataclass(repr=False)
641class Markdown(Element):
642    proto: MarkdownProto = field(repr=False)
643
644    is_caption: bool
645    allow_html: bool
646    key: None
647
648    def __init__(self, proto: MarkdownProto, root: ElementTree):
649        self.proto = proto
650        self.key = None
651        self.root = root
652        self.type = "markdown"
653
654    @property
655    def value(self) -> str:
656        return self.proto.body
657
658
659@dataclass(repr=False)
660class Caption(Markdown):
661    def __init__(self, proto: MarkdownProto, root: ElementTree):
662        super().__init__(proto, root)
663        self.type = "caption"
664
665
666@dataclass(repr=False)
667class Divider(Markdown):
668    def __init__(self, proto: MarkdownProto, root: ElementTree):
669        super().__init__(proto, root)
670        self.type = "divider"
671
672
673@dataclass(repr=False)
674class Latex(Markdown):
675    def __init__(self, proto: MarkdownProto, root: ElementTree):
676        super().__init__(proto, root)
677        self.type = "latex"
678
679
680@dataclass(repr=False)
681class Metric(Element):
682    proto: MetricProto
683    label: str
684    delta: str
685    color: str
686    help: str
687
688    def __init__(self, proto: MetricProto, root: ElementTree):
689        self.proto = proto
690        self.key = None
691        self.root = root
692        self.type = "metric"
693
694    @property
695    def value(self) -> str:
696        return self.proto.body
697
698
699@dataclass(repr=False)
700class Multiselect(Widget, Generic[T]):
701    """A representation of ``st.multiselect``."""
702
703    _value: list[T] | None
704
705    proto: MultiSelectProto = field(repr=False)
706    label: str
707    options: list[str]
708    max_selections: int
709    help: str
710    form_id: str
711
712    def __init__(self, proto: MultiSelectProto, root: ElementTree):
713        super().__init__(proto, root)
714        self.type = "multiselect"
715        self.options = list(proto.options)
716
717    @property
718    def _widget_state(self) -> WidgetState:
719        """Protobuf message representing the state of the widget, including
720        any interactions that have happened.
721        Should be the same as the frontend would produce for those interactions.
722        """
723        ws = WidgetState()
724        ws.id = self.id
725        ws.int_array_value.data[:] = self.indices
726        return ws
727
728    @property
729    def value(self) -> list[T]:
730        """The currently selected values from the options. (list)"""
731        if self._value is not None:
732            return self._value
733        else:
734            state = self.root.session_state
735            assert state
736            return cast(List[T], state[self.id])
737
738    @property
739    def indices(self) -> Sequence[int]:
740        """The indices of the currently selected values from the options. (list)"""
741        return [self.options.index(str(v)) for v in self.value]
742
743    def set_value(self, v: list[T]) -> Multiselect[T]:
744        """Set the value of the multiselect widget. (list)"""
745
746        # Implementation note: set_value not work correctly if `format_func` is also
747        # passed to the multiselect. This is because we send options via proto with
748        # applied `format_func`, but keep original values in session state
749        # as widget value.
750
751        self._value = v
752        return self
753
754    def select(self, v: T) -> Multiselect[T]:
755        """
756        Add a selection to the widget. Do nothing if the value is already selected.\
757        If testing a multiselect widget with repeated options, use ``set_value``\
758        instead.
759        """
760        current = self.value
761        if v in current:
762            return self
763        else:
764            new = current.copy()
765            new.append(v)
766            self.set_value(new)
767            return self
768
769    def unselect(self, v: T) -> Multiselect[T]:
770        """
771        Remove a selection from the widget. Do nothing if the value is not\
772        already selected. If a value is selected multiple times, the first\
773        instance is removed.
774        """
775        current = self.value
776        if v not in current:
777            return self
778        else:
779            new = current.copy()
780            while v in new:
781                new.remove(v)
782            self.set_value(new)
783            return self
784
785
786Number = Union[int, float]
787
788
789@dataclass(repr=False)
790class NumberInput(Widget):
791    """A representation of ``st.number_input``."""
792
793    _value: Number | None | InitialValue
794    proto: NumberInputProto = field(repr=False)
795    label: str
796    min: Number | None
797    max: Number | None
798    step: Number
799    help: str
800    form_id: str