Improved AppTest/ElementTree formatting (streamlit#7658)

lib/streamlit/testing/v1/app_test.py

@@ -24,7 +24,7 @@
 from unittest.mock import MagicMock
 from urllib import parse
 
-from streamlit import source_util, util
+from streamlit import source_util
 from streamlit.proto.WidgetStates_pb2 import WidgetStates
 from streamlit.runtime import Runtime
 from streamlit.runtime.caching.storage.dummy_cache_storage import (
@@ -78,6 +78,7 @@
     Toggle,
     Warning,
     WidgetList,
+    repr_,
 )
 from streamlit.testing.v1.local_script_runner import LocalScriptRunner
 from streamlit.util import HASHLIB_KWARGS
@@ -930,4 +931,4 @@ def get(self, element_type: str) -> Sequence[Node]:
         return self._tree.get(element_type)
 
     def __repr__(self) -> str:
-        return util.repr_(self)
+        return repr_(self)

lib/streamlit/testing/v1/element_tree.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 from __future__ import annotations
 
+import textwrap
 from abc import ABC, abstractmethod
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, fields, is_dataclass
 from datetime import date, datetime, time, timedelta
 from typing import (
     TYPE_CHECKING,
@@ -1556,7 +1557,52 @@ def run(self, *, timeout: float | None = None) -> AppTest:
         return self.root.run(timeout=timeout)
 
     def __repr__(self):
-        return util.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(), dict()]
+
+    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, value in fields_vals:
+        if isinstance(value, dict):
+            line = f"{field}={format_dict(value)}"
+        else:
+            line = f"{field}={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)
@@ -1729,7 +1775,7 @@ def run(self, *, timeout: float | None = None) -> AppTest:
         return self._runner._run(widget_states, timeout=timeout)
 
     def __repr__(self):
-        return repr(self.children)
+        return format_dict(self.children)
 
 
 def parse_tree_from_messages(messages: list[ForwardMsg]) -> ElementTree: