diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/mod.rs index e89f14c31680..dad72cab74f3 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/mod.rs @@ -55,16 +55,21 @@ impl NuLazyFrame { } pub fn collect(self, span: Span) -> Result { - self.to_polars() - .collect() - .map_err(|e| ShellError::GenericError { - error: "Error collecting lazy frame".into(), - msg: e.to_string(), - span: Some(span), - help: None, - inner: vec![], - }) - .map(|df| NuDataFrame::new(true, df)) + crate::handle_panic( + || { + self.to_polars() + .collect() + .map_err(|e| ShellError::GenericError { + error: "Error collecting lazy frame".into(), + msg: e.to_string(), + span: Some(span), + help: None, + inner: vec![], + }) + .map(|df| NuDataFrame::new(true, df)) + }, + span, + ) } pub fn apply_with_expr(self, expr: NuExpression, f: F) -> Self diff --git a/crates/nu_plugin_polars/src/lib.rs b/crates/nu_plugin_polars/src/lib.rs index f70ab2b6efd5..342e5cd28c9f 100644 --- a/crates/nu_plugin_polars/src/lib.rs +++ b/crates/nu_plugin_polars/src/lib.rs @@ -1,4 +1,7 @@ -use std::cmp::Ordering; +use std::{ + cmp::Ordering, + panic::{catch_unwind, AssertUnwindSafe}, +}; use cache::cache_commands; pub use cache::{Cache, Cacheable}; @@ -209,6 +212,22 @@ impl Plugin for PolarsPlugin { } } +pub(crate) fn handle_panic(f: F, span: Span) -> Result +where + F: FnOnce() -> Result, +{ + match catch_unwind(AssertUnwindSafe(f)) { + Ok(inner_result) => inner_result, + Err(_) => Err(ShellError::GenericError { + error: "Panic occurred".into(), + msg: "".into(), + span: Some(span), + help: None, + inner: vec![], + }), + } +} + #[cfg(test)] pub mod test { use super::*;