Bidirectional communication and streams for plugins (#11911)

crates/nu_plugin_stream_example/src/example.rs

@@ -0,0 +1,67 @@
+use nu_plugin::{EvaluatedCall, LabeledError};
+use nu_protocol::{ListStream, PipelineData, RawStream, Value};
+
+pub struct Example;
+
+mod int_or_float;
+use self::int_or_float::IntOrFloat;
+
+impl Example {
+    pub fn seq(
+        &self,
+        call: &EvaluatedCall,
+        _input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let first: i64 = call.req(0)?;
+        let last: i64 = call.req(1)?;
+        let span = call.head;
+        let iter = (first..=last).map(move |number| Value::int(number, span));
+        let list_stream = ListStream::from_stream(iter, None);
+        Ok(PipelineData::ListStream(list_stream, None))
+    }
+
+    pub fn sum(
+        &self,
+        call: &EvaluatedCall,
+        input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let mut acc = IntOrFloat::Int(0);
+        let span = input.span();
+        for value in input {
+            if let Ok(n) = value.as_i64() {
+                acc.add_i64(n);
+            } else if let Ok(n) = value.as_f64() {
+                acc.add_f64(n);
+            } else {
+                return Err(LabeledError {
+                    label: "Stream only accepts ints and floats".into(),
+                    msg: format!("found {}", value.get_type()),
+                    span,
+                });
+            }
+        }
+        Ok(PipelineData::Value(acc.to_value(call.head), None))
+    }
+
+    pub fn collect_external(
+        &self,
+        call: &EvaluatedCall,
+        input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let stream = input.into_iter().map(|value| {
+            value
+                .as_str()
+                .map(|str| str.as_bytes())
+                .or_else(|_| value.as_binary())
+                .map(|bin| bin.to_vec())
+        });
+        Ok(PipelineData::ExternalStream {
+            stdout: Some(RawStream::new(Box::new(stream), None, call.head, None)),
+            stderr: None,
+            exit_code: None,
+            span: call.head,
+            metadata: None,
+            trim_end_newline: false,
+        })
+    }
+}

crates/nu_plugin_stream_example/src/example/int_or_float.rs

@@ -0,0 +1,42 @@
+use nu_protocol::Value;
+
+use nu_protocol::Span;
+
+/// Accumulates numbers into either an int or a float. Changes type to float on the first
+/// float received.
+#[derive(Clone, Copy)]
+pub(crate) enum IntOrFloat {
+    Int(i64),
+    Float(f64),
+}
+
+impl IntOrFloat {
+    pub(crate) fn add_i64(&mut self, n: i64) {
+        match self {
+            IntOrFloat::Int(ref mut v) => {
+                *v += n;
+            }
+            IntOrFloat::Float(ref mut v) => {
+                *v += n as f64;
+            }
+        }
+    }
+
+    pub(crate) fn add_f64(&mut self, n: f64) {
+        match self {
+            IntOrFloat::Int(v) => {
+                *self = IntOrFloat::Float(*v as f64 + n);
+            }
+            IntOrFloat::Float(ref mut v) => {
+                *v += n;
+            }
+        }
+    }
+
+    pub(crate) fn to_value(self, span: Span) -> Value {
+        match self {
+            IntOrFloat::Int(v) => Value::int(v, span),
+            IntOrFloat::Float(v) => Value::float(v, span),
+        }
+    }
+}

crates/nu_plugin_stream_example/src/lib.rs

@@ -0,0 +1,4 @@
+mod example;
+mod nu;
+
+pub use example::Example;

crates/nu_plugin_stream_example/src/main.rs

@@ -0,0 +1,30 @@
+use nu_plugin::{serve_plugin, MsgPackSerializer};
+use nu_plugin_stream_example::Example;
+
+fn main() {
+    // When defining your plugin, you can select the Serializer that could be
+    // used to encode and decode the messages. The available options are
+    // MsgPackSerializer and JsonSerializer. Both are defined in the serializer
+    // folder in nu-plugin.
+    serve_plugin(&mut Example {}, MsgPackSerializer {})
+
+    // Note
+    // When creating plugins in other languages one needs to consider how a plugin
+    // is added and used in nushell.
+    // The steps are:
+    // - The plugin is register. In this stage nushell calls the binary file of
+    //      the plugin sending information using the encoded PluginCall::PluginSignature object.
+    //      Use this encoded data in your plugin to design the logic that will return
+    //      the encoded signatures.
+    //      Nushell is expecting and encoded PluginResponse::PluginSignature with all the
+    //      plugin signatures
+    // - When calling the plugin, nushell sends to the binary file the encoded
+    //      PluginCall::CallInfo which has all the call information, such as the
+    //      values of the arguments, the name of the signature called and the input
+    //      from the pipeline.
+    //      Use this data to design your plugin login and to create the value that
+    //      will be sent to nushell
+    //      Nushell expects an encoded PluginResponse::Value from the plugin
+    // - If an error needs to be sent back to nushell, one can encode PluginResponse::Error.
+    //      This is a labeled error that nushell can format for pretty printing
+}

crates/nu_plugin_stream_example/src/nu/mod.rs

@@ -0,0 +1,86 @@
+use crate::Example;
+use nu_plugin::{EvaluatedCall, LabeledError, StreamingPlugin};
+use nu_protocol::{
+    Category, PipelineData, PluginExample, PluginSignature, Span, SyntaxShape, Type, Value,
+};
+
+impl StreamingPlugin for Example {
+    fn signature(&self) -> Vec<PluginSignature> {
+        let span = Span::unknown();
+        vec![
+            PluginSignature::build("stream_example")
+                .usage("Examples for streaming plugins")
+                .search_terms(vec!["example".into()])
+                .category(Category::Experimental),
+            PluginSignature::build("stream_example seq")
+                .usage("Example stream generator for a list of values")
+                .search_terms(vec!["example".into()])
+                .required("first", SyntaxShape::Int, "first number to generate")
+                .required("last", SyntaxShape::Int, "last number to generate")
+                .input_output_type(Type::Nothing, Type::List(Type::Int.into()))
+                .plugin_examples(vec![PluginExample {
+                    example: "stream_example seq 1 3".into(),
+                    description: "generate a sequence from 1 to 3".into(),
+                    result: Some(Value::list(
+                        vec![
+                            Value::int(1, span),
+                            Value::int(2, span),
+                            Value::int(3, span),
+                        ],
+                        span,
+                    )),
+                }])
+                .category(Category::Experimental),
+            PluginSignature::build("stream_example sum")
+                .usage("Example stream consumer for a list of values")
+                .search_terms(vec!["example".into()])
+                .input_output_types(vec![
+                    (Type::List(Type::Int.into()), Type::Int),
+                    (Type::List(Type::Float.into()), Type::Float),
+                ])
+                .plugin_examples(vec![PluginExample {
+                    example: "seq 1 5 | stream_example sum".into(),
+                    description: "sum values from 1 to 5".into(),
+                    result: Some(Value::int(15, span)),
+                }])
+                .category(Category::Experimental),
+            PluginSignature::build("stream_example collect-external")
+                .usage("Example transformer to raw external stream")
+                .search_terms(vec!["example".into()])
+                .input_output_types(vec![
+                    (Type::List(Type::String.into()), Type::String),
+                    (Type::List(Type::Binary.into()), Type::Binary),
+                ])
+                .plugin_examples(vec![PluginExample {
+                    example: "[a b] | stream_example collect-external".into(),
+                    description: "collect strings into one stream".into(),
+                    result: Some(Value::string("ab", span)),
+                }])
+                .category(Category::Experimental),
+        ]
+    }
+
+    fn run(
+        &mut self,
+        name: &str,
+        _config: &Option<Value>,
+        call: &EvaluatedCall,
+        input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        match name {
+            "stream_example" => Err(LabeledError {
+                label: "No subcommand provided".into(),
+                msg: "add --help here to see usage".into(),
+                span: Some(call.head)
+            }),
+            "stream_example seq" => self.seq(call, input),
+            "stream_example sum" => self.sum(call, input),
+            "stream_example collect-external" => self.collect_external(call, input),
+            _ => Err(LabeledError {
+                label: "Plugin call with wrong name signature".into(),
+                msg: "the signature used to call the plugin does not match any name in the plugin signature vector".into(),
+                span: Some(call.head),
+            }),
+        }
+    }
+}