Work on NuDataType

crates/nu_plugin_polars/src/cache/list.rs

@@ -132,7 +132,22 @@ impl PluginCommand for ListDF {
                         "reference_count" => Value::int(value.reference_count as i64, call.head),
                     },
                     call.head,
-                )))
+                ))),
+                PolarsPluginObject::NuDataType(_) => Ok(Some(Value::record(
+                    record! {
+                        "key" => Value::string(key.to_string(), call.head),
+                        "created" => Value::date(value.created, call.head),
+                        "columns" => Value::nothing(call.head),
+                        "rows" => Value::nothing(call.head),
+                        "type" => Value::string("DataType", call.head),
+                        "estimated_size" => Value::nothing(call.head),
+                        "span_contents" =>  Value::string(span_contents, value.span),
+                        "span_start" => Value::int(value.span.start as i64, call.head),
+                        "span_end" => Value::int(value.span.end as i64, call.head),
+                        "reference_count" => Value::int(value.reference_count as i64, call.head),
+                    },
+                    call.head,
+                ))),
             }
         })?;
         let vals = vals.into_iter().flatten().collect();

crates/nu_plugin_polars/src/dataframe/values/mod.rs

@@ -1,12 +1,12 @@
 mod file_type;
 mod nu_dataframe;
+mod nu_dtype;
 mod nu_expression;
 mod nu_lazyframe;
 mod nu_lazygroupby;
 mod nu_schema;
 mod nu_when;
 pub mod utils;
-mod nu_dtype;
 
 use crate::{Cacheable, PolarsPlugin};
 use nu_plugin::EngineInterface;
@@ -18,11 +18,11 @@ use uuid::Uuid;
 
 pub use file_type::PolarsFileType;
 pub use nu_dataframe::{Axis, Column, NuDataFrame, NuDataFrameCustomValue};
+pub use nu_dtype::NuDataType;
 pub use nu_expression::{NuExpression, NuExpressionCustomValue};
 pub use nu_lazyframe::{NuLazyFrame, NuLazyFrameCustomValue};
 pub use nu_lazygroupby::{NuLazyGroupBy, NuLazyGroupByCustomValue};
 pub use nu_schema::{str_to_dtype, NuSchema};
-pub use nu_dtype::NuDataType;
 pub use nu_when::{NuWhen, NuWhenCustomValue, NuWhenType};
 
 #[derive(Debug, Clone)]
@@ -33,6 +33,7 @@ pub enum PolarsPluginType {
     NuLazyGroupBy,
     NuWhen,
     NuPolarsTestData,
+    NuDataType,
 }
 
 impl fmt::Display for PolarsPluginType {
@@ -44,6 +45,7 @@ impl fmt::Display for PolarsPluginType {
             Self::NuLazyGroupBy => write!(f, "NuLazyGroupBy"),
             Self::NuWhen => write!(f, "NuWhen"),
             Self::NuPolarsTestData => write!(f, "NuPolarsTestData"),
+            Self::NuDataType => write!(f, "NuDataType"),
         }
     }
 }
@@ -56,6 +58,7 @@ pub enum PolarsPluginObject {
     NuLazyGroupBy(NuLazyGroupBy),
     NuWhen(NuWhen),
     NuPolarsTestData(Uuid, String),
+    NuDataType(NuDataType),
 }
 
 impl PolarsPluginObject {
@@ -104,6 +107,7 @@ impl PolarsPluginObject {
             Self::NuLazyGroupBy(_) => PolarsPluginType::NuLazyGroupBy,
             Self::NuWhen(_) => PolarsPluginType::NuWhen,
             Self::NuPolarsTestData(_, _) => PolarsPluginType::NuPolarsTestData,
+            Self::NuDataType(_) => PolarsPluginType::NuDataType,
         }
     }
 
@@ -115,6 +119,7 @@ impl PolarsPluginObject {
             PolarsPluginObject::NuLazyGroupBy(lg) => lg.id,
             PolarsPluginObject::NuWhen(w) => w.id,
             PolarsPluginObject::NuPolarsTestData(id, _) => *id,
+            PolarsPluginObject::NuDataType(dt) => dt.id,
         }
     }
 
@@ -128,6 +133,7 @@ impl PolarsPluginObject {
             PolarsPluginObject::NuPolarsTestData(id, s) => {
                 Value::string(format!("{id}:{s}"), Span::test_data())
             }
+            PolarsPluginObject::NuDataType(dt) => dt.into_value(span),
         }
     }
 

crates/nu_plugin_polars/src/dataframe/values/nu_dtype.rs

@@ -1,18 +0,0 @@
-use nu_protocol::{ShellError, Span};
-use polars::prelude::{DataType, UnknownKind};
-
-#[derive(Debug, Clone)]
-pub struct NuDataType {
-    dtype: DataType,
-}
-
-impl NuDataType {
-    pub fn new(dtype: DataType) -> Self {
-        Self { dtype }
-    }
-
-    pub fn to_polars(&self) -> DataType {
-        self.dtype.clone()
-    }
-}
-

crates/nu_plugin_polars/src/dataframe/values/nu_dtype/custom_value.rs

@@ -0,0 +1,65 @@
+use nu_protocol::{CustomValue, ShellError, Span, Value};
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+use crate::values::{CustomValueSupport, PolarsPluginCustomValue};
+
+use super::NuDataType;
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct NuDataTypeCustomValue {
+    pub id: Uuid,
+    #[serde(skip)]
+    pub datatype: Option<NuDataType>,
+}
+
+#[typetag::serde]
+impl CustomValue for NuDataTypeCustomValue {
+    fn clone_value(&self, span: nu_protocol::Span) -> Value {
+        Value::custom(Box::new(self.clone()), span)
+    }
+
+    fn type_name(&self) -> String {
+        "NuDataFrame".into()
+    }
+
+    fn to_base_value(&self, span: Span) -> Result<Value, ShellError> {
+        Ok(Value::string(
+            "NuDataFrameValue: custom_value_to_base_value should've been called",
+            span,
+        ))
+    }
+
+    fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn notify_plugin_on_drop(&self) -> bool {
+        true
+    }
+}
+
+impl PolarsPluginCustomValue for NuDataTypeCustomValue {
+    type PolarsPluginObjectType = NuDataType;
+
+    fn id(&self) -> &Uuid {
+        &self.id
+    }
+
+    fn internal(&self) -> &Option<Self::PolarsPluginObjectType> {
+        &self.datatype
+    }
+
+    fn custom_value_to_base_value(
+        &self,
+        plugin: &crate::PolarsPlugin,
+        _engine: &nu_plugin::EngineInterface,
+    ) -> Result<Value, ShellError> {
+        let dtype = NuDataType::try_from_custom_value(plugin, self)?;
+        dtype.base_value(Span::unknown())
+    }
+}

crates/nu_plugin_polars/src/dataframe/values/nu_dtype/mod.rs

@@ -0,0 +1,108 @@
+mod custom_value;
+
+use custom_value::NuDataTypeCustomValue;
+use nu_protocol::{ShellError, Span, Value};
+use polars::prelude::DataType;
+use uuid::Uuid;
+
+use crate::Cacheable;
+
+use super::{
+    nu_schema::dtype_to_value, str_to_dtype, CustomValueSupport, PolarsPluginObject,
+    PolarsPluginType,
+};
+
+#[derive(Debug, Clone)]
+pub struct NuDataType {
+    pub id: uuid::Uuid,
+    dtype: DataType,
+}
+
+impl NuDataType {
+    pub fn new(dtype: DataType) -> Self {
+        Self {
+            id: uuid::Uuid::new_v4(),
+            dtype,
+        }
+    }
+
+    pub fn new_with_str(dtype: &str, span: Span) -> Result<Self, ShellError> {
+        let dtype = str_to_dtype(dtype, span)?;
+        Ok(Self {
+            id: uuid::Uuid::new_v4(),
+            dtype,
+        })
+    }
+
+    pub fn to_polars(&self) -> DataType {
+        self.dtype.clone()
+    }
+}
+
+impl TryFrom<&Value> for NuDataType {
+    type Error = ShellError;
+
+    fn try_from(value: &Value) -> Result<Self, Self::Error> {
+        match value {
+            Value::String { val, internal_span } => NuDataType::new_with_str(val, *internal_span),
+            _ => Err(ShellError::GenericError {
+                error: format!("Unsupported value: {:?}", value),
+                msg: "".into(),
+                span: Some(value.span()),
+                help: None,
+                inner: vec![],
+            }),
+        }
+    }
+}
+
+impl From<NuDataType> for Value {
+    fn from(nu_dtype: NuDataType) -> Self {
+        Value::String {
+            val: nu_dtype.dtype.to_string(),
+            internal_span: Span::unknown(),
+        }
+    }
+}
+
+impl Cacheable for NuDataType {
+    fn cache_id(&self) -> &Uuid {
+        &self.id
+    }
+
+    fn to_cache_value(&self) -> Result<super::PolarsPluginObject, ShellError> {
+        Ok(PolarsPluginObject::NuDataType(self.clone()))
+    }
+
+    fn from_cache_value(cv: super::PolarsPluginObject) -> Result<Self, ShellError> {
+        match cv {
+            PolarsPluginObject::NuDataType(dt) => Ok(dt),
+            _ => Err(ShellError::GenericError {
+                error: "Cache value is not a dataframe".into(),
+                msg: "".into(),
+                span: None,
+                help: None,
+                inner: vec![],
+            }),
+        }
+    }
+}
+
+impl CustomValueSupport for NuDataType {
+    type CV = NuDataTypeCustomValue;
+
+    fn get_type_static() -> super::PolarsPluginType {
+        PolarsPluginType::NuDataType
+    }
+
+    fn custom_value(self) -> Self::CV {
+        NuDataTypeCustomValue {
+            id: self.id,
+            datatype: Some(self),
+        }
+    }
+
+    fn base_value(self, span: Span) -> Result<Value, ShellError> {
+        Ok(dtype_to_value(&self.dtype, span))
+    }
+}

crates/nu_plugin_polars/src/dataframe/values/nu_schema.rs

@@ -55,7 +55,7 @@ fn fields_to_value(fields: impl Iterator<Item = Field>, span: Span) -> Value {
     Value::record(record, Span::unknown())
 }
 
-fn dtype_to_value(dtype: &DataType, span: Span) -> Value {
+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),