diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/mod.rs b/crates/nu_plugin_polars/src/dataframe/command/core/mod.rs index c2b3d26b18..a199e96d36 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/mod.rs @@ -13,6 +13,7 @@ mod to_dtype; mod to_lazy; mod to_nu; mod to_repr; +mod to_schema; pub use self::open::OpenDataFrame; use crate::PolarsPlugin; @@ -42,5 +43,6 @@ pub(crate) fn core_commands() -> Vec 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_lazy.rs b/crates/nu_plugin_polars/src/dataframe/command/core/to_lazy.rs index ccffde21ff..6355fcd3bc 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 @@ -4,8 +4,11 @@ use crate::values::{NuDataFrame, NuLazyFrame}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ - record, Category, Example, LabeledError, PipelineData, Signature, SyntaxShape, Type, Value, + record, Category, Example, LabeledError, PipelineData, Signature, Span, SyntaxShape, Type, + Value, }; +use polars::prelude::NamedFrom; +use polars::series::Series; #[derive(Clone)] pub struct ToLazyFrame; @@ -25,7 +28,7 @@ impl PluginCommand for ToLazyFrame { Signature::build(self.name()) .named( "schema", - SyntaxShape::Record(vec![]), + SyntaxShape::Any, r#"Polars Schema in format [{name: str}]."#, Some('s'), ) @@ -34,16 +37,28 @@ impl PluginCommand for ToLazyFrame { } fn examples(&self) -> Vec { - vec![Example { - description: "Takes a table and creates a lazyframe", - example: "[[a b];[1 2] [3 4]] | polars into-lazy", - result: None, - }, - Example { - description: "Takes a table, creates a lazyframe, assigns column 'b' type str, displays the schema", - example: "[[a b];[1 2] [3 4]] | polars into-lazy --schema {b: str} | polars schema", - result: Some(Value::test_record(record! {"b" => Value::test_string("str")})), - }, + vec![ + Example { + description: "Takes a table and creates a lazyframe", + example: "[[a b];[1 2] [3 4]] | polars into-lazy", + result: None, + }, + Example { + description: "Takes a table, creates a lazyframe, assigns column 'b' type str, displays the schema", + example: "[[a b];[1 2] [3 4]] | polars into-lazy --schema {b: str} | polars schema", + result: Some(Value::test_record(record! {"b" => Value::test_string("str")})), + }, + Example { + description: "Use a predefined schama", + example: r#"let schema = {a: str, b: str}; [[a b]; [1 "foo"] [2 "bar"]] | polars into-lazy -s $schema"#, + result: Some(NuDataFrame::try_from_series_vec(vec![ + Series::new("a".into(), ["1", "2"]), + Series::new("b".into(), ["foo", "bar"]), + ], Span::test_data()) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, ] } 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 new file mode 100644 index 0000000000..13d90cc135 --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/command/core/to_schema.rs @@ -0,0 +1,74 @@ +use nu_plugin::PluginCommand; +use nu_protocol::{record, Category, Example, ShellError, Signature, Span, Type, Value}; + +use crate::{ + values::{CustomValueSupport, NuSchema}, + PolarsPlugin, +}; + +pub struct ToSchema; + +impl PluginCommand for ToSchema { + type Plugin = PolarsPlugin; + + fn name(&self) -> &str { + "polars into-schema" + } + + fn description(&self) -> &str { + "Convert a value to a polars schema object" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .input_output_type(Type::Any, Type::Custom("schema".into())) + .category(Category::Custom("dataframe".into())) + } + + fn examples(&self) -> Vec { + vec![Example { + 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()), + "b" => Value::string("u8", Span::test_data()), + }, + Span::test_data(), + )), + }] + } + + fn run( + &self, + plugin: &Self::Plugin, + engine: &nu_plugin::EngineInterface, + call: &nu_plugin::EvaluatedCall, + input: nu_protocol::PipelineData, + ) -> Result { + command(plugin, engine, call, input).map_err(nu_protocol::LabeledError::from) + } +} + +fn command( + plugin: &PolarsPlugin, + engine: &nu_plugin::EngineInterface, + call: &nu_plugin::EvaluatedCall, + input: nu_protocol::PipelineData, +) -> Result { + 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) + } +}