Minor DataType refactor (#15728)

# Description
This is needed for the enum work. The recent polars changes have broken
my enum work, so I am breaking it into smaller pull requests.
This commit is contained in:
Jack Wright 2025-05-12 08:11:17 -07:00 committed by GitHub
parent 0f25641722
commit 1a0986903f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 25 additions and 22 deletions

View File

@ -2,12 +2,12 @@ pub mod custom_value;
use custom_value::NuDataTypeCustomValue;
use nu_protocol::{record, ShellError, Span, Value};
use polars::prelude::{DataType, PlSmallStr, TimeUnit, UnknownKind};
use polars::prelude::{DataType, Field, PlSmallStr, TimeUnit, UnknownKind};
use uuid::Uuid;
use crate::{Cacheable, PolarsPlugin};
use super::{nu_schema::dtype_to_value, CustomValueSupport, PolarsPluginObject, PolarsPluginType};
use super::{CustomValueSupport, PolarsPluginObject, PolarsPluginType};
#[derive(Debug, Clone)]
pub struct NuDataType {
@ -300,6 +300,18 @@ pub fn str_to_dtype(dtype: &str, span: Span) -> Result<DataType, ShellError> {
}
}
pub(crate) fn fields_to_value(fields: impl Iterator<Item = Field>, span: Span) -> Value {
let record = fields
.map(|field| {
let col = field.name().to_string();
let val = dtype_to_value(field.dtype(), span);
(col, val)
})
.collect();
Value::record(record, Span::unknown())
}
fn str_to_time_unit(ts_string: &str, span: Span) -> Result<TimeUnit, ShellError> {
match ts_string {
"ms" => Ok(TimeUnit::Milliseconds),
@ -314,3 +326,10 @@ fn str_to_time_unit(ts_string: &str, span: Span) -> Result<TimeUnit, ShellError>
}),
}
}
pub(crate) fn dtype_to_value(dtype: &DataType, span: Span) -> Value {
match dtype {
DataType::Struct(fields) => fields_to_value(fields.iter().cloned(), span),
_ => Value::string(dtype.to_string().replace('[', "<").replace(']', ">"), span),
}
}

View File

@ -9,7 +9,10 @@ use uuid::Uuid;
use crate::{Cacheable, PolarsPlugin};
use super::{str_to_dtype, CustomValueSupport, NuDataType, PolarsPluginObject, PolarsPluginType};
use super::{
nu_dtype::fields_to_value, str_to_dtype, CustomValueSupport, NuDataType, PolarsPluginObject,
PolarsPluginType,
};
#[derive(Debug, Clone)]
pub struct NuSchema {
@ -98,25 +101,6 @@ impl CustomValueSupport for NuSchema {
}
}
fn fields_to_value(fields: impl Iterator<Item = Field>, span: Span) -> Value {
let record = fields
.map(|field| {
let col = field.name().to_string();
let val = dtype_to_value(field.dtype(), span);
(col, val)
})
.collect();
Value::record(record, Span::unknown())
}
pub fn dtype_to_value(dtype: &DataType, span: Span) -> Value {
match dtype {
DataType::Struct(fields) => fields_to_value(fields.iter().cloned(), span),
_ => Value::string(dtype.to_string().replace('[', "<").replace(']', ">"), span),
}
}
fn value_to_schema(plugin: &PolarsPlugin, value: &Value, span: Span) -> Result<Schema, ShellError> {
let fields = value_to_fields(plugin, value, span)?;
let schema = Schema::from_iter(fields);