Implement trame backend as st component

Author: edsaac

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "stpyvista"
-version = "0.1.1"
+version = "0.1.2"
 authors = [
   { name = "Edwin Saavedra C.", email = "esaavedrac@u.northwestern.edu" },
 ]

src/stpyvista/panel_backend.py

@@ -105,7 +105,7 @@ def stpyvista(
         Pyvista plotter object to render.
     
     use_container_width : bool = True
-        If True, set the dataframe width to the width of the parent container. \
+        If True, set the 3D view width to the width of the parent container. \
         This takes precedence over the `horizontal_align` argument. \
         Defaults to `True`.
     
@@ -171,7 +171,6 @@ def stpyvista(
         use_container_width=1 if use_container_width else 0,
         bgcolor=plotter.background_color.hex_rgba,
         key=key,
-        default=0,
     )
 
     return component_value

src/stpyvista/trame_backend.py

@@ -1,6 +1,6 @@
 from multiprocessing import Process, Queue
 from warnings import warn
-
+from pathlib import Path
 import streamlit.components.v1 as components
 
 from pyvista.plotting import Plotter
@@ -10,7 +10,47 @@ def _export_html(queue: Queue, plotter: Plotter):
     queue.put(plotter.export_html(filename=None))
 
 
-def stpyvista(plotter: Plotter, **kwargs) -> None:
+def _as_html(plotter: Plotter, **kwargs) -> None:
+    """Plan to remove this function in the future"""
+    if not isinstance(plotter, Plotter):
+        raise TypeError(f"{plotter} is not a `pyvista.Plotter` instance.")
+
+    if "panel_kwargs" in kwargs:
+        warn(
+            "panel_kwargs is not supported by the trame backend.\n"
+            "They will be ignored"
+        )
+
+    if "horizontal_align" in kwargs:
+        warn(
+            "horizontal_align is not supported by the trame backend.\n"
+            "It will be ignored"
+        )
+
+    queue = Queue(maxsize=1)
+    process = Process(target=_export_html, args=(queue, plotter))
+
+    process.start()
+    html_plotter = queue.get().read()
+    process.join()
+
+    if kwargs.get("use_container_width", True):
+        width = None
+    else:
+        width = plotter.window_size[0]
+
+    components.html(html_plotter, height=plotter.window_size[1], width=width)
+
+
+## Using component declaration
+
+frontend_dir = (Path(__file__).parent / "trame_based").absolute()
+_component_func = components.declare_component(
+    "stpyvista_trame", path=str(frontend_dir)
+)
+
+
+def stpyvista(plotter: Plotter, use_container_width=True, key=None, **kwargs) -> None:
     """
     Renders an interactive Pyvista Plotter in streamlit using the
     trame backend.
@@ -19,7 +59,17 @@ def stpyvista(plotter: Plotter, **kwargs) -> None:
     ----------
     plotter: pv.Plotter
         Pyvista plotter object to render.
+    use_container_width: bool = True
+        If True, set the 3D view width to the width of the parent container. \
+        If False, the width is taken from the plotter window size.    
+    key: Optional[str] = None
+        An optional key that uniquely identifies this component. If this is
+        None, and the component's arguments are changed, the component will
+        be re-mounted in the Streamlit frontend and lose its current state.
+
     """
+
+    ## Checks
     if not isinstance(plotter, Plotter):
         raise TypeError(f"{plotter} is not a `pyvista.Plotter` instance.")
 
@@ -35,16 +85,22 @@ def stpyvista(plotter: Plotter, **kwargs) -> None:
             "It will be ignored"
         )
 
+    ## Get HTML of plotter
     queue = Queue(maxsize=1)
     process = Process(target=_export_html, args=(queue, plotter))
 
     process.start()
     html_plotter = queue.get().read()
     process.join()
 
-    if kwargs.get("use_container_width", True):
-        width = None
-    else:
-        width = plotter.window_size[0]
+    ## Set dimensions
+    width = None if use_container_width else plotter.window_size[0]
+    height = plotter.window_size[1]
 
-    components.html(html_plotter, height=plotter.window_size[1], width=width)
+    _component_func(
+        trame_html=html_plotter,
+        height=height,
+        width=width,
+        use_container_width=use_container_width,
+        key=key,
+    )

src/stpyvista/trame_based/index.html

@@ -0,0 +1,17 @@
+<!DOCTYPE html>
+<html>
+
+<head>
+    <meta http-equiv="Content-type" content="text/html; charset=utf-8" />
+    <meta name="viewport"
+        content="width=device-width, height=device-height, initial-scale=1, maximum-scale=1, minimum-scale=1, user-scalable=no">
+    <title>🧊 stpyvista 🧊</title>
+    <script type="text/javascript" src="./streamlit-component-lib.js"></script>
+    <script type="text/javascript" src="./main.js"></script>
+</head>
+
+<body>
+    <iframe id="stpyvistaframe"
+        sandbox="allow-forms allow-modals allow-popups allow-popups-to-escape-sandbox allow-same-origin allow-scripts allow-downloads"
+        frameborder="0" allowfullscreen allowtransparency="true"></iframe>
+</body>

src/stpyvista/trame_based/main.js

@@ -0,0 +1,61 @@
+// The `Streamlit` object exists because our html file includes
+// `streamlit-component-lib.js`.
+// If you get an error about "Streamlit" not being defined, that
+// means you're missing that file.
+
+function sendValue(value) {
+    Streamlit.setComponentValue(value)
+}
+
+/**
+ * The component's render function. This will be called immediately after
+ * the component is initially loaded, and then again every time the
+ * component gets new data from Python.
+ */
+function onRender(event) {
+
+    // Only run the render code the first time the component is loaded.
+    if (!window.rendered) {
+
+        // You most likely want to get the data passed in like this
+        const { trame_html, height, width, use_container_width, key } = event.detail.args;
+
+        const stpyvistaframe = document.getElementById("stpyvistaframe");
+
+        function updateFrameWidth() {
+            stpyvistaframe.width = document.body.offsetWidth;
+        }
+
+        if (Boolean(use_container_width)) {
+            updateFrameWidth();
+
+            window.onresize = function (event) {
+                updateFrameWidth();
+            }
+        } else {
+            stpyvistaframe.width = width;
+        }
+
+        stpyvistaframe.srcdoc = trame_html;
+        stpyvistaframe.scrolling = "yes";
+
+        stpyvistaframe.height = height;
+        Streamlit.setFrameHeight(height + 5);
+
+        // Send some value to python 
+        // Not very useful at the moment but keep it for later
+        // stpyvistadiv.addEventListener('click', event => sendValue(50), false);
+
+        window.rendered = true;
+
+    }
+}
+
+// Render the component whenever python send a "render event"
+Streamlit.events.addEventListener(Streamlit.RENDER_EVENT, onRender)
+
+// Tell Streamlit that the component is ready to receive events
+Streamlit.setComponentReady()
+
+// Render with the correct height, if this is a fixed-height component
+Streamlit.setFrameHeight()

src/stpyvista/trame_based/streamlit-component-lib.js

@@ -0,0 +1,35 @@
+
+// Borrowed minimalistic Streamlit API from Thiago
+// https://discuss.streamlit.io/t/code-snippet-create-components-without-any-frontend-tooling-no-react-babel-webpack-etc/13064
+function sendMessageToStreamlitClient(type, data) {
+  console.log(type, data)
+  const outData = Object.assign({
+      isStreamlitMessage: true,
+      type: type,
+  }, data);
+  window.parent.postMessage(outData, "*");
+}
+
+const Streamlit = {
+  setComponentReady: function() {
+      sendMessageToStreamlitClient("streamlit:componentReady", {apiVersion: 1});
+  },
+  setFrameHeight: function(height) {
+      sendMessageToStreamlitClient("streamlit:setFrameHeight", {height: height});
+  },
+  setComponentValue: function(value) {
+    sendMessageToStreamlitClient("streamlit:setComponentValue", {value: value});
+    // sendMessageToStreamlitClient("streamlit:setComponentValue", {value});  
+  },
+  RENDER_EVENT: "streamlit:render",
+  events: {
+      addEventListener: function(type, callback) {
+          window.addEventListener("message", function(event) {
+              if (event.data.type === type) {
+                  event.detail = event.data
+                  callback(event);
+              }
+          });
+      }
+  }
+}

test/cube.py

@@ -5,15 +5,14 @@
 
 st.title("🧊 `stpyvista`")
 st.sidebar.header("Show PyVista 3D visualizations in Streamlit")
-backend = st.sidebar.radio(
-    "Select backend",
-    ["panel", "trame"],
-)
+backend = st.sidebar.radio("Select backend", ["panel", "trame_html", "trame"], index=2)
 st.sidebar.divider()
 
 if backend == "panel":
     from stpyvista.panel_backend import stpyvista
-else:
+elif backend == "trame_html":
+    from stpyvista.trame_backend import _as_html as stpyvista
+elif backend == "trame":
     from stpyvista.trame_backend import stpyvista
 
 
@@ -55,3 +54,8 @@
 for i, col in enumerate(cols):
     with col:
         stpyvista(plotter, use_container_width=chk, key=f"cube_{i}")
+
+stpyvista(plotter, use_container_width=chk)
+print(hash(plotter))
+
+st.sidebar.button("Rerun")

test/dataview.py

@@ -1,7 +1,8 @@
 import streamlit as st
 import pyvista as pv
 import numpy as np
-from stpyvista import stpyvista, dataview
+from stpyvista import dataview
+from stpyvista.trame_backend import stpyvista
 
 
 def put_in_plotter(actor: pv.DataSet):