Support Graphviz layout engines (#7505)

* Support Graphviz layout engines

* Ensure radio option and engine match in test

* Update unit test

* Update snapshots

* Update comment in proto to remove 'Dot language'

* Update offending snapshot

* Delete all chromium snapshots and regenerate

* Add missing chromium snapshots

* Reintroduce 'Dot language' comment in proto

* Call engine in jest.mock

* Use default renderer when provided only spec as string

* Test if chart renders when input is a dot string

* Delete obsolete e2e test

* Upload last remaining snapshots

e2e/scripts/st_graphviz_chart.py → e2e_playwright/st_graphviz_chart.py

@@ -38,7 +38,12 @@
     c.node("anode")
 
 # complex graph
-finite = graphviz.Digraph("finite_state_machine", filename="fsm.gv")
+engine = st.sidebar.radio(
+    "Select engine",
+    ["dot", "neato", "twopi", "circo", "fdp", "osage", "patchwork"],
+)
+st.sidebar.write(engine)
+finite = graphviz.Digraph("finite_state_machine", filename="fsm.gv", engine=engine)
 finite.attr(rankdir="LR", size="8,5")
 
 finite.attr("node", shape="doublecircle")
@@ -85,3 +90,11 @@
 
 with col2:
     st.graphviz_chart(right_graph)
+
+
+dot_code = """
+digraph Dot {
+  A -> {B, C, D} -> {F}
+}
+"""
+st.graphviz_chart(dot_code)

e2e_playwright/st_graphviz_chart_test.py

@@ -0,0 +1,65 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022)
+#
+# 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 playwright.sync_api import Page, expect
+
+from e2e_playwright.conftest import ImageCompareFunction, wait_for_app_run
+
+
+def test_initial_setup(app: Page):
+    """Initial setup: ensure charts are loaded."""
+
+    wait_for_app_run(app)
+    title_count = len(app.locator(".stGraphVizChart > svg > g > title").all())
+    assert title_count == 6
+
+
+def test_shows_left_and_right_graph(app: Page):
+    """Test if it shows left and right graph."""
+
+    left_text = app.locator(".stGraphVizChart > svg > g > title").nth(3).text_content()
+    right_text = app.locator(".stGraphVizChart > svg > g > title").nth(4).text_content()
+    assert "Left" in left_text and "Right" in right_text
+
+
+def test_renders_with_specified_engines(
+    app: Page, assert_snapshot: ImageCompareFunction
+):
+    """Test if it renders with specified engines."""
+
+    engines = ["dot", "neato", "twopi", "circo", "fdp", "osage", "patchwork"]
+
+    radios = app.query_selector_all('label[data-baseweb="radio"]')
+
+    for idx, engine in enumerate(engines):
+        radios[idx].click(force=True)
+        wait_for_app_run(app)
+        expect(app.get_by_test_id("stMarkdown").nth(0)).to_have_text(engine)
+
+        assert_snapshot(
+            app.locator(".stGraphVizChart > svg").nth(2),
+            name=f"st_graphviz_chart_engine-{engine}",
+        )
+
+
+def test_dot_string(app: Page, assert_snapshot: ImageCompareFunction):
+    """Test if it renders charts when input is a string (dot language)."""
+
+    title = app.locator(".stGraphVizChart > svg > g > title").nth(5)
+    expect(title).to_have_text("Dot")
+
+    assert_snapshot(
+        app.locator(".stGraphVizChart > svg").nth(5),
+        name="st_graphviz_chart_dot_string",
+    )

frontend/lib/src/components/elements/GraphVizChart/GraphVizChart.test.tsx

@@ -28,8 +28,10 @@ jest.mock("d3-graphviz", () => ({
     zoom: () => ({
       fit: () => ({
         scale: () => ({
-          renderDot: () => ({
-            on: jest.fn(),
+          engine: () => ({
+            renderDot: () => ({
+              on: jest.fn(),
+            }),
           }),
         }),
       }),

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

@@ -16,7 +16,7 @@
 
 import React, { ReactElement, useEffect } from "react"
 import { select } from "d3"
-import { graphviz } from "d3-graphviz"
+import { graphviz, Engine } from "d3-graphviz"
 import { logError } from "@streamlit/lib/src/util/log"
 import withFullScreenWrapper from "@streamlit/lib/src/hocs/withFullScreenWrapper"
 import { GraphVizChart as GraphVizChartProto } from "@streamlit/lib/src/proto"
@@ -73,6 +73,7 @@ export function GraphVizChart({
         .zoom(false)
         .fit(true)
         .scale(1)
+        .engine(element.engine as Engine)
         .renderDot(getChartData())
         .on("end", () => {
           const node = select(`#${chartId} > svg`).node() as SVGGraphicsElement

lib/streamlit/elements/graphviz_chart.py

@@ -130,13 +130,16 @@ def marshall(
 
     if type_util.is_graphviz_chart(figure_or_dot):
         dot = figure_or_dot.source
+        engine = figure_or_dot.engine
     elif isinstance(figure_or_dot, str):
         dot = figure_or_dot
+        engine = "dot"
     else:
         raise StreamlitAPIException(
             "Unhandled type for graphviz chart: %s" % type(figure_or_dot)
         )
 
     proto.spec = dot
+    proto.engine = engine
     proto.use_container_width = use_container_width
     proto.element_id = element_id

lib/tests/streamlit/elements/graphviz_test.py

@@ -58,3 +58,18 @@ def test_use_container_width_true(self):
 
         c = self.get_delta_from_queue().new_element.graphviz_chart
         self.assertEqual(c.use_container_width, True)
+
+    def test_engines(self):
+        """Test that it can be called with engines."""
+        engines = ["dot", "neato", "twopi", "circo", "fdp", "osage", "patchwork"]
+        for engine in engines:
+            graph = graphviz.Graph(comment="The Round Table", engine=engine)
+            graph.node("A", "King Arthur")
+            graph.node("B", "Sir Bedevere the gWise")
+            graph.edges(["AB"])
+
+            st.graphviz_chart(graph)
+
+            c = self.get_delta_from_queue().new_element.graphviz_chart
+            self.assertEqual(hasattr(c, "engine"), True)
+            self.assertEqual(c.engine, engine)

proto/streamlit/proto/GraphVizChart.proto

@@ -26,5 +26,8 @@ message GraphVizChart {
   // A unique ID of this element.
   string element_id = 5;
 
+  // The engine used to layout and render the graph.
+  string engine = 6;
+
   reserved 2, 3;
 }