Simplfied _redirect_output() by raising exception instead of returning bool

4 files changed, 48 insertions, 65 deletions

diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index d6cb0ee1..7eec5d06 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -49,7 +49,7 @@ from . import utils
 from .argparse_custom import CompletionItem, DEFAULT_ARGUMENT_PARSER
 from .clipboard import can_clip, get_paste_buffer, write_to_paste_buffer
 from .decorators import with_argparser
-from .exceptions import Cmd2ArgparseError, Cmd2ShlexError, EmbeddedConsoleExit, EmptyStatement
+from .exceptions import Cmd2ArgparseError, Cmd2ShlexError, EmbeddedConsoleExit, EmptyStatement, RedirectionError
 from .history import History, HistoryItem
 from .parsing import StatementParser, Statement, Macro, MacroArg, shlex_split
 from .rl_utils import rl_type, RlType, rl_get_point, rl_set_prompt, vt100_support, rl_make_safe_prompt, rl_warning
@@ -1634,42 +1634,40 @@ class Cmd(cmd.Cmd):
                         # Start saving command's stdout at this point
                         self.stdout.pause_storage = False
 
-                    redir_error, saved_state = self._redirect_output(statement)
+                    saved_state = self._redirect_output(statement)
                     self._cur_pipe_proc_reader = saved_state.pipe_proc_reader
 
-                # Do not continue if an error occurred while trying to redirect
-                if not redir_error:
-                    # See if we need to update self._redirecting
-                    if not already_redirecting:
-                        self._redirecting = saved_state.redirecting
+                # See if we need to update self._redirecting
+                if not already_redirecting:
+                    self._redirecting = saved_state.redirecting
 
-                    timestart = datetime.datetime.now()
+                timestart = datetime.datetime.now()
 
-                    # precommand hooks
-                    data = plugin.PrecommandData(statement)
-                    for func in self._precmd_hooks:
-                        data = func(data)
-                    statement = data.statement
+                # precommand hooks
+                data = plugin.PrecommandData(statement)
+                for func in self._precmd_hooks:
+                    data = func(data)
+                statement = data.statement
 
-                    # call precmd() for compatibility with cmd.Cmd
-                    statement = self.precmd(statement)
+                # call precmd() for compatibility with cmd.Cmd
+                statement = self.precmd(statement)
 
-                    # go run the command function
-                    stop = self.onecmd(statement, add_to_history=add_to_history)
+                # go run the command function
+                stop = self.onecmd(statement, add_to_history=add_to_history)
 
-                    # postcommand hooks
-                    data = plugin.PostcommandData(stop, statement)
-                    for func in self._postcmd_hooks:
-                        data = func(data)
+                # postcommand hooks
+                data = plugin.PostcommandData(stop, statement)
+                for func in self._postcmd_hooks:
+                    data = func(data)
 
-                    # retrieve the final value of stop, ignoring any statement modification from the hooks
-                    stop = data.stop
+                # retrieve the final value of stop, ignoring any statement modification from the hooks
+                stop = data.stop
 
-                    # call postcmd() for compatibility with cmd.Cmd
-                    stop = self.postcmd(stop, statement)
+                # call postcmd() for compatibility with cmd.Cmd
+                stop = self.postcmd(stop, statement)
 
-                    if self.timing:
-                        self.pfeedback('Elapsed: {}'.format(datetime.datetime.now() - timestart))
+                if self.timing:
+                    self.pfeedback('Elapsed: {}'.format(datetime.datetime.now() - timestart))
             finally:
                 # Get sigint protection while we restore stuff
                 with self.sigint_protection:
@@ -1690,6 +1688,8 @@ class Cmd(cmd.Cmd):
             pass
         except Cmd2ShlexError as ex:
             self.perror("Invalid syntax: {}".format(ex))
+        except RedirectionError as ex:
+            self.perror(ex)
         except Exception as ex:
             self.pexcept(ex)
         finally:
@@ -1902,22 +1902,21 @@ class Cmd(cmd.Cmd):
         # Restore any terminator, suffix, redirection, etc.
         return resolved + statement.post_command
 
-    def _redirect_output(self, statement: Statement) -> Tuple[bool, utils.RedirectionSavedState]:
+    def _redirect_output(self, statement: Statement) -> utils.RedirectionSavedState:
         """Handles output redirection for >, >>, and |.
 
         :param statement: a parsed statement from the user
         :return: A bool telling if an error occurred and a utils.RedirectionSavedState object
+        :raises RedirectionError if an error occurs trying to pipe or redirect
         """
         import io
         import subprocess
 
-        redir_error = False
-
         # Initialize the saved state
         saved_state = utils.RedirectionSavedState(self.stdout, sys.stdout, self._cur_pipe_proc_reader)
 
         if not self.allow_redirection:
-            return redir_error, saved_state
+            return saved_state
 
         if statement.pipe_to:
             # Create a pipe with read and write sides
@@ -1955,10 +1954,9 @@ class Cmd(cmd.Cmd):
 
             # Check if the pipe process already exited
             if proc.returncode is not None:
-                self.perror('Pipe process exited with code {} before command could run'.format(proc.returncode))
                 subproc_stdin.close()
                 new_stdout.close()
-                redir_error = True
+                raise RedirectionError('Pipe process exited with code {} before command could run'.format(proc.returncode))
             else:
                 saved_state.redirecting = True
                 saved_state.pipe_proc_reader = utils.ProcReader(proc, self.stdout, sys.stderr)
@@ -1967,8 +1965,7 @@ class Cmd(cmd.Cmd):
         elif statement.output:
             import tempfile
             if (not statement.output_to) and (not self._can_clip):
-                self.perror("Cannot redirect to paste buffer; missing 'pyperclip' and/or pyperclip dependencies")
-                redir_error = True
+                raise RedirectionError("Cannot redirect to paste buffer; missing 'pyperclip' and/or pyperclip dependencies")
 
             # Redirecting to a file
             elif statement.output_to:
@@ -1984,8 +1981,7 @@ class Cmd(cmd.Cmd):
                     saved_state.redirecting = True
                     sys.stdout = self.stdout = new_stdout
                 except OSError as ex:
-                    self.pexcept('Failed to redirect because - {}'.format(ex))
-                    redir_error = True
+                    raise RedirectionError('Failed to redirect because - {}'.format(ex))
 
             # Redirecting to a paste buffer
             else:
@@ -1997,7 +1993,7 @@ class Cmd(cmd.Cmd):
                     self.stdout.write(get_paste_buffer())
                     self.stdout.flush()
 
-        return redir_error, saved_state
+        return saved_state
 
     def _restore_output(self, statement: Statement, saved_state: utils.RedirectionSavedState) -> None:
         """Handles restoring state after output redirection as well as
diff --git a/cmd2/exceptions.py b/cmd2/exceptions.py
index 15787177..635192e1 100644
--- a/cmd2/exceptions.py
+++ b/cmd2/exceptions.py
@@ -24,3 +24,8 @@ class EmbeddedConsoleExit(SystemExit):
 class EmptyStatement(Exception):
     """Custom exception class for handling behavior when the user just presses <Enter>."""
     pass
+
+
+class RedirectionError(Exception):
+    """Custom exception class for when redirecting or piping output fails"""
+    pass
diff --git a/cmd2/utils.py b/cmd2/utils.py
index 03ede2a3..8ad8fc67 100644
--- a/cmd2/utils.py
+++ b/cmd2/utils.py
@@ -674,7 +674,7 @@ class RedirectionSavedState:
         self.saved_sys_stdout = sys_stdout
         self.saved_pipe_proc_reader = pipe_proc_reader
 
-        # Tells if the command is redirecting
+        # Tells if the command is redirecting or piping
         self.redirecting = False
 
         # If the command created a process to pipe to, then then is its reader
diff --git a/tests/test_cmd2.py b/tests/test_cmd2.py
index b86ddfa6..2afcf701 100755
--- a/tests/test_cmd2.py
+++ b/tests/test_cmd2.py
@@ -482,20 +482,11 @@ def test_output_redirection(base_app):
 def test_output_redirection_to_nonexistent_directory(base_app):
     filename = '~/fakedir/this_does_not_exist.txt'
 
-    # Verify that writing to a file in a non-existent directory doesn't work
-    run_cmd(base_app, 'help > {}'.format(filename))
-    with pytest.raises(FileNotFoundError):
-        with open(filename) as f:
-            content = f.read()
-        verify_help_text(base_app, content)
+    out, err = run_cmd(base_app, 'help > {}'.format(filename))
+    assert 'Failed to redirect' in err[0]
 
-    # Verify that appending to a file also works
-    run_cmd(base_app, 'help history >> {}'.format(filename))
-    with pytest.raises(FileNotFoundError):
-        with open(filename) as f:
-            appended_content = f.read()
-        verify_help_text(base_app, appended_content)
-        assert len(appended_content) > len(content)
+    out, err = run_cmd(base_app, 'help >> {}'.format(filename))
+    assert 'Failed to redirect' in err[0]
 
 def test_output_redirection_to_too_long_filename(base_app):
     filename = '~/sdkfhksdjfhkjdshfkjsdhfkjsdhfkjdshfkjdshfkjshdfkhdsfkjhewfuihewiufhweiufhiweufhiuewhiuewhfiuwehfia' \
@@ -504,20 +495,11 @@ def test_output_redirection_to_too_long_filename(base_app):
                'fheiufhieuwhfewiuhfeiufhiuewfhiuewheiwuhfiuewhfiuewhfeiuwfhewiufhiuewhiuewhfeiuwhfiuwehfuiwehfiuehie' \
                'whfieuwfhieufhiuewhfeiuwfhiuefhueiwhfw'
 
-    # Verify that writing to a file in a non-existent directory doesn't work
-    run_cmd(base_app, 'help > {}'.format(filename))
-    with pytest.raises(OSError):
-        with open(filename) as f:
-            content = f.read()
-        verify_help_text(base_app, content)
+    out, err = run_cmd(base_app, 'help > {}'.format(filename))
+    assert 'Failed to redirect' in err[0]
 
-    # Verify that appending to a file also works
-    run_cmd(base_app, 'help history >> {}'.format(filename))
-    with pytest.raises(OSError):
-        with open(filename) as f:
-            appended_content = f.read()
-        verify_help_text(base_app, content)
-        assert len(appended_content) > len(content)
+    out, err = run_cmd(base_app, 'help >> {}'.format(filename))
+    assert 'Failed to redirect' in err[0]
 
 
 def test_feedback_to_output_true(base_app):