Introducing polars into-schema

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<Box<dyn PluginCommand<Plugin = PolarsPlugin
         Box::new(ToLazyFrame),
         Box::new(ToRepr),
         Box::new(to_dtype::ToDataType),
+        Box::new(to_schema::ToSchema),
     ]
 }

crates/nu_plugin_polars/src/dataframe/command/core/open.rs

@@ -86,7 +86,7 @@ impl PluginCommand for OpenDataFrame {
             )
             .named(
                 "schema",
-                SyntaxShape::Record(vec![]),
+                SyntaxShape::Any,
                 r#"Polars Schema in format [{name: str}]. CSV, JSON, and JSONL files"#,
                 Some('s')
             )
@@ -103,7 +103,7 @@ impl PluginCommand for OpenDataFrame {
             )
             .named(
                 "hive-schema",
-                SyntaxShape::Record(vec![]),
+                SyntaxShape::Any,
                 r#"Hive schema in format [{name: str}]. Parquet and Arrow files"#,
                 None,
             )

crates/nu_plugin_polars/src/dataframe/command/core/to_df.rs

@@ -34,7 +34,7 @@ impl PluginCommand for ToDataFrame {
         Signature::build(self.name())
             .named(
                 "schema",
-                SyntaxShape::Record(vec![]),
+                SyntaxShape::Any,
                 r#"Polars Schema in format [{name: str}]."#,
                 Some('s'),
             )
@@ -203,7 +203,18 @@ impl PluginCommand for ToDataFrame {
                     .expect("simple df for test should not fail")
                     .into_value(Span::test_data()),
                 ),
-            }
+            },
+            Example {
+                description: "Use a predefined schama",
+                example: r#"let schema = {a: str, b: str}; [[a b]; [1 "foo"] [2 "bar"]] | polars into-df -s {a: str, b: str}"#,
+                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()),
+                ),
+            },
         ]
     }
 

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<Example> {
-        vec![Example {
+        vec![
-            description: "Takes a table and creates a lazyframe",
+            Example {
-            example: "[[a b];[1 2] [3 4]] | polars into-lazy",
+                description: "Takes a table and creates a lazyframe",
-            result: None,
+                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 {
-            example: "[[a b];[1 2] [3 4]] | polars into-lazy --schema {b: str} | polars schema",
+                description: "Takes a table, creates a lazyframe, assigns column 'b' type str, displays the schema",
-            result: Some(Value::test_record(record! {"b" => Value::test_string("str")})),
+                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 {a: str, b: str}"#,
+                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()),
+                ),
+            },
         ]
     }
 

crates/nu_plugin_polars/src/dataframe/command/core/to_schema.rs

@@ -0,0 +1,61 @@
+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<Example> {
+        vec![Example {
+            description: "Convert a record into a schema",
+            example: r#"{a: str, b: u8} | polars into-schema"#,
+            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<nu_protocol::PipelineData, nu_protocol::LabeledError> {
+        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<nu_protocol::PipelineData, ShellError> {
+    NuSchema::try_from_pipeline(plugin, input, call.head)?
+        .to_pipeline_data(plugin, engine, call.head)
+}