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 0cbad49dd0..4c4cabc78a 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/core/open.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/core/open.rs @@ -163,7 +163,7 @@ fn command( }); } - let hive_options = build_hive_options(call)?; + let hive_options = build_hive_options(plugin, call)?; match type_option { Some((ext, blamed)) => match PolarsFileType::from(ext.as_str()) { @@ -388,7 +388,7 @@ fn from_json( })?; let maybe_schema = call .get_flag("schema")? - .map(|schema| NuSchema::try_from(&schema)) + .map(|schema| NuSchema::try_from_value(plugin, &schema)) .transpose()?; let buf_reader = BufReader::new(file); @@ -429,11 +429,7 @@ fn from_ndjson( NonZeroUsize::new(DEFAULT_INFER_SCHEMA) .expect("The default infer-schema should be non zero"), ); - let maybe_schema = call - .get_flag("schema")? - .map(|schema| NuSchema::try_from(&schema)) - .transpose()?; - + let maybe_schema = get_schema(plugin, call)?; if !is_eager { let start_time = std::time::Instant::now(); @@ -511,10 +507,7 @@ fn from_csv( .unwrap_or(DEFAULT_INFER_SCHEMA); let skip_rows: Option = call.get_flag("skip-rows")?; let columns: Option> = call.get_flag("columns")?; - let maybe_schema = call - .get_flag("schema")? - .map(|schema| NuSchema::try_from(&schema)) - .transpose()?; + let maybe_schema = get_schema(plugin, call)?; let truncate_ragged_lines: bool = call.has_flag("truncate-ragged-lines")?; if !is_eager { @@ -627,12 +620,15 @@ fn cloud_not_supported(file_type: PolarsFileType, span: Span) -> ShellError { } } -fn build_hive_options(call: &EvaluatedCall) -> Result { +fn build_hive_options( + plugin: &PolarsPlugin, + call: &EvaluatedCall, +) -> Result { let enabled: Option = call.get_flag("hive-enabled")?; let hive_start_idx: Option = call.get_flag("hive-start-idx")?; let schema: Option = call .get_flag::("hive-schema")? - .map(|schema| NuSchema::try_from(&schema)) + .map(|schema| NuSchema::try_from_value(plugin, &schema)) .transpose()?; let try_parse_dates: bool = call.has_flag("hive-try-parse-dates")?; @@ -643,3 +639,11 @@ fn build_hive_options(call: &EvaluatedCall) -> Result { try_parse_dates, }) } + +fn get_schema(plugin: &PolarsPlugin, call: &EvaluatedCall) -> Result, ShellError> { + let schema: Option = call + .get_flag("schema")? + .map(|schema| NuSchema::try_from_value(plugin, &schema)) + .transpose()?; + Ok(schema) +} 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 f8fd9f4572..e60f33acf2 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 @@ -206,7 +206,7 @@ impl PluginCommand for ToDataFrame { ) -> Result { let maybe_schema = call .get_flag("schema")? - .map(|schema| NuSchema::try_from(&schema)) + .map(|schema| NuSchema::try_from_value(plugin, &schema)) .transpose()?; debug!("schema: {:?}", maybe_schema); 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 4d8bf5bfcf..0347c7f9da 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 @@ -54,7 +54,7 @@ impl PluginCommand for ToLazyFrame { ) -> Result { let maybe_schema = call .get_flag("schema")? - .map(|schema| NuSchema::try_from(&schema)) + .map(|schema| NuSchema::try_from_value(plugin, &schema)) .transpose()?; let df = NuDataFrame::try_from_iter(plugin, input.into_iter(), maybe_schema)?; diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_schema/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_schema/mod.rs index 4011f21572..db0bb8172b 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/nu_schema/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/nu_schema/mod.rs @@ -7,7 +7,7 @@ use nu_protocol::{ShellError, Span, Value}; use polars::prelude::{DataType, Field, Schema, SchemaExt, SchemaRef}; use uuid::Uuid; -use crate::Cacheable; +use crate::{Cacheable, PolarsPlugin}; use super::{str_to_dtype, CustomValueSupport, PolarsPluginObject, PolarsPluginType}; @@ -26,14 +26,6 @@ impl NuSchema { } } -impl TryFrom<&Value> for NuSchema { - type Error = ShellError; - fn try_from(value: &Value) -> Result { - let schema = value_to_schema(value, Span::unknown())?; - Ok(Self::new(Arc::new(schema))) - } -} - impl From for SchemaRef { fn from(val: NuSchema) -> Self { Arc::clone(&val.schema) @@ -86,6 +78,24 @@ impl CustomValueSupport for NuSchema { fn base_value(self, span: Span) -> Result { Ok(fields_to_value(self.schema.iter_fields(), span)) } + + fn try_from_value(plugin: &PolarsPlugin, value: &Value) -> Result { + if let Value::Custom { val, .. } = value { + if let Some(cv) = val.as_any().downcast_ref::() { + Self::try_from_custom_value(plugin, cv) + } else { + Err(ShellError::CantConvert { + to_type: Self::get_type_static().to_string(), + from_type: value.get_type().to_string(), + span: value.span(), + help: None, + }) + } + } else { + let schema = value_to_schema(value, Span::unknown())?; + Ok(Self::new(Arc::new(schema))) + } + } } fn fields_to_value(fields: impl Iterator, span: Span) -> Value {