Allow configurable expansion depth for st.json (streamlit#8984)

## Describe your changes

Allows users to specify the depth to which the JSON element should be
expanded, collapsing deeper levels.

```python
data = {
    "level_1": {
        "level_2": {
            "level_3a": {
                "item1": "value1",
                "item2": "value2",
            },
            "level_3b": {
                "item3": "value3",
                "level_4": {
                    "item4": "value4",
                    "level_5": {
                        "item5": "value5"
                    }
                }
            }
        },
        "simple_item": "simple_value"
    },
    "top_level_item": "top_level_value"
}

st.json(data) # Default expands all nodes (to max depth). Same behavior as today
st.json(data, expanded=True) # Same as above
st.json(data, expanded=False) # Collapses all nodes. Same behavior as today
st.json(data, expanded=0) # Collapses all nodes. Same as expanded=False
```

## GitHub Issue Link (if applicable)

- Closes: streamlit#8825

## Testing Plan

- Added unit and e2e tests.

---

**Contribution License Agreement**

By submitting this pull request you agree that all contributions to this
project are made under the Apache 2.0 license.

e2e_playwright/st_json.py

@@ -41,3 +41,16 @@
 
 st.subheader("Empty dict:")
 st.json({})
+
+st.subheader("Expand to depth of 2:")
+st.json(
+    {
+        "level1": {
+            "level2": {"level3": {"a": "b"}},
+            "c": "d",
+            "list": [{"list_item": "value"}],
+        },
+        "string": "Hello World",
+    },
+    expanded=2,
+)

e2e_playwright/st_json_test.py

@@ -20,14 +20,15 @@
 def test_st_json_displays_correctly(app: Page, assert_snapshot: ImageCompareFunction):
     """Test st.json renders the data correctly."""
     json_elements = app.get_by_test_id("stJson")
-    expect(json_elements).to_have_count(6)
+    expect(json_elements).to_have_count(7)
 
     assert_snapshot(json_elements.nth(0), name="st_json-simple_dict")
     assert_snapshot(json_elements.nth(1), name="st_json-collapsed")
     assert_snapshot(json_elements.nth(2), name="st_json-with_white_spaces")
     # The complex dict is screenshot tested in the themed test below
     assert_snapshot(json_elements.nth(4), name="st_json-simple_list")
     assert_snapshot(json_elements.nth(5), name="st_json-empty_dict")
+    assert_snapshot(json_elements.nth(6), name="st_json-expanded_2")
 
 
 def test_st_json_displays_correctly_when_themed(

frontend/lib/src/components/elements/Json/Json.tsx

@@ -16,10 +16,10 @@
 
 import React, { ReactElement, useRef } from "react"
 
-import { useTheme } from "@emotion/react"
 import JSON5 from "json5"
-import ReactJson from "react-json-view"
 import Clipboard from "clipboard"
+import ReactJson from "react-json-view"
+import { useTheme } from "@emotion/react"
 
 import ErrorElement from "@streamlit/lib/src/components/shared/ErrorElement"
 import { Json as JsonProto } from "@streamlit/lib/src/proto"
@@ -76,7 +76,7 @@ export default function Json({ width, element }: JsonProps): ReactElement {
     <div data-testid="stJson" style={styleProp} ref={elementRef}>
       <ReactJson
         src={bodyObject}
-        collapsed={!element.expanded}
+        collapsed={element.maxExpandDepth ?? !element.expanded}
         displayDataTypes={false}
         displayObjectSize={false}
         name={false}

lib/streamlit/elements/json.py

@@ -28,10 +28,9 @@
 
 
 def _ensure_serialization(o: object) -> str | list[Any]:
-    """A repr function for json.dumps default arg, which tries to serialize sets as lists"""
-    if isinstance(o, set):
-        return list(o)
-    return repr(o)
+    """A repr function for json.dumps default arg, which tries to serialize sets
+    as lists."""
+    return list(o) if isinstance(o, set) else repr(o)
 
 
 class JsonMixin:
@@ -40,20 +39,23 @@ def json(
         self,
         body: object,
         *,  # keyword-only arguments:
-        expanded: bool = True,
+        expanded: bool | int = True,
     ) -> DeltaGenerator:
         """Display object or string as a pretty-printed JSON string.
 
         Parameters
         ----------
+
         body : object or str
             The object to print as JSON. All referenced objects should be
             serializable to JSON as well. If object is a string, we assume it
             contains serialized JSON.
 
-        expanded : bool
-            An optional boolean that allows the user to set whether the initial
-            state of this json element should be expanded. Defaults to True.
+        expanded : bool or int
+            Controls the initial expansion state of the json element.
+            If bool, ``True`` expands all levels, ``False`` collapses all levels.
+            If int, specifies the depth to which the json should be expanded,
+            collapsing deeper levels. Defaults to ``True``.
 
         Example
         -------
@@ -62,14 +64,14 @@ def json(
         >>> st.json(
         ...     {
         ...         "foo": "bar",
-        ...         "baz": "boz",
         ...         "stuff": [
         ...             "stuff 1",
         ...             "stuff 2",
         ...             "stuff 3",
-        ...             "stuff 5",
         ...         ],
-        ...     }
+        ...         "level1": {"level2": {"level3": {"a": "b"}}},
+        ...     },
+        ...     expanded=2,
         ... )
 
         .. output::
@@ -104,7 +106,18 @@ def json(
 
         json_proto = JsonProto()
         json_proto.body = body
-        json_proto.expanded = expanded
+
+        if isinstance(expanded, bool):
+            json_proto.expanded = expanded
+        elif isinstance(expanded, int):
+            json_proto.expanded = True
+            json_proto.max_expand_depth = expanded
+        else:
+            raise TypeError(
+                f"The type {str(type(expanded))} of `expanded` is not supported"
+                ", must be bool or int."
+            )
+
         return self.dg._enqueue("json", json_proto)
 
     @property

lib/tests/streamlit/elements/json_test.py

@@ -26,7 +26,9 @@ def test_st_json(self):
         st.json('{"some": "json"}')
 
         el = self.get_delta_from_queue().new_element
-        self.assertEqual(el.json.body, '{"some": "json"}')
+        assert el.json.body == '{"some": "json"}'
+        assert el.json.expanded is True
+        assert el.json.HasField("max_expand_depth") is False
 
         # Test that an object containing non-json-friendly keys can still
         # be displayed.  Resultant json body will be missing those keys.
@@ -36,4 +38,25 @@ def test_st_json(self):
         st.json(data)
 
         el = self.get_delta_from_queue().new_element
-        self.assertEqual(el.json.body, '{"array": "array([1, 2, 3, 4, 5])"}')
+        assert el.json.body == '{"array": "array([1, 2, 3, 4, 5])"}'
+
+    def test_expanded_param(self):
+        """Test expanded paramter for `st.json`"""
+        st.json(
+            {
+                "level1": {"level2": {"level3": {"a": "b"}}, "c": "d"},
+            },
+            expanded=2,
+        )
+
+        el = self.get_delta_from_queue().new_element
+        assert el.json.expanded is True
+        assert el.json.max_expand_depth == 2
+
+        with self.assertRaises(TypeError):
+            st.json(
+                {
+                    "level1": {"level2": {"level3": {"a": "b"}}, "c": "d"},
+                },
+                expanded=["foo"],  # type: ignore
+            )

proto/streamlit/proto/Json.proto

@@ -22,6 +22,9 @@ option java_outer_classname = "JsonProto";
 message Json {
   // Content to display.
   string body = 1;
-
+  // Controls the initial expansion state of the json element.
+  // Is superseded by max_expand_depth if provided.
   bool expanded = 2;
+  // The maximum depth to expand the JSON object.
+  optional int32 max_expand_depth = 3;
 }