step three

Signed-off-by: Andrei Gherghescu <[email protected]>

Author: andrei-ng

plotly/src/plot.rs

@@ -1091,6 +1091,7 @@ mod tests {
         assert!(file_size > 0,);
         #[cfg(not(feature = "debug"))]
         assert!(std::fs::remove_file(&dst).is_ok());
+        exporter.close();
     }
 
     #[test]
@@ -1110,6 +1111,7 @@ mod tests {
         assert!(file_size > 0,);
         #[cfg(not(feature = "debug"))]
         assert!(std::fs::remove_file(&dst).is_ok());
+        exporter.close();
     }
 
     #[test]
@@ -1129,6 +1131,7 @@ mod tests {
         assert!(file_size > 0,);
         #[cfg(not(feature = "debug"))]
         assert!(std::fs::remove_file(&dst).is_ok());
+        exporter.close();
     }
 
     #[test]
@@ -1156,6 +1159,7 @@ mod tests {
         assert!(file_size > 0,);
         #[cfg(not(feature = "debug"))]
         assert!(std::fs::remove_file(&dst).is_ok());
+        exporter.close();
     }
 
     #[test]
@@ -1175,6 +1179,7 @@ mod tests {
         assert!(file_size > 0,);
         #[cfg(not(feature = "debug"))]
         assert!(std::fs::remove_file(&dst).is_ok());
+        exporter.close();
     }
 
     #[test]
@@ -1200,6 +1205,7 @@ mod tests {
         // Limit the comparison to the first characters;
         // As image contents seem to be slightly inconsistent across platforms
         assert_eq!(expected_decoded[..2], result_decoded[..2]);
+        exporter.close();
     }
 
     #[test]
@@ -1221,6 +1227,7 @@ mod tests {
         // seem to contain uniquely generated IDs
         const LEN: usize = 10;
         assert_eq!(expected[..LEN], image_svg[..LEN]);
+        exporter.close();
     }
 
     #[test]
@@ -1261,5 +1268,6 @@ mod tests {
         assert!(file_size > 0,);
         #[cfg(not(feature = "debug"))]
         assert!(std::fs::remove_file(&dst).is_ok());
+        exporter.close();
     }
 }

plotly_static/src/lib.rs

@@ -622,27 +622,36 @@ impl StaticExporterBuilder {
 
     /// Create a new WebDriver instance based on the spawn_webdriver flag
     fn create_webdriver(&self) -> Result<WebDriver> {
-        match self.spawn_webdriver {
-            // Try to connect to existing WebDriver or spawn new if not available
-            true => WebDriver::connect_or_spawn(self.webdriver_port),
-            // Create the WebDriver instance without spawning
-            false => WebDriver::new(self.webdriver_port),
-        }
+        let port = self.webdriver_port;
+        let in_async = tokio::runtime::Handle::try_current().is_ok();
+
+        let run_create_fn = |spawn: bool| -> Result<WebDriver> {
+            let work = move || {
+                if spawn {
+                    WebDriver::connect_or_spawn(port)
+                } else {
+                    WebDriver::new(port)
+                }
+            };
+            if in_async {
+                std::thread::spawn(work)
+                    .join()
+                    .map_err(|_| anyhow!("failed to join webdriver thread"))?
+            } else {
+                work()
+            }
+        };
+
+        run_create_fn(self.spawn_webdriver)
     }
 }
 
-// Async builder for async-first exporter (added without reordering existing items)
+// Async builder for async-first exporter (added without reordering existing
+// items)
 impl StaticExporterBuilder {
     /// Build an async exporter for use within async contexts.
     pub fn build_async(&self) -> Result<AsyncStaticExporter> {
-        let wd = if self.spawn_webdriver {
-            let port = self.webdriver_port;
-            std::thread::spawn(move || WebDriver::connect_or_spawn(port))
-                .join()
-                .map_err(|_| anyhow!("failed to join webdriver spawn thread"))??
-        } else {
-            WebDriver::new(self.webdriver_port)?
-        };
+        let wd = self.create_webdriver()?;
         Ok(AsyncStaticExporter {
             webdriver_port: self.webdriver_port,
             webdriver_url: self.webdriver_url.clone(),
@@ -760,6 +769,13 @@ impl StaticExporter {
         height: usize,
         scale: f64,
     ) -> Result<(), Box<dyn std::error::Error>> {
+        if tokio::runtime::Handle::try_current().is_ok() {
+            return Err(anyhow!(
+                "StaticExporter sync methods cannot be used inside an async context. \
+             Use StaticExporterBuilder::build_async() and the associated AsyncStaticExporter::write_fig(...)."
+            )
+            .into());
+        }
         let rt = self.runtime.clone();
         rt.block_on(
             self.inner
@@ -818,17 +834,20 @@ impl StaticExporter {
         height: usize,
         scale: f64,
     ) -> Result<String, Box<dyn std::error::Error>> {
+        if tokio::runtime::Handle::try_current().is_ok() {
+            return Err(anyhow!(
+                "StaticExporter sync methods cannot be used inside an async context. \
+             Use StaticExporterBuilder::build_async() and the associated AsyncStaticExporter::write_to_string(...)."
+            )
+            .into());
+        }
         let rt = self.runtime.clone();
         rt.block_on(
             self.inner
                 .write_to_string(plot, format, width, height, scale),
         )
     }
 
-    /// Convert the Plotly graph to a static image using Kaleido and return the
-    /// result as a String
-    // Removed internal export/extract in favor of delegating to AsyncStaticExporter
-
     fn extract_plain(payload: &str, format: &ImageFormat) -> Result<String> {
         match payload.split_once(",") {
             Some((type_info, data)) => {
@@ -930,16 +949,11 @@ impl StaticExporter {
     }
 
     /// Explicitly close the WebDriver session and stop the driver.
-    /// Prefer calling this in long-running applications to ensure deterministic cleanup.
+    /// Prefer calling this in long-running applications to ensure deterministic
+    /// cleanup.
     pub fn close(&mut self) {
-        if let Some(client) = self.inner.webdriver_client.take() {
-            let runtime = self.runtime.clone();
-            runtime.block_on(async {
-                if let Err(e) = client.close().await {
-                    error!("Failed to close WebDriver client: {e}");
-                }
-            });
-        }
+        let runtime = self.runtime.clone();
+        runtime.block_on(self.inner.close());
     }
 }