From b1c4fdfe1dde5d63fdf0ffcb15ec77cd5f036bb8 Mon Sep 17 00:00:00 2001
From: Jack Wright <ayax79@gmail.com>
Date: Wed, 9 Apr 2025 13:18:30 -0700
Subject: [PATCH] Introducing `polars into-schema`

---
 .../src/dataframe/command/core/mod.rs         |  2 +
 .../src/dataframe/command/core/open.rs        |  4 +-
 .../src/dataframe/command/core/to_df.rs       | 15 ++++-
 .../src/dataframe/command/core/to_lazy.rs     | 39 ++++++++----
 .../src/dataframe/command/core/to_schema.rs   | 61 +++++++++++++++++++
 5 files changed, 105 insertions(+), 16 deletions(-)
 create mode 100644 crates/nu_plugin_polars/src/dataframe/command/core/to_schema.rs

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<Box<dyn PluginCommand<Plugin = PolarsPlugin
         Box::new(ToLazyFrame),
         Box::new(ToRepr),
         Box::new(to_dtype::ToDataType),
+        Box::new(to_schema::ToSchema),
     ]
 }
diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/open.rs b/crates/nu_plugin_polars/src/dataframe/command/core/open.rs
index 4c4cabc78a..f01fc76824 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/core/open.rs
+++ b/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,
             )
diff --git a/crates/nu_plugin_polars/src/dataframe/command/core/to_df.rs b/crates/nu_plugin_polars/src/dataframe/command/core/to_df.rs
index 72da002333..9280a22fab 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/core/to_df.rs
+++ b/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()),
+                ),
+            },
         ]
     }
 
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..1ff53b535b 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<Example> {
-        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 {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()),
+                ),
+            },
         ]
     }
 
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..7430008653
--- /dev/null
+++ b/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)
+}