Add --data flag to metadata and set-metadata to allow persisting me…

…tadata This should help work around issues like nushell#7368 without having to propagate metadata everywhere (which would be a much larger refactor).
KoviRobi · Aug 1, 2023 · 03eb9e8 · 03eb9e8
1 parent 7e3becd
commit 03eb9e8
Show file tree

Hide file tree

Showing 5 changed files with 150 additions and 68 deletions.
diff --git a/crates/nu-command/src/debug/metadata.rs b/crates/nu-command/src/debug/metadata.rs
@@ -2,7 +2,7 @@ use nu_engine::CallExt;
 use nu_protocol::ast::{Call, Expr, Expression};
 use nu_protocol::engine::{Command, EngineState, Stack};
 use nu_protocol::{
-    Category, DataSource, Example, IntoPipelineData, PipelineData, PipelineMetadata, ShellError,
+    Category, Example, IntoPipelineData, PipelineData, ShellError,
     Signature, Span, SyntaxShape, Type, Value,
 };
 
@@ -27,6 +27,11 @@ impl Command for Metadata {
                 SyntaxShape::Any,
                 "the expression you want metadata for",
             )
+            .switch(
+                "data",
+                "also add the data to the output record, under the key `data`",
+                Some('d'),
+            )
             .category(Category::Debug)
     }
 
@@ -39,6 +44,7 @@ impl Command for Metadata {
     ) -> Result<PipelineData, ShellError> {
         let arg = call.positional_nth(0);
         let head = call.head;
+        let include_data = call.has_flag("data");
 
         match arg {
             Some(Expression {
@@ -53,58 +59,43 @@ impl Command for Metadata {
                             ..
                         } => {
                             let origin = stack.get_var_with_origin(*var_id, *span)?;
-
-                            Ok(build_metadata_record(&origin, &input.metadata(), head)
-                                .into_pipeline_data())
+                            Ok(build_metadata_record(
+                                origin.span().ok(),
+                                input,
+                                include_data,
+                                head,
+                            ))
                         }
                         _ => {
                             let val: Value = call.req(engine_state, stack, 0)?;
-                            Ok(build_metadata_record(&val, &input.metadata(), head)
-                                .into_pipeline_data())
+                            Ok(build_metadata_record(
+                                val.span().ok(),
+                                input,
+                                include_data,
+                                head,
+                            ))
                         }
                     }
                 } else {
                     let val: Value = call.req(engine_state, stack, 0)?;
-                    Ok(build_metadata_record(&val, &input.metadata(), head).into_pipeline_data())
+                    Ok(build_metadata_record(
+                        val.span().ok(),
+                        input,
+                        include_data,
+                        head,
+                    ))
                 }
             }
             Some(_) => {
                 let val: Value = call.req(engine_state, stack, 0)?;
-                Ok(build_metadata_record(&val, &input.metadata(), head).into_pipeline_data())
-            }
-            None => {
-                let mut cols = vec![];
-                let mut vals = vec![];
-                if let Some(x) = input.metadata().as_deref() {
-                    match x {
-                        PipelineMetadata {
-                            data_source: DataSource::Ls,
-                        } => {
-                            cols.push("source".into());
-                            vals.push(Value::string("ls", head))
-                        }
-                        PipelineMetadata {
-                            data_source: DataSource::HtmlThemes,
-                        } => {
-                            cols.push("source".into());
-                            vals.push(Value::string("into html --list", head))
-                        }
-                        PipelineMetadata {
-                            data_source: DataSource::Profiling(values),
-                        } => {
-                            cols.push("profiling".into());
-                            vals.push(Value::list(values.clone(), head))
-                        }
-                    }
-                }
-
-                Ok(Value::Record {
-                    cols,
-                    vals,
-                    span: head,
-                }
-                .into_pipeline_data())
+                Ok(build_metadata_record(
+                    val.span().ok(),
+                    input,
+                    include_data,
+                    head,
+                ))
             }
+            None => Ok(build_metadata_record(None, input, include_data, head)),
         }
     }
 
@@ -120,19 +111,31 @@ impl Command for Metadata {
                 example: "ls | metadata",
                 result: None,
             },
+            Example {
+                description: "Get the metadata of the input, along with the data",
+                example: "ls | metadata --data",
+                result: None,
+            },
         ]
     }
 }
 
-fn build_metadata_record(
-    arg: &Value,
-    metadata: &Option<Box<PipelineMetadata>>,
+pub fn build_metadata_record(
+    arg_span: Option<Span>,
+    pipeline: PipelineData,
+    include_data: bool,
     head: Span,
-) -> Value {
+) -> PipelineData {
     let mut cols = vec![];
     let mut vals = vec![];
+    let metadata = pipeline.metadata();
 
-    if let Ok(span) = arg.span() {
+    if include_data {
+        cols.push("data".into());
+        vals.push(pipeline.into_value(head));
+    }
+
+    if let Some(span) = arg_span {
         cols.push("span".into());
         vals.push(Value::Record {
             cols: vec!["start".into(), "end".into()],
@@ -150,34 +153,17 @@ fn build_metadata_record(
         });
     }
 
-    if let Some(x) = metadata.as_deref() {
-        match x {
-            PipelineMetadata {
-                data_source: DataSource::Ls,
-            } => {
-                cols.push("source".into());
-                vals.push(Value::string("ls", head))
-            }
-            PipelineMetadata {
-                data_source: DataSource::HtmlThemes,
-            } => {
-                cols.push("source".into());
-                vals.push(Value::string("into html --list", head))
-            }
-            PipelineMetadata {
-                data_source: DataSource::Profiling(values),
-            } => {
-                cols.push("profiling".into());
-                vals.push(Value::list(values.clone(), head))
-            }
-        }
+    if let Some(x) = &metadata {
+        cols.push("source".into());
+        vals.push(Value::string(format!("{}", x), head))
     }
 
     Value::Record {
         cols,
         vals,
         span: head,
     }
+    .into_pipeline_data()
 }
 
 #[cfg(test)]

diff --git a/crates/nu-command/src/debug/mod.rs b/crates/nu-command/src/debug/mod.rs
@@ -5,6 +5,7 @@ mod inspect;
 mod inspect_table;
 mod metadata;
 mod profile;
+mod set_metadata;
 mod timeit;
 mod view;
 mod view_files;
@@ -18,6 +19,7 @@ pub use inspect::Inspect;
 pub use inspect_table::build_table;
 pub use metadata::Metadata;
 pub use profile::Profile;
+pub use set_metadata::SetMetadata;
 pub use timeit::TimeIt;
 pub use view::View;
 pub use view_files::ViewFiles;

diff --git a/crates/nu-command/src/debug/set_metadata.rs b/crates/nu-command/src/debug/set_metadata.rs
@@ -0,0 +1,75 @@
+use nu_engine::CallExt;
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{
+    Category, Example, PipelineData, PipelineMetadata, Signature, SyntaxShape, Type, Value,
+};
+
+#[derive(Clone)]
+pub struct SetMetadata;
+
+impl Command for SetMetadata {
+    fn name(&self) -> &str {
+        "set-metadata"
+    }
+
+    fn usage(&self) -> &str {
+        "Assigns the metadata from the metadata argument into the stream"
+    }
+
+    fn signature(&self) -> nu_protocol::Signature {
+        Signature::build("set-metadata")
+            .input_output_types(vec![(Type::Any, Type::Any)])
+            .allow_variants_without_examples(true)
+            .required(
+                "metadata_var_name",
+                SyntaxShape::Record(vec![]),
+                "metadata variable name",
+            )
+            .category(Category::Debug)
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        let val: Value = call.req(engine_state, stack, 0)?;
+        let metadata = get_source(val);
+        Ok(input.set_metadata(metadata.map(Box::new)))
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Set the metadata of a variable",
+            example: "let-with-metadata a md = ls; $a | set-metadata $md",
+            result: None,
+        }]
+    }
+}
+
+pub fn get_source(arg: Value) -> Option<PipelineMetadata> {
+    let source_val = match arg {
+        Value::Record { cols, vals, .. } => cols
+            .iter()
+            .zip(vals)
+            .find_map(|(col, val)| (col == "source").then_some(val)),
+        _ => return None,
+    }?;
+
+    source_val.as_string().ok()?.parse().ok()
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(SetMetadata {})
+    }
+}
diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs
@@ -136,6 +136,7 @@ pub fn add_shell_command_context(mut engine_state: EngineState) -> EngineState {
             Explain,
             Inspect,
             Metadata,
+            SetMetadata,
             Profile,
             TimeIt,
             View,

diff --git a/crates/nu-protocol/src/pipeline_data.rs b/crates/nu-protocol/src/pipeline_data.rs
@@ -60,9 +60,27 @@ pub struct PipelineMetadata {
     pub data_source: DataSource,
 }
 
-#[derive(Debug, Clone)]
+impl std::fmt::Display for PipelineMetadata {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.data_source.fmt(f)
+    }
+}
+
+impl std::str::FromStr for PipelineMetadata {
+    type Err = strum::ParseError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(PipelineMetadata {
+            data_source: s.parse()?,
+        })
+    }
+}
+
+#[derive(Debug, Clone, strum_macros::Display, strum_macros::EnumString)]
 pub enum DataSource {
+    #[strum(serialize = "ls")]
     Ls,
+    #[strum(serialize = "to html --list")]
     HtmlThemes,
     Profiling(Vec<Value>),
 }