diff --git a/example/derive_expression/run.py b/example/derive_expression/run.py index 58cf5c4..c22ce14 100644 --- 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() -) \ No newline at end of file +try: + out.with_columns( + pl.col("names").panic.panic() + ) +except pl.ComputeError as e: + assert "the plugin panicked" in str(e) + +print("finished") \ No newline at end of file diff --git a/pyo3-polars-derive/src/lib.rs b/pyo3-polars-derive/src/lib.rs index 978930d..a7d07fd 100644 --- 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 index 8616a07..b1aa705 100644 --- 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