diff --git a/crates/nu_plugin_polars/src/dataframe/series/string/replace.rs b/crates/nu_plugin_polars/src/dataframe/series/string/replace.rs index 9d7d097f60..c0fa1d2358 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/string/replace.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/string/replace.rs @@ -1,4 +1,10 @@ -use crate::{missing_flag_error, values::CustomValueSupport, PolarsPlugin}; +use crate::{ + missing_flag_error, + values::{ + cant_convert_err, CustomValueSupport, NuExpression, PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use super::super::super::values::{Column, NuDataFrame}; @@ -7,7 +13,7 @@ use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Value, }; -use polars::prelude::{IntoSeries, StringNameSpaceImpl}; +use polars::prelude::{lit, IntoSeries, StringNameSpaceImpl}; #[derive(Clone)] pub struct Replace; @@ -37,33 +43,61 @@ impl PluginCommand for Replace { "replacing string", Some('r'), ) - .input_output_type( - Type::Custom("dataframe".into()), - Type::Custom("dataframe".into()), - ) + .input_output_types(vec![ + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ]) .category(Category::Custom("dataframe".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Replaces string", - example: "[abc abc abc] | polars into-df | polars replace --pattern ab --replace AB", - result: Some( - NuDataFrame::try_from_columns( - vec![Column::new( - "0".to_string(), - vec![ - Value::test_string("ABc"), - Value::test_string("ABc"), - Value::test_string("ABc"), - ], - )], - None, - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + vec![ + Example { + description: "Replaces string in column", + example: + "[[a]; [abc] [abcabc]] | polars into-df | polars select (polars col a | polars replace --pattern ab --replace AB) | polars collect", + result: Some( + NuDataFrame::try_from_columns( + vec![Column::new( + "a".to_string(), + vec![ + Value::test_string("ABc"), + Value::test_string("ABcabc"), + ], + )], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Replaces string", + example: + "[abc abc abc] | polars into-df | polars replace --pattern ab --replace AB", + result: Some( + NuDataFrame::try_from_columns( + vec![Column::new( + "0".to_string(), + vec![ + Value::test_string("ABc"), + Value::test_string("ABc"), + Value::test_string("ABc"), + ], + )], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -73,15 +107,56 @@ impl PluginCommand for Replace { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), + PolarsPluginObject::NuLazyFrame(lazy) => { + command_df(plugin, engine, call, lazy.collect(call.head)?) + } + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } + .map_err(LabeledError::from) } } -fn command( +fn command_expr( plugin: &PolarsPlugin, engine: &EngineInterface, call: &EvaluatedCall, - input: PipelineData, + expr: NuExpression, +) -> Result { + let pattern: String = call + .get_flag("pattern")? + .ok_or_else(|| missing_flag_error("pattern", call.head))?; + let pattern = lit(pattern); + + let replace: String = call + .get_flag("replace")? + .ok_or_else(|| missing_flag_error("replace", call.head))?; + let replace = lit(replace); + + let res: NuExpression = expr + .into_polars() + .str() + .replace(pattern, replace, false) + .into(); + + res.to_pipeline_data(plugin, engine, call.head) +} + +fn command_df( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, ) -> Result { let pattern: String = call .get_flag("pattern")? @@ -90,7 +165,6 @@ fn command( .get_flag("replace")? .ok_or_else(|| missing_flag_error("replace", call.head))?; - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; let series = df.as_series(call.head)?; let chunked = series.str().map_err(|e| ShellError::GenericError { error: "Error conversion to string".into(), diff --git a/crates/nu_plugin_polars/src/dataframe/series/string/replace_all.rs b/crates/nu_plugin_polars/src/dataframe/series/string/replace_all.rs index d88b581451..13799818ee 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/string/replace_all.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/string/replace_all.rs @@ -1,4 +1,10 @@ -use crate::{missing_flag_error, values::CustomValueSupport, PolarsPlugin}; +use crate::{ + missing_flag_error, + values::{ + cant_convert_err, CustomValueSupport, NuExpression, PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use super::super::super::values::{Column, NuDataFrame}; @@ -7,7 +13,7 @@ use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Value, }; -use polars::prelude::{IntoSeries, StringNameSpaceImpl}; +use polars::prelude::{lit, IntoSeries, StringNameSpaceImpl}; #[derive(Clone)] pub struct ReplaceAll; @@ -37,34 +43,62 @@ impl PluginCommand for ReplaceAll { "replacing string", Some('r'), ) - .input_output_type( - Type::Custom("dataframe".into()), - Type::Custom("dataframe".into()), - ) + .input_output_types(vec![ + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ]) .category(Category::Custom("dataframe".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Replaces string", - example: - "[abac abac abac] | polars into-df | polars replace-all --pattern a --replace A", - result: Some( - NuDataFrame::try_from_columns( - vec![Column::new( - "0".to_string(), - vec![ - Value::test_string("AbAc"), - Value::test_string("AbAc"), - Value::test_string("AbAc"), - ], - )], - None, - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + vec![ + Example { + description: "Replaces string in a column", + example: + "[[a]; [abac] [abac] [abac]] | polars into-df | polars select (polars col a | polars replace-all --pattern a --replace A) | polars collect", + result: Some( + NuDataFrame::try_from_columns( + vec![Column::new( + "a".to_string(), + vec![ + Value::test_string("AbAc"), + Value::test_string("AbAc"), + Value::test_string("AbAc"), + ], + )], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Replaces string", + example: + "[abac abac abac] | polars into-df | polars replace-all --pattern a --replace A", + result: Some( + NuDataFrame::try_from_columns( + vec![Column::new( + "0".to_string(), + vec![ + Value::test_string("AbAc"), + Value::test_string("AbAc"), + Value::test_string("AbAc"), + ], + )], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -74,15 +108,55 @@ impl PluginCommand for ReplaceAll { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), + PolarsPluginObject::NuLazyFrame(lazy) => { + command_df(plugin, engine, call, lazy.collect(call.head)?) + } + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } + .map_err(LabeledError::from) } } -fn command( +fn command_expr( plugin: &PolarsPlugin, engine_state: &EngineInterface, call: &EvaluatedCall, - input: PipelineData, + expr: NuExpression, +) -> Result { + let pattern: String = call + .get_flag("pattern")? + .ok_or_else(|| missing_flag_error("pattern", call.head))?; + let pattern = lit(pattern); + let replace: String = call + .get_flag("replace")? + .ok_or_else(|| missing_flag_error("replace", call.head))?; + let replace = lit(replace); + + let res: NuExpression = expr + .into_polars() + .str() + .replace_all(pattern, replace, false) + .into(); + + res.to_pipeline_data(plugin, engine_state, call.head) +} + +fn command_df( + plugin: &PolarsPlugin, + engine_state: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, ) -> Result { let pattern: String = call .get_flag("pattern")? @@ -91,7 +165,6 @@ fn command( .get_flag("replace")? .ok_or_else(|| missing_flag_error("replace", call.head))?; - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; let series = df.as_series(call.head)?; let chunked = series.str().map_err(|e| ShellError::GenericError { error: "Error conversion to string".into(), diff --git a/crates/nu_plugin_polars/src/dataframe/values/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/mod.rs index 41a119cddc..42fb8dace1 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/mod.rs @@ -185,8 +185,8 @@ impl CustomValueType { } } -pub fn cant_convert_err(value: &Value, types: &[PolarsPluginType]) -> ShellError { - let type_string = types +pub fn cant_convert_err(value: &Value, expected_types: &[PolarsPluginType]) -> ShellError { + let type_string = expected_types .iter() .map(ToString::to_string) .collect::>()