diff --git a/crates/nu_plugin_polars/Cargo.toml b/crates/nu_plugin_polars/Cargo.toml index 2d3e1ee312..0c75a1fc56 100644 --- a/crates/nu_plugin_polars/Cargo.toml +++ b/crates/nu_plugin_polars/Cargo.toml @@ -73,6 +73,7 @@ features = [ "strings", "string_to_integer", "streaming", + "temporal", "to_dummies", ] optional = false diff --git a/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs b/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs index 85d3aa84a0..ed48b9b905 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs @@ -1,4 +1,9 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuExpression, PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use super::super::super::values::{Column, NuDataFrame}; @@ -26,34 +31,60 @@ impl PluginCommand for StrFTime { fn signature(&self) -> Signature { Signature::build(self.name()) .required("fmt", SyntaxShape::String, "Format rule") - .input_output_type( - Type::Custom("dataframe".into()), - Type::Custom("dataframe".into()), - ) + .input_output_types(vec![ + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ]) .category(Category::Custom("dataframe".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Formats date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Formats date column as a string", + example: r#"let date = '2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'; + let df = ([[a]; [$date]] | polars into-df); + let df2 = $df | polars with-column [(polars col a | polars strftime "%Y/%m/%d" | polars as b)] | polars collect; + $df2.b"#, + result: Some( + NuDataFrame::try_from_columns( + vec![Column::new( + "b".to_string(), + vec![Value::test_string("2020/08/04")], + )], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Formats date", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); let df = ([$dt $dt] | polars into-df); $df | polars strftime "%Y/%m/%d""#, - result: Some( - NuDataFrame::try_from_columns( - vec![Column::new( - "0".to_string(), - vec![ - Value::test_string("2020/08/04"), - Value::test_string("2020/08/04"), - ], - )], - None, - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + result: Some( + NuDataFrame::try_from_columns( + vec![Column::new( + "0".to_string(), + vec![ + Value::test_string("2020/08/04"), + Value::test_string("2020/08/04"), + ], + )], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -63,19 +94,45 @@ impl PluginCommand for StrFTime { 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 fmt: String = call.req(0)?; + let res: NuExpression = expr.into_polars().dt().strftime(&fmt).into(); + res.to_pipeline_data(plugin, engine, call.head) +} + +fn command_df( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, ) -> Result { let fmt: String = call.req(0)?; - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError {