forked from extern/nushell
b8d37a7541
<!--
if this PR closes one or more issues, you can automatically link the PR
with
them by using one of the [*linking
keywords*](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword),
e.g.
- this PR should close #xxxx
- fixes #xxxx
you can also mention related issues, PRs or discussions!
-->
# Description
<!--
Thank you for improving Nushell. Please, check our [contributing
guide](../CONTRIBUTING.md) and talk to the core team before making major
changes.
Description of your pull request goes here. **Provide examples and/or
screenshots** if your changes affect the user experience.
-->
Fixes https://github.com/nushell/nushell/issues/11716
The problem is in our [record creation
API](0d518bf813/crates/nu-protocol/src/value/record.rs (L33))
which panics if the numbers of columns and values are different. I added
a safe variant that returns a `Result` and used it in the `rotate`
command.
## TODO in another PR:
Go through all `from_raw_cols_vals_unchecked()` (this includes the
`record!` macro which uses the unchecked version) and make sure that
either
a) it is guaranteed the number of cols and vals is the same, or
b) convert the call to `from_raw_cols_vals()`
Reason: Nushell should never panic.
# User-Facing Changes
<!-- List of all changes that impact the user experience here. This
helps us keep track of breaking changes. -->
# Tests + Formatting
<!--
Don't forget to add tests that cover your changes.
Make sure you've run and fixed any issues with these commands:
- `cargo fmt --all -- --check` to check standard code formatting (`cargo
fmt --all` applies these changes)
- `cargo clippy --workspace -- -D warnings -D clippy::unwrap_used` to
check that you're using the standard code style
- `cargo test --workspace` to check that all tests pass (on Windows make
sure to [enable developer
mode](https://learn.microsoft.com/en-us/windows/apps/get-started/developer-mode-features-and-debugging))
- `cargo run -- -c "use std testing; testing run-tests --path
crates/nu-std"` to run the tests for the standard library
> **Note**
> from `nushell` you can also use the `toolkit` as follows
> ```bash
> use toolkit.nu # or use an `env_change` hook to activate it
automatically
> toolkit check pr
> ```
-->
# After Submitting
<!-- If your PR had any user-facing changes, update [the
documentation](https://github.com/nushell/nushell.github.io) after the
PR is merged, if necessary. This will help us keep the docs up to date.
-->
398 lines
14 KiB
Rust
398 lines
14 KiB
Rust
use std::sync::Arc;
|
|
|
|
use nu_protocol::{Record, ShellError, Span, Value};
|
|
use polars::prelude::{DataType, Field, Schema, SchemaRef, TimeUnit};
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct NuSchema {
|
|
pub schema: SchemaRef,
|
|
}
|
|
|
|
impl NuSchema {
|
|
pub fn new(schema: Schema) -> Self {
|
|
Self {
|
|
schema: Arc::new(schema),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl TryFrom<&Value> for NuSchema {
|
|
type Error = ShellError;
|
|
fn try_from(value: &Value) -> Result<Self, Self::Error> {
|
|
let schema = value_to_schema(value, Span::unknown())?;
|
|
Ok(Self::new(schema))
|
|
}
|
|
}
|
|
|
|
impl From<NuSchema> for Value {
|
|
fn from(schema: NuSchema) -> Self {
|
|
fields_to_value(schema.schema.iter_fields(), Span::unknown())
|
|
}
|
|
}
|
|
|
|
impl From<NuSchema> for SchemaRef {
|
|
fn from(val: NuSchema) -> Self {
|
|
Arc::clone(&val.schema)
|
|
}
|
|
}
|
|
|
|
fn fields_to_value(fields: impl Iterator<Item = Field>, span: Span) -> Value {
|
|
let (cols, vals) = fields
|
|
.map(|field| {
|
|
let val = dtype_to_value(field.data_type(), span);
|
|
let col = field.name().to_string();
|
|
(col, val)
|
|
})
|
|
.unzip();
|
|
|
|
let record = Record::from_raw_cols_vals_unchecked(cols, vals);
|
|
Value::record(record, Span::unknown())
|
|
}
|
|
|
|
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(value: &Value, span: Span) -> Result<Schema, ShellError> {
|
|
let fields = value_to_fields(value, span)?;
|
|
let schema = Schema::from_iter(fields);
|
|
Ok(schema)
|
|
}
|
|
|
|
fn value_to_fields(value: &Value, span: Span) -> Result<Vec<Field>, ShellError> {
|
|
let fields = value
|
|
.as_record()?
|
|
.into_iter()
|
|
.map(|(col, val)| match val {
|
|
Value::Record { .. } => {
|
|
let fields = value_to_fields(val, span)?;
|
|
let dtype = DataType::Struct(fields);
|
|
Ok(Field::new(col, dtype))
|
|
}
|
|
_ => {
|
|
let dtype = dtype_str_to_schema(&val.as_string()?, span)?;
|
|
Ok(Field::new(col, dtype))
|
|
}
|
|
})
|
|
.collect::<Result<Vec<Field>, ShellError>>()?;
|
|
Ok(fields)
|
|
}
|
|
|
|
fn dtype_str_to_schema(dtype: &str, span: Span) -> Result<DataType, ShellError> {
|
|
match dtype {
|
|
"bool" => Ok(DataType::Boolean),
|
|
"u8" => Ok(DataType::UInt8),
|
|
"u16" => Ok(DataType::UInt16),
|
|
"u32" => Ok(DataType::UInt32),
|
|
"u64" => Ok(DataType::UInt64),
|
|
"i8" => Ok(DataType::Int8),
|
|
"i16" => Ok(DataType::Int16),
|
|
"i32" => Ok(DataType::Int32),
|
|
"i64" => Ok(DataType::Int64),
|
|
"f32" => Ok(DataType::Float32),
|
|
"f64" => Ok(DataType::Float64),
|
|
"str" => Ok(DataType::String),
|
|
"binary" => Ok(DataType::Binary),
|
|
"date" => Ok(DataType::Date),
|
|
"time" => Ok(DataType::Time),
|
|
"null" => Ok(DataType::Null),
|
|
"unknown" => Ok(DataType::Unknown),
|
|
"object" => Ok(DataType::Object("unknown", None)),
|
|
_ if dtype.starts_with("list") => {
|
|
let dtype = dtype
|
|
.trim_start_matches("list")
|
|
.trim_start_matches('<')
|
|
.trim_end_matches('>')
|
|
.trim();
|
|
let dtype = dtype_str_to_schema(dtype, span)?;
|
|
Ok(DataType::List(Box::new(dtype)))
|
|
}
|
|
_ if dtype.starts_with("datetime") => {
|
|
let dtype = dtype
|
|
.trim_start_matches("datetime")
|
|
.trim_start_matches('<')
|
|
.trim_end_matches('>');
|
|
let mut split = dtype.split(',');
|
|
let next = split
|
|
.next()
|
|
.ok_or_else(|| ShellError::GenericError {
|
|
error: "Invalid polars data type".into(),
|
|
msg: "Missing time unit".into(),
|
|
span: Some(span),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.trim();
|
|
let time_unit = str_to_time_unit(next, span)?;
|
|
let next = split
|
|
.next()
|
|
.ok_or_else(|| ShellError::GenericError {
|
|
error: "Invalid polars data type".into(),
|
|
msg: "Missing time zone".into(),
|
|
span: Some(span),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.trim();
|
|
let timezone = if "*" == next {
|
|
None
|
|
} else {
|
|
Some(next.to_string())
|
|
};
|
|
Ok(DataType::Datetime(time_unit, timezone))
|
|
}
|
|
_ if dtype.starts_with("duration") => {
|
|
let inner = dtype.trim_start_matches("duration<").trim_end_matches('>');
|
|
let next = inner
|
|
.split(',')
|
|
.next()
|
|
.ok_or_else(|| ShellError::GenericError {
|
|
error: "Invalid polars data type".into(),
|
|
msg: "Missing time unit".into(),
|
|
span: Some(span),
|
|
help: None,
|
|
inner: vec![],
|
|
})?
|
|
.trim();
|
|
let time_unit = str_to_time_unit(next, span)?;
|
|
Ok(DataType::Duration(time_unit))
|
|
}
|
|
_ => Err(ShellError::GenericError {
|
|
error: "Invalid polars data type".into(),
|
|
msg: format!("Unknown type: {dtype}"),
|
|
span: Some(span),
|
|
help: None,
|
|
inner: vec![],
|
|
}),
|
|
}
|
|
}
|
|
|
|
fn str_to_time_unit(ts_string: &str, span: Span) -> Result<TimeUnit, ShellError> {
|
|
match ts_string {
|
|
"ms" => Ok(TimeUnit::Milliseconds),
|
|
"us" | "μs" => Ok(TimeUnit::Microseconds),
|
|
"ns" => Ok(TimeUnit::Nanoseconds),
|
|
_ => Err(ShellError::GenericError {
|
|
error: "Invalid polars data type".into(),
|
|
msg: "Invalid time unit".into(),
|
|
span: Some(span),
|
|
help: None,
|
|
inner: vec![],
|
|
}),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {
|
|
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_value_to_schema() {
|
|
let value = Value::Record {
|
|
val: Record::from_raw_cols_vals_unchecked(
|
|
vec!["name".into(), "age".into(), "address".into()],
|
|
vec![
|
|
Value::String {
|
|
val: "str".into(),
|
|
internal_span: Span::test_data(),
|
|
},
|
|
Value::String {
|
|
val: "i32".into(),
|
|
internal_span: Span::test_data(),
|
|
},
|
|
Value::Record {
|
|
val: Record::from_raw_cols_vals_unchecked(
|
|
vec!["street".into(), "city".into()],
|
|
vec![
|
|
Value::String {
|
|
val: "str".into(),
|
|
internal_span: Span::test_data(),
|
|
},
|
|
Value::String {
|
|
val: "str".into(),
|
|
internal_span: Span::test_data(),
|
|
},
|
|
],
|
|
),
|
|
internal_span: Span::test_data(),
|
|
},
|
|
],
|
|
),
|
|
internal_span: Span::test_data(),
|
|
};
|
|
let schema = value_to_schema(&value, Span::unknown()).unwrap();
|
|
let expected = Schema::from_iter(vec![
|
|
Field::new("name", DataType::String),
|
|
Field::new("age", DataType::Int32),
|
|
Field::new(
|
|
"address",
|
|
DataType::Struct(vec![
|
|
Field::new("street", DataType::String),
|
|
Field::new("city", DataType::String),
|
|
]),
|
|
),
|
|
]);
|
|
assert_eq!(schema, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dtype_str_to_schema_simple_types() {
|
|
let dtype = "bool";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Boolean;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "u8";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::UInt8;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "u16";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::UInt16;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "u32";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::UInt32;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "u64";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::UInt64;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "i8";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Int8;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "i16";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Int16;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "i32";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Int32;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "i64";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Int64;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "str";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::String;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "binary";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Binary;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "date";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Date;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "time";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Time;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "null";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Null;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "unknown";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Unknown;
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "object";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Object("unknown", None);
|
|
assert_eq!(schema, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dtype_str_schema_datetime() {
|
|
let dtype = "datetime<ms, *>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Datetime(TimeUnit::Milliseconds, None);
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "datetime<us, *>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Datetime(TimeUnit::Microseconds, None);
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "datetime<μs, *>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Datetime(TimeUnit::Microseconds, None);
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "datetime<ns, *>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Datetime(TimeUnit::Nanoseconds, None);
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "datetime<ms, UTC>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Datetime(TimeUnit::Milliseconds, Some("UTC".into()));
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "invalid";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown());
|
|
assert!(schema.is_err())
|
|
}
|
|
|
|
#[test]
|
|
fn test_dtype_str_schema_duration() {
|
|
let dtype = "duration<ms>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Duration(TimeUnit::Milliseconds);
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "duration<us>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Duration(TimeUnit::Microseconds);
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "duration<μs>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Duration(TimeUnit::Microseconds);
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "duration<ns>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::Duration(TimeUnit::Nanoseconds);
|
|
assert_eq!(schema, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dtype_str_to_schema_list_types() {
|
|
let dtype = "list<i32>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::List(Box::new(DataType::Int32));
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "list<duration<ms>>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::List(Box::new(DataType::Duration(TimeUnit::Milliseconds)));
|
|
assert_eq!(schema, expected);
|
|
|
|
let dtype = "list<datetime<ms, *>>";
|
|
let schema = dtype_str_to_schema(dtype, Span::unknown()).unwrap();
|
|
let expected = DataType::List(Box::new(DataType::Datetime(TimeUnit::Milliseconds, None)));
|
|
assert_eq!(schema, expected);
|
|
}
|
|
}
|