diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/to_dtype.rs b/crates/nu_plugin_polars/src/dataframe/command/core/to_dtype.rs index f62d0c2b23..1cbab2ae2f 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/to_dtype.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/to_dtype.rs @@ -27,8 +27,8 @@ impl PluginCommand for ToDataType { fn examples(&self) -> Vec { vec![Example { - description: "Convert a string to a specific datatype", - example: r#""i64" | polars into-dtype"#, + description: "Convert a string to a specific datatype and back to a nu object", + example: r#"'i64' | polars into-dtype | polars into-nu"#, result: Some(Value::string("i64", Span::test_data())), }] } @@ -53,3 +53,16 @@ fn command( NuDataType::try_from_pipeline(plugin, input, call.head)? .to_pipeline_data(plugin, engine, call.head) } + +#[cfg(test)] +mod test { + use crate::test::test_polars_plugin_command; + + use super::*; + use nu_protocol::ShellError; + + #[test] + fn test_into_dtype() -> Result<(), ShellError> { + test_polars_plugin_command(&ToDataType) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/to_nu.rs b/crates/nu_plugin_polars/src/dataframe/command/core/to_nu.rs index 07427ca63d..fc0a912e0a 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/to_nu.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/to_nu.rs @@ -5,8 +5,7 @@ use nu_protocol::{ }; use crate::{ - dataframe::values::NuExpression, - values::{CustomValueSupport, NuLazyFrame}, + values::{cant_convert_err, CustomValueSupport, PolarsPluginObject, PolarsPluginType}, PolarsPlugin, }; @@ -39,6 +38,8 @@ impl PluginCommand for ToNu { .input_output_types(vec![ (Type::Custom("expression".into()), Type::Any), (Type::Custom("dataframe".into()), Type::table()), + (Type::Custom("datatype".into()), Type::Any), + (Type::Custom("schema".into()), Type::Any), ]) .category(Category::Custom("dataframe".into())) } @@ -86,31 +87,54 @@ impl PluginCommand for ToNu { fn run( &self, plugin: &Self::Plugin, - _engine: &EngineInterface, + engine: &EngineInterface, call: &EvaluatedCall, input: PipelineData, ) -> Result { - let value = input.into_value(call.head)?; - if NuDataFrame::can_downcast(&value) || NuLazyFrame::can_downcast(&value) { - dataframe_command(plugin, call, value) - } else { - expression_command(plugin, call, value) - } - .map_err(|e| e.into()) + command(plugin, engine, call, input).map_err(LabeledError::from) } } -fn dataframe_command( +fn command( plugin: &PolarsPlugin, + _engine: &EngineInterface, call: &EvaluatedCall, - input: Value, + input: PipelineData, ) -> Result { + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuDataFrame(df) => dataframe_command(call, df), + PolarsPluginObject::NuLazyFrame(lazy) => dataframe_command(call, lazy.collect(call.head)?), + PolarsPluginObject::NuExpression(expr) => { + let value = expr.to_value(call.head)?; + Ok(PipelineData::Value(value, None)) + } + PolarsPluginObject::NuDataType(dt) => { + let value = dt.base_value(call.head)?; + Ok(PipelineData::Value(value, None)) + } + PolarsPluginObject::NuSchema(schema) => { + let value = schema.base_value(call.head)?; + Ok(PipelineData::Value(value, None)) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + PolarsPluginType::NuDataType, + PolarsPluginType::NuSchema, + ], + )), + } +} + +fn dataframe_command(call: &EvaluatedCall, df: NuDataFrame) -> Result { let rows: Option = call.get_flag("rows")?; let tail: bool = call.has_flag("tail")?; let index: bool = call.has_flag("index")?; - let df = NuDataFrame::try_from_value_coerce(plugin, &input, call.head)?; - let values = if tail { df.tail(rows, index, call.head)? } else { @@ -127,17 +151,6 @@ fn dataframe_command( Ok(PipelineData::Value(value, None)) } -fn expression_command( - plugin: &PolarsPlugin, - call: &EvaluatedCall, - input: Value, -) -> Result { - let expr = NuExpression::try_from_value(plugin, &input)?; - let value = expr.to_value(call.head)?; - - Ok(PipelineData::Value(value, None)) -} - #[cfg(test)] mod test { use super::*; diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/to_schema.rs b/crates/nu_plugin_polars/src/dataframe/command/core/to_schema.rs index 7430008653..13d90cc135 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/to_schema.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/to_schema.rs @@ -27,8 +27,8 @@ impl PluginCommand for ToSchema { fn examples(&self) -> Vec { vec![Example { - description: "Convert a record into a schema", - example: r#"{a: str, b: u8} | polars into-schema"#, + description: "Convert a record into a schema and back to a nu object", + example: r#"{a: str, b: u8} | polars into-schema | polars into-nu"#, result: Some(Value::record( record! { "a" => Value::string("str", Span::test_data()), @@ -59,3 +59,16 @@ fn command( NuSchema::try_from_pipeline(plugin, input, call.head)? .to_pipeline_data(plugin, engine, call.head) } + +#[cfg(test)] +mod test { + use crate::test::test_polars_plugin_command; + + use super::*; + use nu_protocol::ShellError; + + #[test] + fn test_into_schema() -> Result<(), ShellError> { + test_polars_plugin_command(&ToSchema) + } +}