From 9abb14b5fd7dfbbfd6740e9108ccb8fb5dd12884 Mon Sep 17 00:00:00 2001
From: JT <547158+jntrnr@users.noreply.github.com>
Date: Thu, 10 Mar 2022 06:29:23 -0500
Subject: [PATCH] ensure exit codes in more cases (#4803)

---
 crates/nu-cli/src/util.rs | 15 ++++++++++++---
 src/eval_file.rs          | 23 ++++++++++++-----------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/crates/nu-cli/src/util.rs b/crates/nu-cli/src/util.rs
index f7d12481da..48ed227d55 100644
--- a/crates/nu-cli/src/util.rs
+++ b/crates/nu-cli/src/util.rs
@@ -19,13 +19,22 @@ pub fn print_pipeline_data(
     let stdout = std::io::stdout();
 
     if let PipelineData::ExternalStream {
-        stdout: Some(stream),
+        stdout: stream,
+        exit_code,
         ..
     } = input
     {
-        for s in stream {
-            let _ = stdout.lock().write_all(s?.as_binary()?);
+        if let Some(stream) = stream {
+            for s in stream {
+                let _ = stdout.lock().write_all(s?.as_binary()?);
+            }
         }
+
+        // Make sure everything has finished
+        if let Some(exit_code) = exit_code {
+            let _: Vec<_> = exit_code.into_iter().collect();
+        }
+
         return Ok(());
     }
 
diff --git a/src/eval_file.rs b/src/eval_file.rs
index e6f77d0eb3..1f879c408d 100644
--- a/src/eval_file.rs
+++ b/src/eval_file.rs
@@ -87,9 +87,14 @@ pub(crate) fn evaluate(
 pub fn print_table_or_error(
     engine_state: &EngineState,
     stack: &mut Stack,
-    pipeline_data: PipelineData,
+    mut pipeline_data: PipelineData,
     config: &Config,
 ) {
+    let exit_code = match &mut pipeline_data {
+        PipelineData::ExternalStream { exit_code, .. } => exit_code.take(),
+        _ => None,
+    };
+
     match engine_state.find_decl("table".as_bytes()) {
         Some(decl_id) => {
             let table = engine_state.get_decl(decl_id).run(
@@ -100,12 +105,7 @@ pub fn print_table_or_error(
             );
 
             match table {
-                Ok(mut table) => {
-                    let exit_code = match &mut table {
-                        PipelineData::ExternalStream { exit_code, .. } => exit_code.take(),
-                        _ => None,
-                    };
-
+                Ok(table) => {
                     for item in table {
                         let stdout = std::io::stdout();
 
@@ -125,10 +125,6 @@ pub fn print_table_or_error(
                             Err(err) => eprintln!("{}", err),
                         };
                     }
-
-                    if let Some(exit_code) = exit_code {
-                        let _: Vec<_> = exit_code.into_iter().collect();
-                    }
                 }
                 Err(error) => {
                     let working_set = StateWorkingSet::new(engine_state);
@@ -161,4 +157,9 @@ pub fn print_table_or_error(
             }
         }
     };
+
+    // Make sure everything has finished
+    if let Some(exit_code) = exit_code {
+        let _: Vec<_> = exit_code.into_iter().collect();
+    }
 }