Add support for stderr and exit code (#4647)

crates/nu-command/src/system/complete.rs

@@ -0,0 +1,100 @@
+use nu_protocol::{
+    ast::Call,
+    engine::{Command, EngineState, Stack},
+    Category, Example, IntoPipelineData, PipelineData, ShellError, Signature, Value,
+};
+
+#[derive(Clone)]
+pub struct Complete;
+
+impl Command for Complete {
+    fn name(&self) -> &str {
+        "complete"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("complete").category(Category::System)
+    }
+
+    fn usage(&self) -> &str {
+        "Complete the external piped in, collecting outputs and exit code"
+    }
+
+    fn run(
+        &self,
+        _engine_state: &EngineState,
+        _stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        match input {
+            PipelineData::ExternalStream {
+                stdout,
+                stderr,
+                exit_code,
+                ..
+            } => {
+                let mut cols = vec!["stdout".to_string()];
+                let mut vals = vec![];
+
+                let stdout = stdout.into_bytes()?;
+                if let Ok(st) = String::from_utf8(stdout.item.clone()) {
+                    vals.push(Value::String {
+                        val: st,
+                        span: stdout.span,
+                    })
+                } else {
+                    vals.push(Value::Binary {
+                        val: stdout.item,
+                        span: stdout.span,
+                    })
+                };
+
+                if let Some(stderr) = stderr {
+                    cols.push("stderr".to_string());
+                    let stderr = stderr.into_bytes()?;
+                    if let Ok(st) = String::from_utf8(stderr.item.clone()) {
+                        vals.push(Value::String {
+                            val: st,
+                            span: stderr.span,
+                        })
+                    } else {
+                        vals.push(Value::Binary {
+                            val: stderr.item,
+                            span: stderr.span,
+                        })
+                    };
+                }
+
+                if let Some(exit_code) = exit_code {
+                    let mut v: Vec<_> = exit_code.collect();
+
+                    if let Some(v) = v.pop() {
+                        cols.push("exit_code".to_string());
+                        vals.push(v);
+                    }
+                }
+
+                Ok(Value::Record {
+                    cols,
+                    vals,
+                    span: call.head,
+                }
+                .into_pipeline_data())
+            }
+            _ => Err(ShellError::SpannedLabeledError(
+                "Complete only works with external streams".to_string(),
+                "complete only works on external streams".to_string(),
+                call.head,
+            )),
+        }
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Run the external completion",
+            example: "^external arg1 | complete",
+            result: None,
+        }]
+    }
+}

crates/nu-command/src/system/mod.rs

@@ -1,4 +1,5 @@
 mod benchmark;
+mod complete;
 mod exec;
 mod ps;
 mod run_external;
@@ -6,6 +7,7 @@ mod sys;
 mod which_;
 
 pub use benchmark::Benchmark;
+pub use complete::Complete;
 pub use exec::Exec;
 pub use ps::Ps;
 pub use run_external::{External, ExternalCommand};

crates/nu-command/src/system/run_external.rs

@@ -8,7 +8,7 @@ use std::sync::mpsc;
 use nu_engine::env_to_strings;
 use nu_protocol::engine::{EngineState, Stack};
 use nu_protocol::{ast::Call, engine::Command, ShellError, Signature, SyntaxShape, Value};
-use nu_protocol::{Category, Example, PipelineData, RawStream, Span, Spanned};
+use nu_protocol::{Category, Example, ListStream, PipelineData, RawStream, Span, Spanned};
 
 use itertools::Itertools;
 
@@ -192,14 +192,51 @@ impl ExternalCommand {
                 let redirect_stderr = self.redirect_stderr;
                 let span = self.name.span;
                 let output_ctrlc = ctrlc.clone();
-                let (tx, rx) = mpsc::channel();
+                let (stdout_tx, stdout_rx) = mpsc::channel();
+                let (stderr_tx, stderr_rx) = mpsc::channel();
+                let (exit_code_tx, exit_code_rx) = mpsc::channel();
 
                 std::thread::spawn(move || {
                     // If this external is not the last expression, then its output is piped to a channel
                     // and we create a ListStream that can be consumed
 
                     if redirect_stderr {
-                        let _ = child.stderr.take();
+                        let stderr = child.stderr.take().ok_or_else(|| {
+                            ShellError::ExternalCommand(
+                                "Error taking stderr from external".to_string(),
+                                "Redirects need access to stderr of an external command"
+                                    .to_string(),
+                                span,
+                            )
+                        })?;
+
+                        // Stderr is read using the Buffer reader. It will do so until there is an
+                        // error or there are no more bytes to read
+                        let mut buf_read = BufReader::with_capacity(OUTPUT_BUFFER_SIZE, stderr);
+                        while let Ok(bytes) = buf_read.fill_buf() {
+                            if bytes.is_empty() {
+                                break;
+                            }
+
+                            // The Cow generated from the function represents the conversion
+                            // from bytes to String. If no replacements are required, then the
+                            // borrowed value is a proper UTF-8 string. The Owned option represents
+                            // a string where the values had to be replaced, thus marking it as bytes
+                            let bytes = bytes.to_vec();
+                            let length = bytes.len();
+                            buf_read.consume(length);
+
+                            if let Some(ctrlc) = &ctrlc {
+                                if ctrlc.load(Ordering::SeqCst) {
+                                    break;
+                                }
+                            }
+
+                            match stderr_tx.send(bytes) {
+                                Ok(_) => continue,
+                                Err(_) => break,
+                            }
+                        }
                     }
 
                     if redirect_stdout {
@@ -234,7 +271,7 @@ impl ExternalCommand {
                                 }
                             }
 
-                            match tx.send(bytes) {
+                            match stdout_tx.send(bytes) {
                                 Ok(_) => continue,
                                 Err(_) => break,
                             }
@@ -247,16 +284,42 @@ impl ExternalCommand {
                             err.to_string(),
                             span,
                         )),
-                        Ok(_) => Ok(()),
+                        Ok(x) => {
+                            if let Some(code) = x.code() {
+                                let _ = exit_code_tx.send(Value::Int {
+                                    val: code as i64,
+                                    span: head,
+                                });
+                            } else if x.success() {
+                                let _ = exit_code_tx.send(Value::Int { val: 0, span: head });
+                            } else {
+                                let _ = exit_code_tx.send(Value::Int {
+                                    val: -1,
+                                    span: head,
+                                });
+                            }
+                            Ok(())
+                        }
                     }
                 });
-                let receiver = ChannelReceiver::new(rx);
+                let stdout_receiver = ChannelReceiver::new(stdout_rx);
+                let stderr_receiver = ChannelReceiver::new(stderr_rx);
+                let exit_code_receiver = ValueReceiver::new(exit_code_rx);
 
-                Ok(PipelineData::RawStream(
-                    RawStream::new(Box::new(receiver), output_ctrlc, head),
-                    head,
-                    None,
-                ))
+                Ok(PipelineData::ExternalStream {
+                    stdout: RawStream::new(Box::new(stdout_receiver), output_ctrlc.clone(), head),
+                    stderr: Some(RawStream::new(
+                        Box::new(stderr_receiver),
+                        output_ctrlc.clone(),
+                        head,
+                    )),
+                    exit_code: Some(ListStream::from_stream(
+                        Box::new(exit_code_receiver),
+                        output_ctrlc,
+                    )),
+                    span: head,
+                    metadata: None,
+                })
             }
         }
     }
@@ -452,8 +515,8 @@ fn trim_enclosing_quotes(input: &str) -> String {
     }
 }
 
-// Receiver used for the ListStream
-// It implements iterator so it can be used as a ListStream
+// Receiver used for the RawStream
+// It implements iterator so it can be used as a RawStream
 struct ChannelReceiver {
     rx: mpsc::Receiver<Vec<u8>>,
 }
@@ -474,3 +537,26 @@ impl Iterator for ChannelReceiver {
         }
     }
 }
+
+// Receiver used for the ListStream
+// It implements iterator so it can be used as a ListStream
+struct ValueReceiver {
+    rx: mpsc::Receiver<Value>,
+}
+
+impl ValueReceiver {
+    pub fn new(rx: mpsc::Receiver<Value>) -> Self {
+        Self { rx }
+    }
+}
+
+impl Iterator for ValueReceiver {
+    type Item = Value;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.rx.recv() {
+            Ok(v) => Some(v),
+            Err(_) => None,
+        }
+    }
+}