From 6bf955a5a5577636395041f3c9f3d96b7071f75a Mon Sep 17 00:00:00 2001 From: Julian Amarilla Date: Fri, 23 May 2025 20:31:28 -0300 Subject: [PATCH] Fix #15571 panic on write to source parquet file (#15601) Fixes #15571 # Description Writing to a source `.parquet` (`polars open some_file.parquet | polars save some_file.parquet`) file made the plugin panic, added a guard to check the data_source path as per [this comment](https://github.com/nushell/nushell/issues/15571#issuecomment-2812707161) Example output now: Screenshot 2025-04-18 at 21 10 30 # User-Facing Changes # Tests + Formatting - Add 1 test - clippy OK - cargo fmt OK # After Submitting No action required --- crates/nu-command/src/debug/metadata.rs | 93 ++++++------------- .../nu-command/src/debug/metadata_access.rs | 39 ++------ crates/nu-command/src/debug/mod.rs | 1 + crates/nu-command/src/debug/util.rs | 34 +++++++ .../command/aggregation/agg_groups.rs | 2 + .../command/aggregation/aggregate.rs | 2 + .../dataframe/command/aggregation/count.rs | 2 + .../command/aggregation/cumulative.rs | 2 + .../dataframe/command/aggregation/groupby.rs | 5 +- .../dataframe/command/aggregation/implode.rs | 2 + .../src/dataframe/command/aggregation/max.rs | 2 + .../src/dataframe/command/aggregation/mean.rs | 2 + .../dataframe/command/aggregation/median.rs | 2 + .../src/dataframe/command/aggregation/min.rs | 2 + .../dataframe/command/aggregation/n_null.rs | 5 +- .../dataframe/command/aggregation/n_unique.rs | 2 + .../src/dataframe/command/aggregation/over.rs | 2 + .../dataframe/command/aggregation/quantile.rs | 2 + .../dataframe/command/aggregation/rolling.rs | 5 +- .../src/dataframe/command/aggregation/std.rs | 2 + .../src/dataframe/command/aggregation/sum.rs | 2 + .../command/aggregation/value_counts.rs | 5 +- .../src/dataframe/command/aggregation/var.rs | 2 + .../dataframe/command/boolean/all_false.rs | 5 +- .../src/dataframe/command/boolean/all_true.rs | 5 +- .../src/dataframe/command/boolean/arg_true.rs | 5 +- .../src/dataframe/command/boolean/expr_not.rs | 2 + .../command/boolean/is_duplicated.rs | 5 +- .../src/dataframe/command/boolean/is_in.rs | 2 + .../dataframe/command/boolean/is_not_null.rs | 2 + .../src/dataframe/command/boolean/is_null.rs | 2 + .../dataframe/command/boolean/is_unique.rs | 5 +- .../src/dataframe/command/boolean/not.rs | 5 +- .../dataframe/command/boolean/otherwise.rs | 2 + .../src/dataframe/command/boolean/set.rs | 5 +- .../src/dataframe/command/boolean/when.rs | 2 + .../src/dataframe/command/core/cache.rs | 2 + .../src/dataframe/command/core/columns.rs | 5 +- .../src/dataframe/command/core/fetch.rs | 2 + .../src/dataframe/command/core/open.rs | 15 ++- .../src/dataframe/command/core/profile.rs | 2 + .../src/dataframe/command/core/save/mod.rs | 84 ++++++++++++++--- .../src/dataframe/command/core/summary.rs | 5 +- .../src/dataframe/command/core/to_df.rs | 2 + .../src/dataframe/command/core/to_lazy.rs | 2 + .../src/dataframe/command/data/alias.rs | 2 + .../src/dataframe/command/data/arg_where.rs | 4 +- .../src/dataframe/command/data/cast.rs | 2 + .../src/dataframe/command/data/col.rs | 4 +- .../src/dataframe/command/data/concat.rs | 5 +- .../src/dataframe/command/data/cut.rs | 5 +- .../src/dataframe/command/data/drop.rs | 5 +- .../dataframe/command/data/drop_duplicates.rs | 5 +- .../src/dataframe/command/data/drop_nulls.rs | 5 +- .../src/dataframe/command/data/explode.rs | 5 +- .../src/dataframe/command/data/fill_nan.rs | 2 + .../src/dataframe/command/data/fill_null.rs | 2 + .../src/dataframe/command/data/filter.rs | 5 +- .../src/dataframe/command/data/filter_with.rs | 2 + .../src/dataframe/command/data/first.rs | 2 + .../src/dataframe/command/data/flatten.rs | 5 +- .../src/dataframe/command/data/get.rs | 5 +- .../src/dataframe/command/data/join.rs | 2 + .../src/dataframe/command/data/last.rs | 2 + .../src/dataframe/command/data/len.rs | 4 +- .../src/dataframe/command/data/pivot.rs | 2 + .../src/dataframe/command/data/qcut.rs | 5 +- .../src/dataframe/command/data/query_df.rs | 7 +- .../src/dataframe/command/data/rename.rs | 2 + .../src/dataframe/command/data/reverse.rs | 2 + .../src/dataframe/command/data/select.rs | 2 + .../src/dataframe/command/data/shift.rs | 2 + .../src/dataframe/command/data/slice.rs | 5 +- .../dataframe/command/data/sort_by_expr.rs | 2 + .../src/dataframe/command/data/take.rs | 5 +- .../src/dataframe/command/data/unique.rs | 2 + .../src/dataframe/command/data/unnest.rs | 2 + .../src/dataframe/command/data/unpivot.rs | 2 + .../src/dataframe/command/data/with_column.rs | 2 + .../src/dataframe/command/datetime/as_date.rs | 5 +- .../dataframe/command/datetime/as_datetime.rs | 5 +- .../command/datetime/convert_time_zone.rs | 2 + .../dataframe/command/datetime/datepart.rs | 2 + .../src/dataframe/command/datetime/get_day.rs | 21 +++-- .../dataframe/command/datetime/get_hour.rs | 5 +- .../dataframe/command/datetime/get_minute.rs | 5 +- .../dataframe/command/datetime/get_month.rs | 5 +- .../command/datetime/get_nanosecond.rs | 5 +- .../dataframe/command/datetime/get_ordinal.rs | 5 +- .../dataframe/command/datetime/get_second.rs | 5 +- .../dataframe/command/datetime/get_week.rs | 5 +- .../dataframe/command/datetime/get_weekday.rs | 5 +- .../dataframe/command/datetime/get_year.rs | 5 +- .../command/datetime/replace_time_zone.rs | 2 + .../dataframe/command/datetime/strftime.rs | 2 + .../dataframe/command/datetime/truncate.rs | 5 +- .../dataframe/command/integer/to_decimal.rs | 2 + .../dataframe/command/integer/to_integer.rs | 2 + .../src/dataframe/command/list/contains.rs | 2 + .../dataframe/command/string/concat_str.rs | 4 +- .../src/dataframe/command/string/contains.rs | 2 + .../src/dataframe/command/string/replace.rs | 2 + .../dataframe/command/string/replace_all.rs | 2 + .../src/dataframe/command/string/str_join.rs | 2 + .../dataframe/command/string/str_lengths.rs | 2 + .../src/dataframe/command/string/str_slice.rs | 2 + .../src/dataframe/command/string/str_split.rs | 2 + .../command/string/str_strip_chars.rs | 2 + .../dataframe/command/string/to_lowercase.rs | 2 + .../dataframe/command/string/to_uppercase.rs | 2 + 110 files changed, 452 insertions(+), 165 deletions(-) create mode 100644 crates/nu-command/src/debug/util.rs diff --git a/crates/nu-command/src/debug/metadata.rs b/crates/nu-command/src/debug/metadata.rs index 61b038fd04..e65afd5d83 100644 --- a/crates/nu-command/src/debug/metadata.rs +++ b/crates/nu-command/src/debug/metadata.rs @@ -1,6 +1,7 @@ +use super::util::{build_metadata_record, extend_record_with_metadata}; use nu_engine::command_prelude::*; use nu_protocol::{ - DataSource, PipelineMetadata, + PipelineMetadata, ast::{Expr, Expression}, }; @@ -56,56 +57,40 @@ impl Command for Metadata { } => { let origin = stack.get_var_with_origin(*var_id, *span)?; - Ok( - build_metadata_record(&origin, input.metadata().as_ref(), head) - .into_pipeline_data(), + Ok(build_metadata_record_value( + &origin, + input.metadata().as_ref(), + head, ) + .into_pipeline_data()) } _ => { let val: Value = call.req(engine_state, stack, 0)?; - Ok(build_metadata_record(&val, input.metadata().as_ref(), head) - .into_pipeline_data()) + Ok( + build_metadata_record_value(&val, input.metadata().as_ref(), head) + .into_pipeline_data(), + ) } } } else { let val: Value = call.req(engine_state, stack, 0)?; - Ok(build_metadata_record(&val, input.metadata().as_ref(), head) - .into_pipeline_data()) + Ok( + build_metadata_record_value(&val, input.metadata().as_ref(), head) + .into_pipeline_data(), + ) } } Some(_) => { let val: Value = call.req(engine_state, stack, 0)?; - Ok(build_metadata_record(&val, input.metadata().as_ref(), head) - .into_pipeline_data()) - } - None => { - let mut record = Record::new(); - if let Some(x) = input.metadata().as_ref() { - match x { - PipelineMetadata { - data_source: DataSource::Ls, - .. - } => record.push("source", Value::string("ls", head)), - PipelineMetadata { - data_source: DataSource::HtmlThemes, - .. - } => record.push("source", Value::string("into html --list", head)), - PipelineMetadata { - data_source: DataSource::FilePath(path), - .. - } => record.push( - "source", - Value::string(path.to_string_lossy().to_string(), head), - ), - _ => {} - } - if let Some(ref content_type) = x.content_type { - record.push("content_type", Value::string(content_type, head)); - } - } - - Ok(Value::record(record, head).into_pipeline_data()) + Ok( + build_metadata_record_value(&val, input.metadata().as_ref(), head) + .into_pipeline_data(), + ) } + None => Ok( + Value::record(build_metadata_record(input.metadata().as_ref(), head), head) + .into_pipeline_data(), + ), } } @@ -125,7 +110,11 @@ impl Command for Metadata { } } -fn build_metadata_record(arg: &Value, metadata: Option<&PipelineMetadata>, head: Span) -> Value { +fn build_metadata_record_value( + arg: &Value, + metadata: Option<&PipelineMetadata>, + head: Span, +) -> Value { let mut record = Record::new(); let span = arg.span(); @@ -140,31 +129,7 @@ fn build_metadata_record(arg: &Value, metadata: Option<&PipelineMetadata>, head: ), ); - if let Some(x) = metadata { - match x { - PipelineMetadata { - data_source: DataSource::Ls, - .. - } => record.push("source", Value::string("ls", head)), - PipelineMetadata { - data_source: DataSource::HtmlThemes, - .. - } => record.push("source", Value::string("into html --list", head)), - PipelineMetadata { - data_source: DataSource::FilePath(path), - .. - } => record.push( - "source", - Value::string(path.to_string_lossy().to_string(), head), - ), - _ => {} - } - if let Some(ref content_type) = x.content_type { - record.push("content_type", Value::string(content_type, head)); - } - } - - Value::record(record, head) + Value::record(extend_record_with_metadata(record, metadata, head), head) } #[cfg(test)] diff --git a/crates/nu-command/src/debug/metadata_access.rs b/crates/nu-command/src/debug/metadata_access.rs index 6abf16dc2a..297fc2e613 100644 --- a/crates/nu-command/src/debug/metadata_access.rs +++ b/crates/nu-command/src/debug/metadata_access.rs @@ -1,9 +1,11 @@ use nu_engine::{command_prelude::*, get_eval_block_with_early_return}; use nu_protocol::{ - DataSource, PipelineData, PipelineMetadata, ShellError, Signature, SyntaxShape, Type, Value, + PipelineData, ShellError, Signature, SyntaxShape, Type, Value, engine::{Call, Closure, Command, EngineState, Stack}, }; +use super::util::build_metadata_record; + #[derive(Clone)] pub struct MetadataAccess; @@ -40,7 +42,10 @@ impl Command for MetadataAccess { // `ClosureEvalOnce` is not used as it uses `Stack::captures_to_stack` rather than // `Stack::captures_to_stack_preserve_out_dest`. This command shouldn't collect streams let mut callee_stack = caller_stack.captures_to_stack_preserve_out_dest(closure.captures); - let metadata_record = build_metadata_record(input.metadata().as_ref(), call.head); + let metadata_record = Value::record( + build_metadata_record(input.metadata().as_ref(), call.head), + call.head, + ); if let Some(var_id) = block.signature.get_positional(0).and_then(|var| var.var_id) { callee_stack.add_var(var_id, metadata_record) @@ -64,36 +69,6 @@ impl Command for MetadataAccess { } } -fn build_metadata_record(metadata: Option<&PipelineMetadata>, head: Span) -> Value { - let mut record = Record::new(); - - if let Some(x) = metadata { - match x { - PipelineMetadata { - data_source: DataSource::Ls, - .. - } => record.push("source", Value::string("ls", head)), - PipelineMetadata { - data_source: DataSource::HtmlThemes, - .. - } => record.push("source", Value::string("into html --list", head)), - PipelineMetadata { - data_source: DataSource::FilePath(path), - .. - } => record.push( - "source", - Value::string(path.to_string_lossy().to_string(), head), - ), - _ => {} - } - if let Some(ref content_type) = x.content_type { - record.push("content_type", Value::string(content_type, head)); - } - } - - Value::record(record, head) -} - #[cfg(test)] mod test { use crate::ToJson; diff --git a/crates/nu-command/src/debug/mod.rs b/crates/nu-command/src/debug/mod.rs index 50ffa46b9a..521e8948c3 100644 --- a/crates/nu-command/src/debug/mod.rs +++ b/crates/nu-command/src/debug/mod.rs @@ -9,6 +9,7 @@ mod metadata_access; mod metadata_set; mod profile; mod timeit; +mod util; mod view; mod view_blocks; mod view_files; diff --git a/crates/nu-command/src/debug/util.rs b/crates/nu-command/src/debug/util.rs new file mode 100644 index 0000000000..07f5f38df9 --- /dev/null +++ b/crates/nu-command/src/debug/util.rs @@ -0,0 +1,34 @@ +use nu_protocol::{DataSource, PipelineMetadata, Record, Span, Value}; + +pub fn extend_record_with_metadata( + mut record: Record, + metadata: Option<&PipelineMetadata>, + head: Span, +) -> Record { + if let Some(PipelineMetadata { + data_source, + content_type, + }) = metadata + { + match data_source { + DataSource::Ls => record.push("source", Value::string("ls", head)), + DataSource::HtmlThemes => { + record.push("source", Value::string("into html --list", head)) + } + DataSource::FilePath(path) => record.push( + "source", + Value::string(path.to_string_lossy().to_string(), head), + ), + DataSource::None => {} + } + if let Some(content_type) = content_type { + record.push("content_type", Value::string(content_type, head)); + } + }; + + record +} + +pub fn build_metadata_record(metadata: Option<&PipelineMetadata>, head: Span) -> Record { + extend_record_with_metadata(Record::new(), metadata, head) +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs index ce3d916e35..e05c1094a7 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/agg_groups.rs @@ -61,12 +61,14 @@ impl PluginCommand for ExprAggGroups { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs index 0692bbe1c4..987cf22922 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs @@ -127,6 +127,7 @@ impl PluginCommand for LazyAggregate { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let vals: Vec = call.rest(0)?; let value = Value::list(vals, call.head); let expressions = NuExpression::extract_exprs(plugin, value)?; @@ -153,6 +154,7 @@ impl PluginCommand for LazyAggregate { let lazy = NuLazyFrame::new(false, polars.agg(&expressions)); lazy.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs index 97b7ebaef9..38a4bb8393 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/count.rs @@ -56,12 +56,14 @@ impl PluginCommand for ExprCount { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/cumulative.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/cumulative.rs index c9fe221347..1587d7eb56 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/cumulative.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/cumulative.rs @@ -156,6 +156,7 @@ impl PluginCommand for Cumulative { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; let cum_type: Spanned = call.req(0)?; let cum_type = CumulativeType::from_str(&cum_type.item, cum_type.span)?; @@ -177,6 +178,7 @@ impl PluginCommand for Cumulative { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs index defdcb6402..67376ca8b1 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs @@ -100,6 +100,7 @@ impl PluginCommand for ToLazyGroupBy { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let vals: Vec = call.rest(0)?; let expr_value = Value::list(vals, call.head); let expressions = NuExpression::extract_exprs(plugin, expr_value)?; @@ -117,7 +118,9 @@ impl PluginCommand for ToLazyGroupBy { let pipeline_value = input.into_value(call.head)?; let lazy = NuLazyFrame::try_from_value_coerce(plugin, &pipeline_value)?; - command(plugin, engine, call, lazy, expressions).map_err(LabeledError::from) + command(plugin, engine, call, lazy, expressions) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs index 2c5bd1f0e5..cf3e828c43 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/implode.rs @@ -56,12 +56,14 @@ impl PluginCommand for ExprImplode { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/max.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/max.rs index cf0380530c..068bd44bfd 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/max.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/max.rs @@ -90,6 +90,7 @@ impl PluginCommand for ExprMax { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), @@ -105,6 +106,7 @@ impl PluginCommand for ExprMax { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/mean.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/mean.rs index 8b848d9a1f..e4593c158c 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/mean.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/mean.rs @@ -90,6 +90,7 @@ impl PluginCommand for ExprMean { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), @@ -105,6 +106,7 @@ impl PluginCommand for ExprMean { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/median.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/median.rs index bd9833bd10..1101ca19ba 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/median.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/median.rs @@ -91,6 +91,7 @@ impl PluginCommand for LazyMedian { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command(plugin, engine, call, df.lazy()), @@ -109,6 +110,7 @@ impl PluginCommand for LazyMedian { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/min.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/min.rs index be6456e381..90c44563ea 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/min.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/min.rs @@ -90,6 +90,7 @@ impl PluginCommand for ExprMin { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), @@ -105,6 +106,7 @@ impl PluginCommand for ExprMin { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/n_null.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/n_null.rs index 4e69d9cdf0..0a6889079a 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/n_null.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/n_null.rs @@ -56,7 +56,10 @@ impl PluginCommand for NNull { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/n_unique.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/n_unique.rs index d9be7a405a..d9e4a1406d 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/n_unique.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/n_unique.rs @@ -70,6 +70,7 @@ impl PluginCommand for NUnique { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command(plugin, engine, call, df), @@ -90,6 +91,7 @@ impl PluginCommand for NUnique { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/over.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/over.rs index 4aaba37004..c284227791 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/over.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/over.rs @@ -83,6 +83,7 @@ impl PluginCommand for Over { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let vals: Vec = call.rest(0)?; let expr_value = Value::list(vals, call.head); let expressions = NuExpression::extract_exprs(plugin, expr_value)?; @@ -103,6 +104,7 @@ impl PluginCommand for Over { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/quantile.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/quantile.rs index 7594afad1f..00d76f5054 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/quantile.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/quantile.rs @@ -99,6 +99,7 @@ impl PluginCommand for LazyQuantile { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; let quantile: f64 = call.req(0)?; match PolarsPluginObject::try_from_value(plugin, &value)? { @@ -123,6 +124,7 @@ impl PluginCommand for LazyQuantile { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/rolling.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/rolling.rs index 17682d08db..5e5b54187b 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/rolling.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/rolling.rs @@ -119,7 +119,10 @@ impl PluginCommand for Rolling { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/std.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/std.rs index 575d624a63..a32a9741ff 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/std.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/std.rs @@ -83,6 +83,7 @@ impl PluginCommand for ExprStd { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), @@ -98,6 +99,7 @@ impl PluginCommand for ExprStd { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/sum.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/sum.rs index e04057d48f..7815d52fcf 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/sum.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/sum.rs @@ -90,6 +90,7 @@ impl PluginCommand for ExprSum { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), @@ -105,6 +106,7 @@ impl PluginCommand for ExprSum { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/value_counts.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/value_counts.rs index 429eaf6b9b..a47aaf0fde 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/value_counts.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/value_counts.rs @@ -74,7 +74,10 @@ impl PluginCommand for ValueCount { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/var.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/var.rs index 9facb40c5d..f310bcc167 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/var.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/var.rs @@ -84,6 +84,7 @@ impl PluginCommand for ExprVar { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), @@ -99,6 +100,7 @@ impl PluginCommand for ExprVar { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/all_false.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/all_false.rs index 233d4abd63..424987a42f 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/all_false.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/all_false.rs @@ -73,7 +73,10 @@ impl PluginCommand for AllFalse { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/all_true.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/all_true.rs index 297439dfa6..14dfe92b4a 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/all_true.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/all_true.rs @@ -73,7 +73,10 @@ impl PluginCommand for AllTrue { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/arg_true.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/arg_true.rs index f602757b31..4078ba72a7 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/arg_true.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/arg_true.rs @@ -60,7 +60,10 @@ impl PluginCommand for ArgTrue { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/expr_not.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/expr_not.rs index 25bf6bdeb6..6da37575d2 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/expr_not.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/expr_not.rs @@ -67,12 +67,14 @@ impl PluginCommand for ExprNot { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/is_duplicated.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/is_duplicated.rs index f0b08f288c..7fe2cfae0d 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/is_duplicated.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/is_duplicated.rs @@ -87,7 +87,10 @@ impl PluginCommand for IsDuplicated { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/is_in.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/is_in.rs index d7bb912278..2bfcc0dafb 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/is_in.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/is_in.rs @@ -114,6 +114,7 @@ impl PluginCommand for ExprIsIn { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), @@ -131,6 +132,7 @@ impl PluginCommand for ExprIsIn { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/is_not_null.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/is_not_null.rs index 49a351ff20..e856f64df0 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/is_not_null.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/is_not_null.rs @@ -78,6 +78,7 @@ impl PluginCommand for IsNotNull { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command(plugin, engine, call, df), @@ -98,6 +99,7 @@ impl PluginCommand for IsNotNull { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/is_null.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/is_null.rs index 5acf8f2f44..8505cb75b8 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/is_null.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/is_null.rs @@ -80,6 +80,7 @@ impl PluginCommand for IsNull { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command(plugin, engine, call, df), @@ -100,6 +101,7 @@ impl PluginCommand for IsNull { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/is_unique.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/is_unique.rs index 5a6f4322eb..8f714d1221 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/is_unique.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/is_unique.rs @@ -87,7 +87,10 @@ impl PluginCommand for IsUnique { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/not.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/not.rs index 35f9ca2bcb..bd55070ec6 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/not.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/not.rs @@ -61,8 +61,11 @@ impl PluginCommand for NotSeries { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; - command(plugin, engine, call, df).map_err(LabeledError::from) + command(plugin, engine, call, df) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/otherwise.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/otherwise.rs index aeddcf5d88..b915af4f88 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/otherwise.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/otherwise.rs @@ -95,6 +95,7 @@ impl PluginCommand for ExprOtherwise { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let otherwise_predicate: Value = call.req(0)?; let otherwise_predicate = NuExpression::try_from_value(plugin, &otherwise_predicate)?; @@ -108,6 +109,7 @@ impl PluginCommand for ExprOtherwise { complete .to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/set.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/set.rs index c1e347381d..3fe013b838 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/set.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/set.rs @@ -72,7 +72,10 @@ impl PluginCommand for SetSeries { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/boolean/when.rs b/crates/nu_plugin_polars/src/dataframe/command/boolean/when.rs index a38c5f8cdb..351277cd7e 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/boolean/when.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/boolean/when.rs @@ -111,6 +111,7 @@ impl PluginCommand for ExprWhen { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let when_predicate: Value = call.req(0)?; let when_predicate = NuExpression::try_from_value(plugin, &when_predicate)?; @@ -137,6 +138,7 @@ impl PluginCommand for ExprWhen { when_then .to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/cache.rs b/crates/nu_plugin_polars/src/dataframe/command/core/cache.rs index d12c3d2776..1c5f460f95 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/cache.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/cache.rs @@ -56,11 +56,13 @@ impl PluginCommand for LazyCache { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let lazy = NuLazyFrame::try_from_pipeline_coerce(plugin, input, call.head) .map_err(LabeledError::from)?; let lazy = NuLazyFrame::new(lazy.from_eager, lazy.to_polars().cache()); lazy.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/columns.rs b/crates/nu_plugin_polars/src/dataframe/command/core/columns.rs index 5a89cfaed1..8dc36f41f6 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/columns.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/columns.rs @@ -44,7 +44,10 @@ impl PluginCommand for ColumnsDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, call, input).map_err(|e| e.into()) + let metadata = input.metadata(); + command(plugin, call, input) + .map_err(|e| e.into()) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/fetch.rs b/crates/nu_plugin_polars/src/dataframe/command/core/fetch.rs index cd52f16dfc..de8bdedfab 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/fetch.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/fetch.rs @@ -66,6 +66,7 @@ impl PluginCommand for LazyFetch { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let rows: i64 = call.req(0)?; let value = input.into_value(call.head)?; let lazy = NuLazyFrame::try_from_value_coerce(plugin, &value)?; @@ -87,6 +88,7 @@ impl PluginCommand for LazyFetch { eager .to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/open.rs b/crates/nu_plugin_polars/src/dataframe/command/core/open.rs index 0cec4efd76..7f450794d1 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/open.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/open.rs @@ -9,8 +9,8 @@ use nu_utils::perf; use nu_plugin::{EvaluatedCall, PluginCommand}; use nu_protocol::{ - Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Spanned, - SyntaxShape, Type, Value, + Category, DataSource, Example, LabeledError, PipelineData, PipelineMetadata, ShellError, + Signature, Span, Spanned, SyntaxShape, Type, Value, shell_error::{self, io::IoError}, }; @@ -131,8 +131,8 @@ impl PluginCommand for OpenDataFrame { plugin: &Self::Plugin, engine: &nu_plugin::EngineInterface, call: &nu_plugin::EvaluatedCall, - _input: nu_protocol::PipelineData, - ) -> Result { + _input: PipelineData, + ) -> Result { command(plugin, engine, call).map_err(|e| e.into()) } } @@ -166,6 +166,11 @@ fn command( let hive_options = build_hive_options(plugin, call)?; + let uri = spanned_file.item.clone(); + let data_source = DataSource::FilePath(uri.into()); + + let metadata = PipelineMetadata::default().with_data_source(data_source); + match type_option { Some((ext, blamed)) => match PolarsFileType::from(ext.as_str()) { PolarsFileType::Csv | PolarsFileType::Tsv => { @@ -200,7 +205,7 @@ fn command( "File without extension", ))), } - .map(|value| PipelineData::Value(value, None)) + .map(|value| PipelineData::Value(value, Some(metadata))) } fn from_parquet( diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/profile.rs b/crates/nu_plugin_polars/src/dataframe/command/core/profile.rs index e05e031391..41355f7f79 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/profile.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/profile.rs @@ -63,6 +63,7 @@ The units of the timings are microseconds."# call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_lazy(plugin, engine, call, df.lazy()), @@ -73,6 +74,7 @@ The units of the timings are microseconds."# )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/save/mod.rs b/crates/nu_plugin_polars/src/dataframe/command/core/save/mod.rs index ec36fe3f86..52f00f2b6d 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/save/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/save/mod.rs @@ -15,8 +15,8 @@ use crate::{ use log::debug; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ - Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Spanned, - SyntaxShape, Type, + Category, DataSource, Example, LabeledError, PipelineData, PipelineMetadata, ShellError, + Signature, Span, Spanned, SyntaxShape, Type, shell_error::{self, io::IoError}, }; use polars::error::PolarsError; @@ -112,11 +112,20 @@ impl PluginCommand for SaveDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let spanned_file: Spanned = call.req(0)?; + debug!("file: {}", spanned_file.item); + + let metadata = input.metadata(); let value = input.into_value(call.head)?; + check_writing_into_source_file( + metadata.as_ref(), + &spanned_file.as_ref().map(PathBuf::from), + )?; + match PolarsPluginObject::try_from_value(plugin, &value)? { po @ PolarsPluginObject::NuDataFrame(_) | po @ PolarsPluginObject::NuLazyFrame(_) => { - command(plugin, engine, call, po) + command(plugin, engine, call, po, spanned_file) } _ => Err(cant_convert_err( &value, @@ -132,10 +141,8 @@ fn command( engine: &EngineInterface, call: &EvaluatedCall, polars_object: PolarsPluginObject, + spanned_file: Spanned, ) -> Result { - let spanned_file: Spanned = call.req(0)?; - debug!("file: {}", spanned_file.item); - let resource = Resource::new(plugin, engine, &spanned_file)?; let type_option: Option<(String, Span)> = call .get_flag("type")? @@ -223,6 +230,28 @@ fn command( Ok(PipelineData::empty()) } +fn check_writing_into_source_file( + metadata: Option<&PipelineMetadata>, + dest: &Spanned, +) -> Result<(), ShellError> { + let Some(DataSource::FilePath(source)) = metadata.map(|meta| &meta.data_source) else { + return Ok(()); + }; + + if &dest.item == source { + return Err(write_into_source_error(dest.span)); + } + + Ok(()) +} + +fn write_into_source_error(span: Span) -> ShellError { + polars_file_save_error( + PolarsError::InvalidOperation("attempted to save into source".into()), + span, + ) +} + pub(crate) fn polars_file_save_error(e: PolarsError, span: Span) -> ShellError { ShellError::GenericError { error: format!("Error saving file: {e}"), @@ -247,17 +276,13 @@ pub fn unknown_file_save_error(span: Span) -> ShellError { pub(crate) mod test { use nu_plugin_test_support::PluginTest; use nu_protocol::{Span, Value}; + use tempfile::TempDir; use uuid::Uuid; use crate::PolarsPlugin; - fn test_save(cmd: &'static str, extension: &str) -> Result<(), Box> { + fn tmp_dir_sandbox() -> Result<(TempDir, PluginTest), Box> { let tmp_dir = tempfile::tempdir()?; - let mut tmp_file = tmp_dir.path().to_owned(); - tmp_file.push(format!("{}.{}", Uuid::new_v4(), extension)); - let tmp_file_str = tmp_file.to_str().expect("should be able to get file path"); - - let cmd = format!("{cmd} {tmp_file_str}"); let mut plugin_test = PluginTest::new("polars", PolarsPlugin::new()?.into())?; plugin_test.engine_state_mut().add_env_var( "PWD".to_string(), @@ -270,6 +295,17 @@ pub(crate) mod test { Span::test_data(), ), ); + + Ok((tmp_dir, plugin_test)) + } + + fn test_save(cmd: &'static str, extension: &str) -> Result<(), Box> { + let (tmp_dir, mut plugin_test) = tmp_dir_sandbox()?; + let mut tmp_file = tmp_dir.path().to_owned(); + tmp_file.push(format!("{}.{}", Uuid::new_v4(), extension)); + let tmp_file_str = tmp_file.to_str().expect("should be able to get file path"); + + let cmd = format!("{cmd} {tmp_file_str}"); let _pipeline_data = plugin_test.eval(&cmd)?; assert!(tmp_file.exists()); @@ -290,4 +326,28 @@ pub(crate) mod test { extension, ) } + + #[test] + fn test_write_to_source_guard() -> Result<(), Box> { + let (tmp_dir, mut plugin_test) = tmp_dir_sandbox()?; + let mut tmp_file = tmp_dir.path().to_owned(); + dbg!(&tmp_dir); + tmp_file.push(format!("{}.{}", Uuid::new_v4(), "parquet")); + let tmp_file_str = tmp_file.to_str().expect("Should be able to get file path"); + + let _setup = plugin_test.eval(&format!( + "[1 2 3] | polars into-df | polars save {tmp_file_str}", + ))?; + + let output = plugin_test.eval(&format!( + "polars open {tmp_file_str} | polars save {tmp_file_str}" + )); + + assert!(output.is_err_and(|e| { + e.to_string() + .contains("Error saving file: attempted to save into source") + })); + + Ok(()) + } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/summary.rs b/crates/nu_plugin_polars/src/dataframe/command/core/summary.rs index c507fa273f..b4ee5f5760 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/summary.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/summary.rs @@ -112,7 +112,10 @@ impl PluginCommand for Summary { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/to_df.rs b/crates/nu_plugin_polars/src/dataframe/command/core/to_df.rs index 55473c5059..35efefe393 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/to_df.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/to_df.rs @@ -251,6 +251,7 @@ impl PluginCommand for ToDataFrame { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let maybe_schema = call .get_flag("schema")? .map(|schema| NuSchema::try_from_value(plugin, &schema)) @@ -301,6 +302,7 @@ impl PluginCommand for ToDataFrame { df.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/to_lazy.rs b/crates/nu_plugin_polars/src/dataframe/command/core/to_lazy.rs index 5a2dbf6c44..23bc354dc6 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/to_lazy.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/to_lazy.rs @@ -75,6 +75,7 @@ impl PluginCommand for ToLazyFrame { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let maybe_schema = call .get_flag("schema")? .map(|schema| NuSchema::try_from_value(plugin, &schema)) @@ -88,6 +89,7 @@ impl PluginCommand for ToLazyFrame { lazy.cache(plugin, engine, call.head)?.into_value(call.head), None, )) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/alias.rs b/crates/nu_plugin_polars/src/dataframe/command/data/alias.rs index 8919a802aa..d61f921c65 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/alias.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/alias.rs @@ -64,6 +64,7 @@ impl PluginCommand for ExprAlias { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let alias: String = call.req(0)?; let expr = NuExpression::try_from_pipeline(plugin, input, call.head)?; @@ -71,6 +72,7 @@ impl PluginCommand for ExprAlias { expr.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/arg_where.rs b/crates/nu_plugin_polars/src/dataframe/command/data/arg_where.rs index 3579b5d0a2..05aed56e4f 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/arg_where.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/arg_where.rs @@ -58,13 +58,15 @@ impl PluginCommand for ExprArgWhere { plugin: &Self::Plugin, engine: &EngineInterface, call: &EvaluatedCall, - _input: PipelineData, + input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value: Value = call.req(0)?; let expr = NuExpression::try_from_value(plugin, &value)?; let expr: NuExpression = arg_where(expr.into_polars()).into(); expr.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/cast.rs b/crates/nu_plugin_polars/src/dataframe/command/data/cast.rs index 9f3b4fd6fa..7274ac84ab 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/cast.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/cast.rs @@ -90,6 +90,7 @@ impl PluginCommand for CastDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuLazyFrame(lazy) => { @@ -116,6 +117,7 @@ impl PluginCommand for CastDF { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/col.rs b/crates/nu_plugin_polars/src/dataframe/command/data/col.rs index 2411d54083..0432262028 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/col.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/col.rs @@ -119,8 +119,9 @@ impl PluginCommand for ExprCol { plugin: &Self::Plugin, engine: &EngineInterface, call: &EvaluatedCall, - _input: PipelineData, + input: PipelineData, ) -> Result { + let metadata = input.metadata(); let mut names: Vec = vec![call.req(0)?]; names.extend(call.rest(1)?); @@ -144,6 +145,7 @@ impl PluginCommand for ExprCol { expr.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/concat.rs b/crates/nu_plugin_polars/src/dataframe/command/data/concat.rs index 538fcb87b7..6f31788228 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/concat.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/concat.rs @@ -98,8 +98,11 @@ impl PluginCommand for ConcatDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let maybe_df = NuLazyFrame::try_from_pipeline_coerce(plugin, input, call.head).ok(); - command_lazy(plugin, engine, call, maybe_df).map_err(LabeledError::from) + command_lazy(plugin, engine, call, maybe_df) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/cut.rs b/crates/nu_plugin_polars/src/dataframe/command/data/cut.rs index e540c13bf5..268a1f704b 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/cut.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/cut.rs @@ -52,7 +52,10 @@ impl PluginCommand for CutSeries { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(|e| e.into()) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(|e| e.into()) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/drop.rs b/crates/nu_plugin_polars/src/dataframe/command/data/drop.rs index a08faaeb5b..48604d8d4a 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/drop.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/drop.rs @@ -59,7 +59,10 @@ impl PluginCommand for DropDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/drop_duplicates.rs b/crates/nu_plugin_polars/src/dataframe/command/data/drop_duplicates.rs index e6aec2f5dc..9f609699f5 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/drop_duplicates.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/drop_duplicates.rs @@ -72,7 +72,10 @@ impl PluginCommand for DropDuplicates { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/drop_nulls.rs b/crates/nu_plugin_polars/src/dataframe/command/data/drop_nulls.rs index f694a49973..ab127b61c4 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/drop_nulls.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/drop_nulls.rs @@ -99,7 +99,10 @@ impl PluginCommand for DropNulls { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/explode.rs b/crates/nu_plugin_polars/src/dataframe/command/data/explode.rs index ec62913631..1a0c218dcc 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/explode.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/explode.rs @@ -119,7 +119,10 @@ impl PluginCommand for LazyExplode { call: &EvaluatedCall, input: PipelineData, ) -> Result { - explode(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + explode(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/fill_nan.rs b/crates/nu_plugin_polars/src/dataframe/command/data/fill_nan.rs index 5ff88a8b11..a629ccc393 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/fill_nan.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/fill_nan.rs @@ -91,6 +91,7 @@ impl PluginCommand for LazyFillNA { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let fill: Value = call.req(0)?; let value = input.into_value(call.head)?; @@ -119,6 +120,7 @@ impl PluginCommand for LazyFillNA { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/fill_null.rs b/crates/nu_plugin_polars/src/dataframe/command/data/fill_null.rs index 8dff660d5b..b8a9b75e9b 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/fill_null.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/fill_null.rs @@ -68,6 +68,7 @@ impl PluginCommand for LazyFillNull { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let fill: Value = call.req(0)?; let value = input.into_value(call.head)?; @@ -85,6 +86,7 @@ impl PluginCommand for LazyFillNull { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/filter.rs b/crates/nu_plugin_polars/src/dataframe/command/data/filter.rs index f47684c430..4581f8c84a 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/filter.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/filter.rs @@ -109,11 +109,14 @@ impl PluginCommand for LazyFilter { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let expr_value: Value = call.req(0)?; let filter_expr = NuExpression::try_from_value(plugin, &expr_value)?; let pipeline_value = input.into_value(call.head)?; let lazy = NuLazyFrame::try_from_value_coerce(plugin, &pipeline_value)?; - command(plugin, engine, call, lazy, filter_expr).map_err(LabeledError::from) + command(plugin, engine, call, lazy, filter_expr) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/filter_with.rs b/crates/nu_plugin_polars/src/dataframe/command/data/filter_with.rs index eeae6a60dc..efdee6bd5e 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/filter_with.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/filter_with.rs @@ -84,6 +84,7 @@ impl PluginCommand for FilterWith { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), @@ -94,6 +95,7 @@ impl PluginCommand for FilterWith { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/first.rs b/crates/nu_plugin_polars/src/dataframe/command/data/first.rs index cd432a0d78..d09f415ccf 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/first.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/first.rs @@ -97,6 +97,7 @@ impl PluginCommand for FirstDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => { @@ -113,6 +114,7 @@ impl PluginCommand for FirstDF { .map_err(LabeledError::from) } } + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/flatten.rs b/crates/nu_plugin_polars/src/dataframe/command/data/flatten.rs index 2a5190a9da..645f986020 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/flatten.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/flatten.rs @@ -117,7 +117,10 @@ impl PluginCommand for LazyFlatten { call: &EvaluatedCall, input: PipelineData, ) -> Result { - explode(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + explode(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/get.rs b/crates/nu_plugin_polars/src/dataframe/command/data/get.rs index 8c9f1786e8..6e687cbd9c 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/get.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/get.rs @@ -59,7 +59,10 @@ impl PluginCommand for GetDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/join.rs b/crates/nu_plugin_polars/src/dataframe/command/data/join.rs index d2500d759d..cedb9dbab5 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/join.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/join.rs @@ -241,6 +241,7 @@ impl PluginCommand for LazyJoin { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let left = call.has_flag("left")?; let full = call.has_flag("full")?; let cross = call.has_flag("cross")?; @@ -333,6 +334,7 @@ impl PluginCommand for LazyJoin { let lazy = NuLazyFrame::new(from_eager, lazy); lazy.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/last.rs b/crates/nu_plugin_polars/src/dataframe/command/data/last.rs index 62bf9810e6..b1a7f6b0c3 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/last.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/last.rs @@ -72,6 +72,7 @@ impl PluginCommand for LastDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => { @@ -88,6 +89,7 @@ impl PluginCommand for LastDF { .map_err(LabeledError::from) } } + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/len.rs b/crates/nu_plugin_polars/src/dataframe/command/data/len.rs index 62cf685dd0..713121fa11 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/len.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/len.rs @@ -53,11 +53,13 @@ impl PluginCommand for ExprLen { plugin: &Self::Plugin, engine: &EngineInterface, call: &EvaluatedCall, - _input: PipelineData, + input: PipelineData, ) -> Result { + let metadata = input.metadata(); let res: NuExpression = polars::prelude::len().into(); res.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/pivot.rs b/crates/nu_plugin_polars/src/dataframe/command/data/pivot.rs index 672e063f7b..f2eb6b596d 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/pivot.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/pivot.rs @@ -193,6 +193,7 @@ impl PluginCommand for PivotDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); match PolarsPluginObject::try_from_pipeline(plugin, input, call.head)? { PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), PolarsPluginObject::NuLazyFrame(lazy) => { @@ -207,6 +208,7 @@ impl PluginCommand for PivotDF { }), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/qcut.rs b/crates/nu_plugin_polars/src/dataframe/command/data/qcut.rs index 779bd90e75..ccf45b7ed7 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/qcut.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/qcut.rs @@ -53,7 +53,10 @@ impl PluginCommand for QCutSeries { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(|e| e.into()) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(|e| e.into()) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/query_df.rs b/crates/nu_plugin_polars/src/dataframe/command/data/query_df.rs index f9f684695b..8cf0ed0790 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/query_df.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/query_df.rs @@ -63,12 +63,15 @@ impl PluginCommand for QueryDf { fn run( &self, - plugin: &PolarsPlugin, + plugin: &Self::Plugin, engine: &EngineInterface, call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/rename.rs b/crates/nu_plugin_polars/src/dataframe/command/data/rename.rs index 0445b20644..62e332de84 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/rename.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/rename.rs @@ -119,6 +119,7 @@ impl PluginCommand for RenameDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value).map_err(LabeledError::from)? { PolarsPluginObject::NuDataFrame(df) => { @@ -130,6 +131,7 @@ impl PluginCommand for RenameDF { _ => Err(LabeledError::new(format!("Unsupported type: {value:?}")) .with_label("Unsupported Type", call.head)), } + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/reverse.rs b/crates/nu_plugin_polars/src/dataframe/command/data/reverse.rs index d3c7648117..9c82cdf7f7 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/reverse.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/reverse.rs @@ -59,11 +59,13 @@ impl PluginCommand for LazyReverse { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let lazy = NuLazyFrame::try_from_pipeline_coerce(plugin, input, call.head) .map_err(LabeledError::from)?; let lazy = NuLazyFrame::new(lazy.from_eager, lazy.to_polars().reverse()); lazy.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/select.rs b/crates/nu_plugin_polars/src/dataframe/command/data/select.rs index ff43b46615..b4d309c562 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/select.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/select.rs @@ -103,6 +103,7 @@ impl PluginCommand for LazySelect { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let vals: Vec = call.rest(0)?; let expr_value = Value::list(vals, call.head); let expressions = NuExpression::extract_exprs(plugin, expr_value)?; @@ -112,6 +113,7 @@ impl PluginCommand for LazySelect { let lazy: NuLazyFrame = lazy.to_polars().select(&expressions).into(); lazy.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/shift.rs b/crates/nu_plugin_polars/src/dataframe/command/data/shift.rs index a4658962b1..bb58283b7e 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/shift.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/shift.rs @@ -92,6 +92,7 @@ impl PluginCommand for Shift { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { @@ -106,6 +107,7 @@ impl PluginCommand for Shift { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/slice.rs b/crates/nu_plugin_polars/src/dataframe/command/data/slice.rs index 071220d2c0..59b5e50524 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/slice.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/slice.rs @@ -58,7 +58,10 @@ impl PluginCommand for SliceDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/sort_by_expr.rs b/crates/nu_plugin_polars/src/dataframe/command/data/sort_by_expr.rs index 3965725a94..ad8eadd7b9 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/sort_by_expr.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/sort_by_expr.rs @@ -115,6 +115,7 @@ impl PluginCommand for LazySortBy { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let vals: Vec = call.rest(0)?; let expr_value = Value::list(vals, call.head); let expressions = NuExpression::extract_exprs(plugin, expr_value)?; @@ -160,6 +161,7 @@ impl PluginCommand for LazySortBy { ); lazy.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/take.rs b/crates/nu_plugin_polars/src/dataframe/command/data/take.rs index bf63a4b670..f3019690c5 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/take.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/take.rs @@ -89,7 +89,10 @@ impl PluginCommand for TakeDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/unique.rs b/crates/nu_plugin_polars/src/dataframe/command/data/unique.rs index 97f3ea30e4..e4fdf90dec 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/unique.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/unique.rs @@ -176,6 +176,7 @@ impl PluginCommand for Unique { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { @@ -200,6 +201,7 @@ impl PluginCommand for Unique { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/unnest.rs b/crates/nu_plugin_polars/src/dataframe/command/data/unnest.rs index 59c86bec2d..3c0ff2f011 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/unnest.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/unnest.rs @@ -83,6 +83,7 @@ impl PluginCommand for UnnestDF { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); match PolarsPluginObject::try_from_pipeline(plugin, input, call.head)? { PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), @@ -95,6 +96,7 @@ impl PluginCommand for UnnestDF { }), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/unpivot.rs b/crates/nu_plugin_polars/src/dataframe/command/data/unpivot.rs index 4fe4c70ca5..0883cf5af4 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/unpivot.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/unpivot.rs @@ -187,6 +187,7 @@ impl PluginCommand for Unpivot { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); match PolarsPluginObject::try_from_pipeline(plugin, input, call.head)? { PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), @@ -199,6 +200,7 @@ impl PluginCommand for Unpivot { }), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs b/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs index b8dcc12ed9..38bafc9f72 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs @@ -148,6 +148,7 @@ impl PluginCommand for WithColumn { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), @@ -160,6 +161,7 @@ impl PluginCommand for WithColumn { }), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/as_date.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/as_date.rs index 8cd8f502a6..876bd984a1 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/as_date.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/as_date.rs @@ -183,7 +183,10 @@ impl PluginCommand for AsDate { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/as_datetime.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/as_datetime.rs index fccce2b9e1..7154839b7d 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/as_datetime.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/as_datetime.rs @@ -236,7 +236,10 @@ impl PluginCommand for AsDateTime { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/convert_time_zone.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/convert_time_zone.rs index ccf641d3db..9de83b29c4 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/convert_time_zone.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/convert_time_zone.rs @@ -137,6 +137,7 @@ impl PluginCommand for ConvertTimeZone { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { @@ -152,6 +153,7 @@ impl PluginCommand for ConvertTimeZone { _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/datepart.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/datepart.rs index 995c7ebbee..85a619ab97 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/datepart.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/datepart.rs @@ -131,6 +131,7 @@ impl PluginCommand for ExprDatePart { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let part: Spanned = call.req(0)?; let expr = NuExpression::try_from_pipeline(plugin, input, call.head)?; @@ -158,6 +159,7 @@ impl PluginCommand for ExprDatePart { }.into(); expr.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_day.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_day.rs index edf42364a6..100b7530ef 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_day.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_day.rs @@ -77,6 +77,14 @@ impl PluginCommand for GetDay { ] } + fn extra_description(&self) -> &str { + "" + } + + fn search_terms(&self) -> Vec<&str> { + vec![] + } + fn run( &self, plugin: &Self::Plugin, @@ -84,15 +92,10 @@ impl PluginCommand for GetDay { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) - } - - fn extra_description(&self) -> &str { - "" - } - - fn search_terms(&self) -> Vec<&str> { - vec![] + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_hour.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_hour.rs index f1d930fa31..af41adc45c 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_hour.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_hour.rs @@ -100,7 +100,10 @@ impl PluginCommand for GetHour { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_minute.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_minute.rs index 145adf7759..c5dee03ab6 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_minute.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_minute.rs @@ -84,7 +84,10 @@ impl PluginCommand for GetMinute { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_month.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_month.rs index 8f992026b6..888428a115 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_month.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_month.rs @@ -84,7 +84,10 @@ impl PluginCommand for GetMonth { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_nanosecond.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_nanosecond.rs index ce0274b85f..55ed679164 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_nanosecond.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_nanosecond.rs @@ -84,7 +84,10 @@ impl PluginCommand for GetNanosecond { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_ordinal.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_ordinal.rs index c19c9e7585..c4114cc958 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_ordinal.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_ordinal.rs @@ -84,7 +84,10 @@ impl PluginCommand for GetOrdinal { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_second.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_second.rs index 9ab11ff15a..33e01bb0bd 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_second.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_second.rs @@ -84,7 +84,10 @@ impl PluginCommand for GetSecond { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_week.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_week.rs index 3793a55828..978310edfd 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_week.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_week.rs @@ -84,7 +84,10 @@ impl PluginCommand for GetWeek { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_weekday.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_weekday.rs index f73a088b55..e0a63b86ee 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_weekday.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_weekday.rs @@ -84,7 +84,10 @@ impl PluginCommand for GetWeekDay { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_year.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_year.rs index 1c1a561846..862dd730f1 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_year.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_year.rs @@ -84,7 +84,10 @@ impl PluginCommand for GetYear { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/replace_time_zone.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/replace_time_zone.rs index 5ea5160bc9..51a9be9da8 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/replace_time_zone.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/replace_time_zone.rs @@ -216,6 +216,7 @@ impl PluginCommand for ReplaceTimeZone { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; let ambiguous = match call.get_flag::("ambiguous")? { @@ -269,6 +270,7 @@ impl PluginCommand for ReplaceTimeZone { _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/strftime.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/strftime.rs index 12866d0511..064ec846cc 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/strftime.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/strftime.rs @@ -94,6 +94,7 @@ impl PluginCommand for StrFTime { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), @@ -111,6 +112,7 @@ impl PluginCommand for StrFTime { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/truncate.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/truncate.rs index 8b0afca537..63d8cfb584 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/truncate.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/truncate.rs @@ -150,7 +150,10 @@ impl PluginCommand for Truncate { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let metadata = input.metadata(); + command(plugin, engine, call, input) + .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } fn extra_description(&self) -> &str { diff --git a/crates/nu_plugin_polars/src/dataframe/command/integer/to_decimal.rs b/crates/nu_plugin_polars/src/dataframe/command/integer/to_decimal.rs index 25e45d5760..5673df8c85 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/integer/to_decimal.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/integer/to_decimal.rs @@ -67,12 +67,14 @@ impl PluginCommand for ToDecimal { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuExpression(expr) => command(plugin, engine, call, expr), _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/integer/to_integer.rs b/crates/nu_plugin_polars/src/dataframe/command/integer/to_integer.rs index be0cfd7420..f03804c6bf 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/integer/to_integer.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/integer/to_integer.rs @@ -61,12 +61,14 @@ impl PluginCommand for ToInteger { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuExpression(expr) => command(plugin, engine, call, expr), _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/list/contains.rs b/crates/nu_plugin_polars/src/dataframe/command/list/contains.rs index 4e4cb61dc8..81e6d8a0aa 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/list/contains.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/list/contains.rs @@ -116,12 +116,14 @@ impl PluginCommand for ListContains { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/concat_str.rs b/crates/nu_plugin_polars/src/dataframe/command/string/concat_str.rs index 697a9a09a5..a7fbec6084 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/concat_str.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/concat_str.rs @@ -84,8 +84,9 @@ impl PluginCommand for ExprConcatStr { plugin: &Self::Plugin, engine: &EngineInterface, call: &EvaluatedCall, - _input: PipelineData, + input: PipelineData, ) -> Result { + let metadata = input.metadata(); let separator: String = call.req(0)?; let value: Value = call.req(1)?; @@ -94,6 +95,7 @@ impl PluginCommand for ExprConcatStr { expr.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/contains.rs b/crates/nu_plugin_polars/src/dataframe/command/string/contains.rs index 1b51a5e887..77fe0c7f54 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/contains.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/contains.rs @@ -100,6 +100,7 @@ impl PluginCommand for Contains { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), @@ -118,6 +119,7 @@ impl PluginCommand for Contains { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/replace.rs b/crates/nu_plugin_polars/src/dataframe/command/string/replace.rs index e43540bb66..3878e4c72f 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/replace.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/replace.rs @@ -101,6 +101,7 @@ impl PluginCommand for Replace { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), @@ -118,6 +119,7 @@ impl PluginCommand for Replace { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/replace_all.rs b/crates/nu_plugin_polars/src/dataframe/command/string/replace_all.rs index db9e06e5da..4400e3fc92 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/replace_all.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/replace_all.rs @@ -105,6 +105,7 @@ impl PluginCommand for ReplaceAll { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), @@ -122,6 +123,7 @@ impl PluginCommand for ReplaceAll { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/str_join.rs b/crates/nu_plugin_polars/src/dataframe/command/string/str_join.rs index 92b87ca3c0..df27937029 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/str_join.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/str_join.rs @@ -93,6 +93,7 @@ impl PluginCommand for StrJoin { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), @@ -103,6 +104,7 @@ impl PluginCommand for StrJoin { _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/str_lengths.rs b/crates/nu_plugin_polars/src/dataframe/command/string/str_lengths.rs index d56c2a3b06..87418a9742 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/str_lengths.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/str_lengths.rs @@ -89,6 +89,7 @@ impl PluginCommand for StrLengths { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), @@ -106,6 +107,7 @@ impl PluginCommand for StrLengths { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/str_slice.rs b/crates/nu_plugin_polars/src/dataframe/command/string/str_slice.rs index 0c75bb8a09..b9a250877f 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/str_slice.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/str_slice.rs @@ -117,6 +117,7 @@ impl PluginCommand for StrSlice { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), @@ -134,6 +135,7 @@ impl PluginCommand for StrSlice { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/str_split.rs b/crates/nu_plugin_polars/src/dataframe/command/string/str_split.rs index b315b1ce03..acd82eb267 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/str_split.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/str_split.rs @@ -59,6 +59,7 @@ impl PluginCommand for StrSplit { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let separator = call.req::>(0).and_then(|sep| { let sep_expr = NuExpression::try_from_value(plugin, &sep.item)?; Ok(Spanned { @@ -75,6 +76,7 @@ impl PluginCommand for StrSplit { .into(); res.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/str_strip_chars.rs b/crates/nu_plugin_polars/src/dataframe/command/string/str_strip_chars.rs index e2d40ead48..c453394b5c 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/str_strip_chars.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/str_strip_chars.rs @@ -132,12 +132,14 @@ impl PluginCommand for StrStripChars { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/to_lowercase.rs b/crates/nu_plugin_polars/src/dataframe/command/string/to_lowercase.rs index e06299cf25..7c126a0f3c 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/to_lowercase.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/to_lowercase.rs @@ -88,6 +88,7 @@ impl PluginCommand for ToLowerCase { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), @@ -105,6 +106,7 @@ impl PluginCommand for ToLowerCase { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } } diff --git a/crates/nu_plugin_polars/src/dataframe/command/string/to_uppercase.rs b/crates/nu_plugin_polars/src/dataframe/command/string/to_uppercase.rs index b94b664b06..58c5dcc56c 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/string/to_uppercase.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/string/to_uppercase.rs @@ -92,6 +92,7 @@ impl PluginCommand for ToUpperCase { call: &EvaluatedCall, input: PipelineData, ) -> Result { + let metadata = input.metadata(); let value = input.into_value(call.head)?; match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), @@ -109,6 +110,7 @@ impl PluginCommand for ToUpperCase { )), } .map_err(LabeledError::from) + .map(|pd| pd.set_metadata(metadata)) } }