Keeping track of redirection for each command

1 files changed, 82 insertions, 69 deletions

diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 39a2af38..bc6cd769 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -289,6 +289,31 @@ def with_argparser(argparser: argparse.ArgumentParser,
     return arg_decorator
 
 
+class Statekeeper(object):
+    """Class used to save and restore state during load and py commands as well as when redirecting output or pipes."""
+    def __init__(self, obj: Any, attribs: Iterable) -> None:
+        """Use the instance attributes as a generic key-value store to copy instance attributes from outer object.
+
+        :param obj: instance of cmd2.Cmd derived class (your application instance)
+        :param attribs: tuple of strings listing attributes of obj to save a copy of
+        """
+        self.obj = obj
+        self.attribs = attribs
+        if self.obj:
+            self._save()
+
+    def _save(self) -> None:
+        """Create copies of attributes from self.obj inside this Statekeeper instance."""
+        for attrib in self.attribs:
+            setattr(self, attrib, getattr(self.obj, attrib))
+
+    def restore(self) -> None:
+        """Overwrite attributes in self.obj with the saved values stored in this Statekeeper instance."""
+        if self.obj:
+            for attrib in self.attribs:
+                setattr(self.obj, attrib, getattr(self, attrib))
+
+
 class EmbeddedConsoleExit(SystemExit):
     """Custom exception class for use with the py command."""
     pass
@@ -302,6 +327,9 @@ class EmptyStatement(Exception):
 # Contains data about a disabled command which is used to restore its original functions when the command is enabled
 DisabledCommand = namedtuple('DisabledCommand', ['command_function', 'help_function'])
 
+# Used to restore state after redirection ends
+RedirectionSavedState = namedtuple('RedirectionSavedState', ['self_stdout', 'sys_stdout', 'pipe_proc'])
+
 
 class Cmd(cmd.Cmd):
     """An easy but powerful framework for writing line-oriented command interpreters.
@@ -412,10 +440,6 @@ class Cmd(cmd.Cmd):
         # Built-in commands don't make use of this.  It is purely there for user-defined commands and convenience.
         self._last_result = None
 
-        # Used to save state during a redirection
-        self.kept_state = None
-        self.kept_sys = None
-
         # Codes used for exit conditions
         self._STOP_AND_EXIT = True  # cmd convention
 
@@ -1717,9 +1741,17 @@ class Cmd(cmd.Cmd):
                 # we need to run the finalization hooks
                 raise EmptyStatement
 
+            # Keep track of whether or not we were already redirecting before this command
+            already_redirecting = self.redirecting
+
+            # Handle any redirection for this command
+            saved_state = self._redirect_output(statement)
+
+            # See if we need to update self.redirecting
+            if not already_redirecting:
+                self.redirecting = all(val is not None for val in saved_state)
+
             try:
-                if self.allow_redirection:
-                    self._redirect_output(statement)
                 timestart = datetime.datetime.now()
                 if self._in_py:
                     self._last_result = None
@@ -1747,8 +1779,10 @@ class Cmd(cmd.Cmd):
                 if self.timing:
                     self.pfeedback('Elapsed: %s' % str(datetime.datetime.now() - timestart))
             finally:
-                if self.allow_redirection and self.redirecting:
-                    self._restore_output(statement)
+                self._restore_output(statement, saved_state)
+                if not already_redirecting:
+                    self.redirecting = False
+
         except EmptyStatement:
             # don't do anything, but do allow command finalization hooks to run
             pass
@@ -1873,48 +1907,47 @@ class Cmd(cmd.Cmd):
             raise EmptyStatement()
         return statement
 
-    def _redirect_output(self, statement: Statement) -> None:
+    def _redirect_output(self, statement: Statement) -> RedirectionSavedState:
         """Handles output redirection for >, >>, and |.
 
         :param statement: a parsed statement from the user
+        :return: A RedirectionSavedState object. All elements will be None if no redirection was done.
         """
         import io
         import subprocess
 
-        if statement.pipe_to:
-            self.kept_state = Statekeeper(self, ('stdout',))
+        # Default to no redirection
+        ret_val = RedirectionSavedState(None, None, None)
+
+        if not self.allow_redirection:
+            return ret_val
 
+        if statement.pipe_to:
             # Create a pipe with read and write sides
             read_fd, write_fd = os.pipe()
 
             # Open each side of the pipe and set stdout accordingly
-            # noinspection PyTypeChecker
-            self.stdout = io.open(write_fd, 'w')
-            self.redirecting = True
-            # noinspection PyTypeChecker
             subproc_stdin = io.open(read_fd, 'r')
+            new_stdout = io.open(write_fd, 'w')
 
             # We want Popen to raise an exception if it fails to open the process.  Thus we don't set shell to True.
             try:
-                self.pipe_proc = subprocess.Popen(statement.pipe_to, stdin=subproc_stdin)
+                pipe_proc = subprocess.Popen(statement.pipe_to, stdin=subproc_stdin)
+                ret_val = RedirectionSavedState(self_stdout=self.stdout,
+                                                sys_stdout=None,
+                                                pipe_proc=self.pipe_proc)
+                self.stdout = new_stdout
+                self.pipe_proc = pipe_proc
             except Exception as ex:
                 self.perror('Not piping because - {}'.format(ex), traceback_war=False)
-
-                # Restore stdout to what it was and close the pipe
-                self.stdout.close()
                 subproc_stdin.close()
-                self.pipe_proc = None
-                self.kept_state.restore()
-                self.kept_state = None
-                self.redirecting = False
+                new_stdout.close()
 
         elif statement.output:
             import tempfile
             if (not statement.output_to) and (not self.can_clip):
                 raise EnvironmentError("Cannot redirect to paste buffer; install 'pyperclip' and re-run to enable")
-            self.kept_state = Statekeeper(self, ('stdout',))
-            self.kept_sys = Statekeeper(sys, ('stdout',))
-            self.redirecting = True
+
             if statement.output_to:
                 # going to a file
                 mode = 'w'
@@ -1923,24 +1956,34 @@ class Cmd(cmd.Cmd):
                 if statement.output == constants.REDIRECTION_APPEND:
                     mode = 'a'
                 try:
-                    sys.stdout = self.stdout = open(statement.output_to, mode)
+                    new_stdout = open(statement.output_to, mode)
+                    ret_val = RedirectionSavedState(self_stdout=self.stdout,
+                                                    sys_stdout=sys.stdout,
+                                                    pipe_proc=None)
+                    sys.stdout = self.stdout = new_stdout
                 except OSError as ex:
                     self.perror('Not redirecting because - {}'.format(ex), traceback_war=False)
-                    self.redirecting = False
             else:
                 # going to a paste buffer
-                sys.stdout = self.stdout = tempfile.TemporaryFile(mode="w+")
+                new_stdout = tempfile.TemporaryFile(mode="w+")
+                ret_val = RedirectionSavedState(self_stdout=self.stdout,
+                                                sys_stdout=sys.stdout,
+                                                pipe_proc=None)
+                sys.stdout = self.stdout = new_stdout
                 if statement.output == constants.REDIRECTION_APPEND:
                     self.poutput(get_paste_buffer())
 
-    def _restore_output(self, statement: Statement) -> None:
+        return ret_val
+
+    def _restore_output(self, statement: Statement, saved_state: RedirectionSavedState) -> None:
         """Handles restoring state after output redirection as well as
         the actual pipe operation if present.
 
         :param statement: Statement object which contains the parsed input from the user
+        :param saved_state: contains information needed to restore state data
         """
-        # If we have redirected output to a file or the clipboard or piped it to a shell command, then restore state
-        if self.kept_state is not None:
+        # Check if self.stdout was redirected
+        if saved_state.self_stdout is not None:
             # If we redirected output to the clipboard
             if statement.output and not statement.output_to:
                 self.stdout.seek(0)
@@ -1952,21 +1995,16 @@ class Cmd(cmd.Cmd):
             except BrokenPipeError:
                 pass
             finally:
-                # Restore self.stdout
-                self.kept_state.restore()
-                self.kept_state = None
+                self.stdout = saved_state.self_stdout
 
-            # If we were piping output to a shell command, then close the subprocess the shell command was running in
-            if self.pipe_proc is not None:
+            # Check if output was being piped to a process
+            if saved_state.pipe_proc is not None:
                 self.pipe_proc.communicate()
-                self.pipe_proc = None
-
-        # Restore sys.stdout if need be
-        if self.kept_sys is not None:
-            self.kept_sys.restore()
-            self.kept_sys = None
+                self.pipe_proc = saved_state.pipe_proc
 
-        self.redirecting = False
+        # Check if sys.stdout was redirected
+        if saved_state.sys_stdout is not None:
+            sys.stdout = saved_state.sys_stdout
 
     def cmd_func(self, command: str) -> Optional[Callable]:
         """
@@ -3952,28 +3990,3 @@ class Cmd(cmd.Cmd):
         """Register a hook to be called after a command is completed, whether it completes successfully or not."""
         self._validate_cmdfinalization_callable(func)
         self._cmdfinalization_hooks.append(func)
-
-
-class Statekeeper(object):
-    """Class used to save and restore state during load and py commands as well as when redirecting output or pipes."""
-    def __init__(self, obj: Any, attribs: Iterable) -> None:
-        """Use the instance attributes as a generic key-value store to copy instance attributes from outer object.
-
-        :param obj: instance of cmd2.Cmd derived class (your application instance)
-        :param attribs: tuple of strings listing attributes of obj to save a copy of
-        """
-        self.obj = obj
-        self.attribs = attribs
-        if self.obj:
-            self._save()
-
-    def _save(self) -> None:
-        """Create copies of attributes from self.obj inside this Statekeeper instance."""
-        for attrib in self.attribs:
-            setattr(self, attrib, getattr(self.obj, attrib))
-
-    def restore(self) -> None:
-        """Overwrite attributes in self.obj with the saved values stored in this Statekeeper instance."""
-        if self.obj:
-            for attrib in self.attribs:
-                setattr(self.obj, attrib, getattr(self, attrib))