diff --git a/crates/nu_plugin_polars/src/cache/get.rs b/crates/nu_plugin_polars/src/cache/get.rs new file mode 100644 index 0000000000..11ade261b9 --- /dev/null +++ b/crates/nu_plugin_polars/src/cache/get.rs @@ -0,0 +1,96 @@ +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, + Value, +}; +use polars::{prelude::NamedFrom, series::Series}; +use uuid::Uuid; + +use crate::{ + values::{CustomValueSupport, NuDataFrame}, + PolarsPlugin, +}; + +#[derive(Clone)] +pub struct CacheGet; + +impl PluginCommand for CacheGet { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars store-get" + } + + fn usage(&self) -> &str { + "Gets a Dataframe or other object from the plugin cache." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required("key", SyntaxShape::String, "Key of objects to get") + .input_output_types(vec![ + (Type::Any, Type::Custom("dataframe".into())), + (Type::Any, Type::Custom("expression".into())), + ]) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Get a stored object", + example: r#"let df = ([[a b];[1 2] [3 4]] | polars into-df); + polars store-ls | get key | first | polars store-get $in"#, + result: Some( + NuDataFrame::try_from_series_vec( + vec![Series::new("a", &[1_i64, 3]), Series::new("b", &[2_i64, 4])], + Span::test_data(), + ) + .expect("could not create dataframe") + .into_value(Span::test_data()), + ), + }] + } + + fn run( + &self, + plugin: &Self::Plugin, + _engine: &EngineInterface, + call: &EvaluatedCall, + _input: PipelineData, + ) -> Result { + let key = call + .req::(0) + .and_then(|ref k| as_uuid(k, call.head))?; + + let value = if let Some(cache_value) = plugin.cache.get(&key, true)? { + let polars_object = cache_value.value; + polars_object.into_value(call.head) + } else { + Value::nothing(call.head) + }; + + Ok(PipelineData::Value(value, None)) + } +} + +fn as_uuid(s: &str, span: Span) -> Result { + Uuid::parse_str(s).map_err(|e| ShellError::GenericError { + error: format!("Failed to convert key string to UUID: {e}"), + msg: "".into(), + span: Some(span), + help: None, + inner: vec![], + }) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::{First, Get}; + + #[test] + fn test_examples() -> Result<(), ShellError> { + test_polars_plugin_command_with_decls(&CacheGet, vec![Box::new(Get), Box::new(First)]) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/eager/list.rs b/crates/nu_plugin_polars/src/cache/list.rs similarity index 85% rename from crates/nu_plugin_polars/src/dataframe/eager/list.rs rename to crates/nu_plugin_polars/src/cache/list.rs index 68390ae2bc..da434901e4 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/list.rs +++ b/crates/nu_plugin_polars/src/cache/list.rs @@ -12,7 +12,7 @@ impl PluginCommand for ListDF { type Plugin = PolarsPlugin; fn name(&self) -> &str { - "polars ls" + "polars store-ls" } fn usage(&self) -> &str { @@ -26,8 +26,8 @@ impl PluginCommand for ListDF { fn examples(&self) -> Vec { vec![Example { description: "Creates a new dataframe and shows it in the dataframe list", - example: r#"let test = ([[a b];[1 2] [3 4]] | dfr into-df); - polars ls"#, + example: r#"let test = ([[a b];[1 2] [3 4]] | polars into-df); + polars store-ls"#, result: None, }] } @@ -49,11 +49,12 @@ impl PluginCommand for ListDF { "created" => Value::date(value.created, call.head), "columns" => Value::int(df.as_ref().width() as i64, call.head), "rows" => Value::int(df.as_ref().height() as i64, call.head), - "type" => Value::string("NuDataFrame", call.head), + "type" => Value::string("DataFrame", call.head), "estimated_size" => Value::filesize(df.to_polars().estimated_size() as i64, call.head), "span_contents" => Value::string(span_contents, value.span), "span_start" => Value::int(value.span.start as i64, call.head), "span_end" => Value::int(value.span.end as i64, call.head), + "reference_count" => Value::int(value.reference_count as i64, call.head), }, call.head, ))), @@ -65,11 +66,12 @@ impl PluginCommand for ListDF { "created" => Value::date(value.created, call.head), "columns" => Value::int(lf.as_ref().width() as i64, call.head), "rows" => Value::int(lf.as_ref().height() as i64, call.head), - "type" => Value::string("NuLazyFrame", call.head), + "type" => Value::string("LazyFrame", call.head), "estimated_size" => Value::filesize(lf.to_polars().estimated_size() as i64, call.head), "span_contents" => Value::string(span_contents, value.span), "span_start" => Value::int(value.span.start as i64, call.head), "span_end" => Value::int(value.span.end as i64, call.head), + "reference_count" => Value::int(value.reference_count as i64, call.head), }, call.head, ))) @@ -80,11 +82,12 @@ impl PluginCommand for ListDF { "created" => Value::date(value.created, call.head), "columns" => Value::nothing(call.head), "rows" => Value::nothing(call.head), - "type" => Value::string("NuExpression", call.head), + "type" => Value::string("Expression", call.head), "estimated_size" => Value::nothing(call.head), "span_contents" => Value::string(span_contents, value.span), "span_start" => Value::int(value.span.start as i64, call.head), "span_end" => Value::int(value.span.end as i64, call.head), + "reference_count" => Value::int(value.reference_count as i64, call.head), }, call.head, ))), @@ -93,11 +96,12 @@ impl PluginCommand for ListDF { "key" => Value::string(key.to_string(), call.head), "columns" => Value::nothing(call.head), "rows" => Value::nothing(call.head), - "type" => Value::string("NuLazyGroupBy", call.head), + "type" => Value::string("LazyGroupBy", call.head), "estimated_size" => Value::nothing(call.head), "span_contents" => Value::string(span_contents, call.head), "span_start" => Value::int(call.head.start as i64, call.head), "span_end" => Value::int(call.head.end as i64, call.head), + "reference_count" => Value::int(value.reference_count as i64, call.head), }, call.head, ))), @@ -106,11 +110,12 @@ impl PluginCommand for ListDF { "key" => Value::string(key.to_string(), call.head), "columns" => Value::nothing(call.head), "rows" => Value::nothing(call.head), - "type" => Value::string("NuWhen", call.head), + "type" => Value::string("When", call.head), "estimated_size" => Value::nothing(call.head), "span_contents" => Value::string(span_contents.to_string(), call.head), "span_start" => Value::int(call.head.start as i64, call.head), "span_end" => Value::int(call.head.end as i64, call.head), + "reference_count" => Value::int(value.reference_count as i64, call.head), }, call.head, ))), diff --git a/crates/nu_plugin_polars/src/cache.rs b/crates/nu_plugin_polars/src/cache/mod.rs similarity index 71% rename from crates/nu_plugin_polars/src/cache.rs rename to crates/nu_plugin_polars/src/cache/mod.rs index d295f449a9..8862f5bb51 100644 --- a/crates/nu_plugin_polars/src/cache.rs +++ b/crates/nu_plugin_polars/src/cache/mod.rs @@ -1,10 +1,15 @@ +mod get; +mod list; +mod rm; + use std::{ collections::HashMap, sync::{Mutex, MutexGuard}, }; use chrono::{DateTime, FixedOffset, Local}; -use nu_plugin::EngineInterface; +pub use list::ListDF; +use nu_plugin::{EngineInterface, PluginCommand}; use nu_protocol::{LabeledError, ShellError, Span}; use uuid::Uuid; @@ -16,6 +21,7 @@ pub struct CacheValue { pub value: PolarsPluginObject, pub created: DateTime, pub span: Span, + pub reference_count: i16, } #[derive(Default)] @@ -35,15 +41,32 @@ impl Cache { } /// Removes an item from the plugin cache. - /// The maybe_engine parameter is required outside of testing + /// + /// * `maybe_engine` - Current EngineInterface reference. Required outside of testing + /// * `key` - The key of the cache entry to remove. + /// * `force` - Delete even if there are multiple references pub fn remove( &self, maybe_engine: Option<&EngineInterface>, - uuid: &Uuid, + key: &Uuid, + force: bool, ) -> Result, ShellError> { let mut lock = self.lock()?; - let removed = lock.remove(uuid); - plugin_debug!("PolarsPlugin: removing {uuid} from cache: {removed:?}"); + + let reference_count = lock.get_mut(key).map(|cache_value| { + cache_value.reference_count -= 1; + cache_value.reference_count + }); + + let removed = if force || reference_count.unwrap_or_default() < 1 { + let removed = lock.remove(key); + plugin_debug!("PolarsPlugin: removing {key} from cache: {removed:?}"); + removed + } else { + plugin_debug!("PolarsPlugin: decrementing reference count for {key}"); + None + }; + // Once there are no more entries in the cache // we can turn plugin gc back on match maybe_engine { @@ -83,15 +106,21 @@ impl Cache { value, created: Local::now().into(), span, + reference_count: 1, }; let result = lock.insert(uuid, cache_value); drop(lock); Ok(result) } - pub fn get(&self, uuid: &Uuid) -> Result, ShellError> { - let lock = self.lock()?; - let result = lock.get(uuid).cloned(); + pub fn get(&self, uuid: &Uuid, increment: bool) -> Result, ShellError> { + let mut lock = self.lock()?; + let result = lock.get_mut(uuid).map(|cv| { + if increment { + cv.reference_count += 1; + } + cv.clone() + }); drop(lock); Ok(result) } @@ -134,10 +163,18 @@ pub trait Cacheable: Sized + Clone { } fn get_cached(plugin: &PolarsPlugin, id: &Uuid) -> Result, ShellError> { - if let Some(cache_value) = plugin.cache.get(id)? { + if let Some(cache_value) = plugin.cache.get(id, false)? { Ok(Some(Self::from_cache_value(cache_value.value)?)) } else { Ok(None) } } } + +pub(crate) fn cache_commands() -> Vec>> { + vec![ + Box::new(ListDF), + Box::new(rm::CacheRemove), + Box::new(get::CacheGet), + ] +} diff --git a/crates/nu_plugin_polars/src/cache/rm.rs b/crates/nu_plugin_polars/src/cache/rm.rs new file mode 100644 index 0000000000..b8b814ba60 --- /dev/null +++ b/crates/nu_plugin_polars/src/cache/rm.rs @@ -0,0 +1,106 @@ +use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; +use nu_protocol::{ + Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, + Value, +}; +use uuid::Uuid; + +use crate::PolarsPlugin; + +#[derive(Clone)] +pub struct CacheRemove; + +impl PluginCommand for CacheRemove { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars store-rm" + } + + fn usage(&self) -> &str { + "Removes a stored Dataframe or other object from the plugin cache." + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .rest("keys", SyntaxShape::String, "Keys of objects to remove") + .input_output_type(Type::Any, Type::List(Box::new(Type::String))) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Removes a stored ", + example: r#"let df = ([[a b];[1 2] [3 4]] | polars into-df); + polars store-ls | get key | first | polars store-rm $in"#, + result: None, + }] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &EngineInterface, + call: &EvaluatedCall, + _input: PipelineData, + ) -> Result { + let msgs: Vec = call + .rest::(0)? + .into_iter() + .map(|ref key| remove_cache_entry(plugin, engine, key, call.head)) + .collect::, ShellError>>()?; + + Ok(PipelineData::Value(Value::list(msgs, call.head), None)) + } +} + +fn remove_cache_entry( + plugin: &PolarsPlugin, + engine: &EngineInterface, + key: &str, + span: Span, +) -> Result { + let key = as_uuid(key, span)?; + let msg = plugin + .cache + .remove(Some(engine), &key, true)? + .map(|_| format!("Removed: {key}")) + .unwrap_or_else(|| format!("No value found for key: {key}")); + Ok(Value::string(msg, span)) +} + +fn as_uuid(s: &str, span: Span) -> Result { + Uuid::parse_str(s).map_err(|e| ShellError::GenericError { + error: format!("Failed to convert key string to UUID: {e}"), + msg: "".into(), + span: Some(span), + help: None, + inner: vec![], + }) +} + +#[cfg(test)] +mod test { + use nu_command::{First, Get}; + use nu_plugin_test_support::PluginTest; + use nu_protocol::Span; + + use super::*; + + #[test] + fn test_remove() -> Result<(), ShellError> { + let plugin = PolarsPlugin::new_test_mode().into(); + let pipeline_data = PluginTest::new("polars", plugin)? + .add_decl(Box::new(First))? + .add_decl(Box::new(Get))? + .eval("let df = ([[a b];[1 2] [3 4]] | polars into-df); polars store-ls | get key | first | polars store-rm $in")?; + let value = pipeline_data.into_value(Span::test_data()); + let msg = value + .as_list()? + .first() + .expect("there should be a first entry") + .as_str()?; + assert!(msg.contains("Removed")); + Ok(()) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/eager/mod.rs b/crates/nu_plugin_polars/src/dataframe/eager/mod.rs index 44dc6b7646..dc50ba7cd2 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/mod.rs @@ -9,7 +9,6 @@ mod filter_with; mod first; mod get; mod last; -mod list; mod melt; mod open; mod query_df; @@ -45,7 +44,6 @@ pub use filter_with::FilterWith; pub use first::FirstDF; pub use get::GetDF; pub use last::LastDF; -pub use list::ListDF; pub use melt::MeltDF; use nu_plugin::PluginCommand; pub use query_df::QueryDf; @@ -82,7 +80,6 @@ pub(crate) fn eager_commands() -> Vec Result<(), ShellError> { - test_polars_plugin_command(&AsDateTime) + test_polars_plugin_command_with_decls(&AsDateTime, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/date/get_day.rs b/crates/nu_plugin_polars/src/dataframe/series/date/get_day.rs index 746a740161..ea26f4999c 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/date/get_day.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/date/get_day.rs @@ -93,10 +93,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&GetDay) + test_polars_plugin_command_with_decls(&GetDay, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/date/get_hour.rs b/crates/nu_plugin_polars/src/dataframe/series/date/get_hour.rs index 91a0ed7be1..d8f6e2e1c1 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/date/get_hour.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/date/get_hour.rs @@ -85,10 +85,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&GetHour) + test_polars_plugin_command_with_decls(&GetHour, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/date/get_minute.rs b/crates/nu_plugin_polars/src/dataframe/series/date/get_minute.rs index dfa4f800f0..bb2f0abc19 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/date/get_minute.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/date/get_minute.rs @@ -83,10 +83,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&GetMinute) + test_polars_plugin_command_with_decls(&GetMinute, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/date/get_month.rs b/crates/nu_plugin_polars/src/dataframe/series/date/get_month.rs index 38a37fcfc8..141665d371 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/date/get_month.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/date/get_month.rs @@ -85,10 +85,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&GetMonth) + test_polars_plugin_command_with_decls(&GetMonth, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/date/get_nanosecond.rs b/crates/nu_plugin_polars/src/dataframe/series/date/get_nanosecond.rs index 182a461a4a..a94f82add5 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/date/get_nanosecond.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/date/get_nanosecond.rs @@ -85,10 +85,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&GetNanosecond) + test_polars_plugin_command_with_decls(&GetNanosecond, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/date/get_ordinal.rs b/crates/nu_plugin_polars/src/dataframe/series/date/get_ordinal.rs index 31465a0f2b..7f074df8d9 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/date/get_ordinal.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/date/get_ordinal.rs @@ -85,10 +85,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&GetOrdinal) + test_polars_plugin_command_with_decls(&GetOrdinal, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/date/get_second.rs b/crates/nu_plugin_polars/src/dataframe/series/date/get_second.rs index 4f5a10ab04..bf11804443 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/date/get_second.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/date/get_second.rs @@ -85,10 +85,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&GetSecond) + test_polars_plugin_command_with_decls(&GetSecond, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/date/get_week.rs b/crates/nu_plugin_polars/src/dataframe/series/date/get_week.rs index 5de61a0f1f..c3b5724bc4 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/date/get_week.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/date/get_week.rs @@ -85,10 +85,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&GetWeek) + test_polars_plugin_command_with_decls(&GetWeek, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/date/get_weekday.rs b/crates/nu_plugin_polars/src/dataframe/series/date/get_weekday.rs index 871ab3fec7..e495370718 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/date/get_weekday.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/date/get_weekday.rs @@ -85,10 +85,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&GetWeekDay) + test_polars_plugin_command_with_decls(&GetWeekDay, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/date/get_year.rs b/crates/nu_plugin_polars/src/dataframe/series/date/get_year.rs index db0af6070a..380cd6367d 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/date/get_year.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/date/get_year.rs @@ -85,10 +85,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&GetYear) + test_polars_plugin_command_with_decls(&GetYear, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs b/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs index 2f3104a809..d793c3910f 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs @@ -104,10 +104,11 @@ fn command( #[cfg(test)] mod test { use super::*; - use crate::test::test_polars_plugin_command; + use crate::test::test_polars_plugin_command_with_decls; + use nu_command::IntoDatetime; #[test] fn test_examples() -> Result<(), ShellError> { - test_polars_plugin_command(&StrFTime) + test_polars_plugin_command_with_decls(&StrFTime, vec![Box::new(IntoDatetime)]) } } diff --git a/crates/nu_plugin_polars/src/dataframe/values/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/mod.rs index d7c26d5942..e0cc3edecb 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/mod.rs @@ -107,6 +107,16 @@ impl PolarsPluginObject { PolarsPluginObject::NuWhen(w) => w.id, } } + + pub fn into_value(self, span: Span) -> Value { + match self { + PolarsPluginObject::NuDataFrame(df) => df.into_value(span), + PolarsPluginObject::NuLazyFrame(lf) => lf.into_value(span), + PolarsPluginObject::NuExpression(e) => e.into_value(span), + PolarsPluginObject::NuLazyGroupBy(lg) => lg.into_value(span), + PolarsPluginObject::NuWhen(w) => w.into_value(span), + } + } } #[derive(Debug, Clone)] diff --git a/crates/nu_plugin_polars/src/lib.rs b/crates/nu_plugin_polars/src/lib.rs index 2691afb010..3baca54ad9 100644 --- a/crates/nu_plugin_polars/src/lib.rs +++ b/crates/nu_plugin_polars/src/lib.rs @@ -1,5 +1,6 @@ use std::cmp::Ordering; +use cache::cache_commands; pub use cache::{Cache, Cacheable}; use dataframe::{stub::PolarsCmd, values::CustomValueType}; use nu_plugin::{EngineInterface, Plugin, PluginCommand}; @@ -40,6 +41,7 @@ impl Plugin for PolarsPlugin { commands.append(&mut lazy_commands()); commands.append(&mut expr_commands()); commands.append(&mut series_commands()); + commands.append(&mut cache_commands()); commands } @@ -50,7 +52,7 @@ impl Plugin for PolarsPlugin { ) -> Result<(), LabeledError> { if !self.disable_cache_drop { let id = CustomValueType::try_from_custom_value(custom_value)?.id(); - let _ = self.cache.remove(Some(engine), &id); + let _ = self.cache.remove(Some(engine), &id, false); } Ok(()) } @@ -178,9 +180,8 @@ impl Plugin for PolarsPlugin { pub mod test { use super::*; use crate::values::PolarsPluginObject; - use nu_command::IntoDatetime; use nu_plugin_test_support::PluginTest; - use nu_protocol::{ShellError, Span}; + use nu_protocol::{engine::Command, ShellError, Span}; impl PolarsPlugin { /// Creates a new polars plugin in test mode @@ -193,6 +194,13 @@ pub mod test { } pub fn test_polars_plugin_command(command: &impl PluginCommand) -> Result<(), ShellError> { + test_polars_plugin_command_with_decls(command, vec![]) + } + + pub fn test_polars_plugin_command_with_decls( + command: &impl PluginCommand, + decls: Vec>, + ) -> Result<(), ShellError> { let plugin = PolarsPlugin::new_test_mode(); let examples = command.examples(); @@ -210,9 +218,12 @@ pub mod test { } } - PluginTest::new("polars", plugin.into())? - .add_decl(Box::new(IntoDatetime))? - .test_examples(&examples)?; + let mut plugin_test = PluginTest::new("polars", plugin.into())?; + + for decl in decls { + let _ = plugin_test.add_decl(decl)?; + } + plugin_test.test_examples(&examples)?; Ok(()) }