Allow saving CustomValues (nushell#16692)

Author: cptpiepmatz

crates/nu-command/src/filesystem/save.rs

@@ -8,6 +8,7 @@ use nu_protocol::{
     byte_stream::copy_with_signals, process::ChildPipe, shell_error::io::IoError,
 };
 use std::{
+    borrow::Cow,
     fs::File,
     io::{self, BufRead, BufReader, Read, Write},
     path::{Path, PathBuf},
@@ -224,8 +225,28 @@ impl Command for Save {
                     )?;
                 }
 
-                let bytes =
-                    input_to_bytes(input, Path::new(&path.item), raw, engine_state, stack, span)?;
+                // Try to convert the input pipeline into another type if we know the extension
+                let ext = extract_extension(&input, &path.item, raw);
+                let converted = match ext {
+                    None => input,
+                    Some(ext) => convert_to_extension(engine_state, &ext, stack, input, span)?,
+                };
+
+                // Save custom value however they implement saving
+                if let PipelineData::Value(Value::Custom { val, internal_span }, ..) = converted {
+                    return val
+                        .save(
+                            Spanned {
+                                item: &path.item,
+                                span: path.span,
+                            },
+                            internal_span,
+                            span,
+                        )
+                        .map(|()| PipelineData::empty());
+                }
+
+                let bytes = value_to_bytes(converted.into_value(span)?)?;
 
                 // Only open file after successful conversion
                 let (mut file, _) =
@@ -321,34 +342,14 @@ fn check_saving_to_source_file(
     Ok(())
 }
 
-/// Convert [`PipelineData`] bytes to write in file, possibly converting
-/// to format of output file
-fn input_to_bytes(
-    input: PipelineData,
-    path: &Path,
-    raw: bool,
-    engine_state: &EngineState,
-    stack: &mut Stack,
-    span: Span,
-) -> Result<Vec<u8>, ShellError> {
-    let ext = if raw {
-        None
-    } else if let PipelineData::ByteStream(..) = input {
-        None
-    } else if let PipelineData::Value(Value::String { .. }, ..) = input {
-        None
-    } else {
-        path.extension()
-            .map(|name| name.to_string_lossy().to_string())
-    };
-
-    let input = if let Some(ext) = ext {
-        convert_to_extension(engine_state, &ext, stack, input, span)?
-    } else {
-        input
-    };
-
-    value_to_bytes(input.into_value(span)?)
+/// Extract extension for conversion.
+fn extract_extension<'e>(input: &PipelineData, path: &'e Path, raw: bool) -> Option<Cow<'e, str>> {
+    match (raw, input) {
+        (true, _)
+        | (_, PipelineData::ByteStream(..))
+        | (_, PipelineData::Value(Value::String { .. }, ..)) => None,
+        _ => path.extension().map(|name| name.to_string_lossy()),
+    }
 }
 
 /// Convert given data into content of file of specified extension if

crates/nu-plugin-engine/src/interface/mod.rs

@@ -17,6 +17,7 @@ use nu_protocol::{
 use nu_utils::SharedCow;
 use std::{
     collections::{BTreeMap, btree_map},
+    path::Path,
     sync::{Arc, OnceLock, mpsc},
 };
 
@@ -1059,6 +1060,32 @@ impl PluginInterface {
         self.custom_value_op_expecting_value(left, CustomValueOp::Operation(operator, right))
     }
 
+    /// Invoke saving operation on a custom value.
+    pub fn custom_value_save(
+        &self,
+        value: Spanned<PluginCustomValueWithSource>,
+        path: Spanned<&Path>,
+        save_call_span: Span,
+    ) -> Result<(), ShellError> {
+        // Check that the value came from the right source
+        value.item.verify_source(value.span, &self.state.source)?;
+
+        let call = PluginCall::CustomValueOp(
+            value.map(|cv| cv.without_source()),
+            CustomValueOp::Save {
+                path: path.map(ToOwned::to_owned),
+                save_call_span,
+            },
+        );
+        match self.plugin_call(call, None)? {
+            PluginCallResponse::Ok => Ok(()),
+            PluginCallResponse::Error(err) => Err(err.into()),
+            _ => Err(ShellError::PluginFailedToDecode {
+                msg: "Received unexpected response to custom value save() call".into(),
+            }),
+        }
+    }
+
     /// Notify the plugin about a dropped custom value.
     pub fn custom_value_dropped(&self, value: PluginCustomValue) -> Result<(), ShellError> {
         // Make sure we don't block here. This can happen on the receiver thread, which would cause a deadlock. We should not try to receive the response - just let it be discarded.
@@ -1253,6 +1280,7 @@ impl CurrentCallState {
                     CustomValueOp::FollowPathString { .. } => Ok(()),
                     CustomValueOp::PartialCmp(value) => self.prepare_value(value, source),
                     CustomValueOp::Operation(_, value) => self.prepare_value(value, source),
+                    CustomValueOp::Save { .. } => Ok(()),
                     CustomValueOp::Dropped => Ok(()),
                 }
             }

crates/nu-plugin-engine/src/plugin_custom_value_with_source/mod.rs

@@ -1,4 +1,4 @@
-use std::{cmp::Ordering, sync::Arc};
+use std::{cmp::Ordering, path::Path, sync::Arc};
 
 use nu_plugin_core::util::with_custom_values_in;
 use nu_plugin_protocol::PluginCustomValue;
@@ -235,6 +235,16 @@ impl CustomValue for PluginCustomValueWithSource {
             )
     }
 
+    fn save(
+        &self,
+        path: Spanned<&Path>,
+        value_span: Span,
+        save_span: Span,
+    ) -> Result<(), ShellError> {
+        self.get_plugin(Some(value_span), "save")?
+            .custom_value_save(self.clone().into_spanned(value_span), path, save_span)
+    }
+
     fn as_any(&self) -> &dyn std::any::Any {
         self
     }

crates/nu-plugin-protocol/src/lib.rs

@@ -28,7 +28,7 @@ use nu_protocol::{
 };
 use nu_utils::SharedCow;
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
+use std::{collections::HashMap, path::PathBuf};
 
 pub use evaluated_call::EvaluatedCall;
 pub use plugin_custom_value::PluginCustomValue;
@@ -212,6 +212,11 @@ pub enum CustomValueOp {
     PartialCmp(Value),
     /// [`operation()`](nu_protocol::CustomValue::operation)
     Operation(Spanned<Operator>, Value),
+    /// [`save()`](nu_protocol::CustomValue::save)
+    Save {
+        path: Spanned<PathBuf>,
+        save_call_span: Span,
+    },
     /// Notify that the custom value has been dropped, if
     /// [`notify_plugin_on_drop()`](nu_protocol::CustomValue::notify_plugin_on_drop) is true
     Dropped,
@@ -226,6 +231,7 @@ impl CustomValueOp {
             CustomValueOp::FollowPathString { .. } => "follow_path_string",
             CustomValueOp::PartialCmp(_) => "partial_cmp",
             CustomValueOp::Operation(_, _) => "operation",
+            CustomValueOp::Save { .. } => "save",
             CustomValueOp::Dropped => "dropped",
         }
     }
@@ -359,6 +365,7 @@ pub enum StreamMessage {
 /// Response to a [`PluginCall`]. The type parameter determines the output type for pipeline data.
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub enum PluginCallResponse<D> {
+    Ok,
     Error(LabeledError),
     Metadata(PluginMetadata),
     Signature(Vec<PluginSignature>),
@@ -374,6 +381,7 @@ impl<D> PluginCallResponse<D> {
         f: impl FnOnce(D) -> Result<T, ShellError>,
     ) -> Result<PluginCallResponse<T>, ShellError> {
         Ok(match self {
+            PluginCallResponse::Ok => PluginCallResponse::Ok,
             PluginCallResponse::Error(err) => PluginCallResponse::Error(err),
             PluginCallResponse::Metadata(meta) => PluginCallResponse::Metadata(meta),
             PluginCallResponse::Signature(sigs) => PluginCallResponse::Signature(sigs),

crates/nu-plugin-protocol/src/plugin_custom_value/mod.rs

@@ -1,6 +1,6 @@
-use std::cmp::Ordering;
+use std::{cmp::Ordering, path::Path};
 
-use nu_protocol::{CustomValue, ShellError, Span, Value, ast::Operator, casing::Casing};
+use nu_protocol::{CustomValue, ShellError, Span, Spanned, Value, ast::Operator, casing::Casing};
 use nu_utils::SharedCow;
 
 use serde::{Deserialize, Serialize};
@@ -93,6 +93,15 @@ impl CustomValue for PluginCustomValue {
         panic!("operation() not available on plugin custom value without source");
     }
 
+    fn save(
+        &self,
+        _path: Spanned<&Path>,
+        _value_span: Span,
+        _save_span: Span,
+    ) -> Result<(), ShellError> {
+        panic!("save() not available on plugin custom value without source");
+    }
+
     fn as_any(&self) -> &dyn std::any::Any {
         self
     }

crates/nu-plugin/src/plugin/interface/mod.rs

@@ -411,6 +411,19 @@ impl EngineInterface {
         })
     }
 
+    /// Write an OK call response or an error.
+    pub(crate) fn write_ok(
+        &self,
+        result: Result<(), impl Into<LabeledError>>,
+    ) -> Result<(), ShellError> {
+        let response = match result {
+            Ok(()) => PluginCallResponse::Ok,
+            Err(err) => PluginCallResponse::Error(err.into()),
+        };
+        self.write(PluginOutput::CallResponse(self.context()?, response))?;
+        self.flush()
+    }
+
     /// Write a call response of either [`PipelineData`] or an error. Returns the stream writer
     /// to finish writing the stream
     pub(crate) fn write_response(

crates/nu-plugin/src/plugin/mod.rs

@@ -5,6 +5,7 @@ use std::{
     ffi::OsString,
     ops::Deref,
     panic::AssertUnwindSafe,
+    path::Path,
     sync::mpsc::{self, TrySendError},
     thread,
 };
@@ -16,8 +17,8 @@ use nu_plugin_core::{
 };
 use nu_plugin_protocol::{CallInfo, CustomValueOp, PluginCustomValue, PluginInput, PluginOutput};
 use nu_protocol::{
-    CustomValue, IntoSpanned, LabeledError, PipelineData, PluginMetadata, ShellError, Spanned,
-    Value, ast::Operator, casing::Casing,
+    CustomValue, IntoSpanned, LabeledError, PipelineData, PluginMetadata, ShellError, Span,
+    Spanned, Value, ast::Operator, casing::Casing,
 };
 use thiserror::Error;
 
@@ -214,6 +215,24 @@ pub trait Plugin: Sync {
             .map_err(LabeledError::from)
     }
 
+    /// Implement saving logic for a custom value.
+    ///
+    /// The default implementation of this method just calls [`CustomValue::save`], but
+    /// the method can be implemented differently if accessing plugin state is desirable.
+    fn custom_value_save(
+        &self,
+        engine: &EngineInterface,
+        value: Spanned<Box<dyn CustomValue>>,
+        path: Spanned<&Path>,
+        save_call_span: Span,
+    ) -> Result<(), LabeledError> {
+        let _ = engine;
+        value
+            .item
+            .save(path, value.span, save_call_span)
+            .map_err(LabeledError::from)
+    }
+
     /// Handle a notification that all copies of a custom value within the engine have been dropped.
     ///
     /// This notification is only sent if [`CustomValue::notify_plugin_on_drop`] was true. Unlike
@@ -652,6 +671,17 @@ fn custom_value_op(
                 .write_response(result)
                 .and_then(|writer| writer.write())
         }
+        CustomValueOp::Save {
+            path,
+            save_call_span,
+        } => {
+            let path = Spanned {
+                item: path.item.as_path(),
+                span: path.span,
+            };
+            let result = plugin.custom_value_save(engine, local_value, path, save_call_span);
+            engine.write_ok(result)
+        }
         CustomValueOp::Dropped => {
             let result = plugin
                 .custom_value_dropped(engine, local_value.item)

crates/nu-protocol/src/value/custom_value.rs

@@ -1,6 +1,6 @@
-use std::{cmp::Ordering, fmt};
+use std::{cmp::Ordering, fmt, path::Path};
 
-use crate::{ShellError, Span, Type, Value, ast::Operator, casing::Casing};
+use crate::{ShellError, Span, Spanned, Type, Value, ast::Operator, casing::Casing};
 
 /// Trait definition for a custom [`Value`](crate::Value) type
 #[typetag::serde(tag = "type")]
@@ -110,6 +110,35 @@ pub trait CustomValue: fmt::Debug + Send + Sync {
         })
     }
 
+    /// Save custom value to disk.
+    ///
+    /// This method is used in `save` to save a custom value to disk.
+    /// This is done before opening any file, so saving can be handled differently.
+    ///
+    /// The default impl just returns an error.
+    fn save(
+        &self,
+        path: Spanned<&Path>,
+        value_span: Span,
+        save_span: Span,
+    ) -> Result<(), ShellError> {
+        let _ = path;
+        Err(ShellError::GenericError {
+            error: "Cannot save custom value".into(),
+            msg: format!("Saving custom value {} failed", self.type_name()),
+            span: Some(save_span),
+            help: None,
+            inner: vec![
+                ShellError::GenericError {
+                error: "Custom value does not implement `save`".into(),
+                msg: format!("{} doesn't implement saving to disk", self.type_name()),
+                span: Some(value_span),
+                help: Some("Check the plugin's documentation for this value type. It might use a different way to save.".into()),
+                inner: vec![],
+            }],
+        })
+    }
+
     /// For custom values in plugins: return `true` here if you would like to be notified when all
     /// copies of this custom value are dropped in the engine.
     ///

crates/nu_plugin_custom_values/src/second_custom_value.rs

@@ -1,7 +1,7 @@
 #![allow(clippy::result_large_err)]
-use nu_protocol::{CustomValue, ShellError, Span, Value};
+use nu_protocol::{CustomValue, ShellError, Span, Spanned, Value, shell_error::io::IoError};
 use serde::{Deserialize, Serialize};
-use std::cmp::Ordering;
+use std::{cmp::Ordering, path::Path};
 
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
 pub struct SecondCustomValue {
@@ -78,4 +78,15 @@ impl CustomValue for SecondCustomValue {
     fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
         self
     }
+
+    fn save(&self, path: Spanned<&Path>, _: Span, save_span: Span) -> Result<(), ShellError> {
+        std::fs::write(path.item, &self.something).map_err(|err| {
+            ShellError::Io(IoError::new_with_additional_context(
+                err,
+                save_span,
+                path.item.to_owned(),
+                format!("Could not save {}", self.type_name()),
+            ))
+        })
+    }
 }