feat: gracefully handle panics (#55)

pola-rs · Dec 30, 2023 · d28aae0 · d28aae0
1 parent 6879b0d
commit d28aae0
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 12 deletions.
diff --git a/example/derive_expression/run.py b/example/derive_expression/run.py
@@ -47,7 +47,11 @@
     assert "the plugin failed with message" in str(e)
 
 
-# For now test if we abort on panic.
-out.with_columns(
-    pl.col("names").panic.panic()
-)
+try:
+    out.with_columns(
+        pl.col("names").panic.panic()
+    )
+except pl.ComputeError as e:
+    assert "the plugin panicked" in str(e)
+
+print("finished")
diff --git a/pyo3-polars-derive/src/lib.rs b/pyo3-polars-derive/src/lib.rs
@@ -175,19 +175,13 @@ fn create_expression_function(ast: syn::ItemFn) -> proc_macro2::TokenStream {
 
 fn get_field_function_name(fn_name: &syn::Ident) -> syn::Ident {
     syn::Ident::new(
-        &format!(
-            "_polars_plugin_field_{}",
-            fn_name,
-        ),
+        &format!("_polars_plugin_field_{}", fn_name,),
         fn_name.span(),
     )
 }
 
 fn get_expression_function_name(fn_name: &syn::Ident) -> syn::Ident {
-    syn::Ident::new(
-        &format!("_polars_plugin_{}", fn_name),
-        fn_name.span(),
-    )
+    syn::Ident::new(&format!("_polars_plugin_{}", fn_name), fn_name.span())
 }
 
 fn get_inputs() -> proc_macro2::TokenStream {

diff --git a/pyo3-polars/src/derive.rs b/pyo3-polars/src/derive.rs
@@ -4,6 +4,7 @@ pub use pyo3_polars_derive::polars_expr;
 use serde::Deserialize;
 use std::cell::RefCell;
 use std::ffi::CString;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 /// Gives the caller extra information on how to execute the expression.
 pub use polars_ffi::version_0::CallerContext;
@@ -39,8 +40,24 @@ pub unsafe extern "C" fn _polars_plugin_get_last_error_message() -> *const std::
     LAST_ERROR.with(|prev| prev.borrow_mut().as_ptr())
 }
 
+static INIT: AtomicBool = AtomicBool::new(false);
+
+fn start_up_init() {
+    // Set a custom panic hook that only shows output if verbose.
+    std::panic::set_hook(Box::new(|info| {
+        let show_message = std::env::var("POLARS_VERBOSE").as_deref().unwrap_or("") == "1";
+        if show_message {
+            eprintln!("{}", info)
+        }
+    }));
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn _polars_plugin_get_version() -> u32 {
+    if !INIT.swap(true, Ordering::Relaxed) {
+        // Plugin version is is always called at least once.
+        start_up_init();
+    }
     let (major, minor) = polars_ffi::get_version();
     // Stack bits together
     ((major as u32) << 16) + minor as u32