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 custom_value::NuDataTypeCustomValue;
use nu_protocol::{record, ShellError, Span, Value}; 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 uuid::Uuid;
use crate::{Cacheable, PolarsPlugin}; use crate::{Cacheable, PolarsPlugin};
use super::{nu_schema::dtype_to_value, CustomValueSupport, PolarsPluginObject, PolarsPluginType}; use super::{CustomValueSupport, PolarsPluginObject, PolarsPluginType};
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct NuDataType { 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> { fn str_to_time_unit(ts_string: &str, span: Span) -> Result<TimeUnit, ShellError> {
match ts_string { match ts_string {
"ms" => Ok(TimeUnit::Milliseconds), "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 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)] #[derive(Debug, Clone)]
pub struct NuSchema { 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> { fn value_to_schema(plugin: &PolarsPlugin, value: &Value, span: Span) -> Result<Schema, ShellError> {
let fields = value_to_fields(plugin, value, span)?; let fields = value_to_fields(plugin, value, span)?;
let schema = Schema::from_iter(fields); let schema = Schema::from_iter(fields);