Merge pull request #655 from python-cmd2/capture_popen

Capture popen

Author: tleonhardt

CHANGELOG.md

@@ -13,6 +13,8 @@
     `argparse.Namespace` object they pass to the `do_*` methods. It is stored in an attribute called `__statement__`.
     This can be useful if a command function needs to know the command line for things like logging.
     * Added a `-t` option to the `load` command for automatically generating a transcript based on a script file
+    * When in a *pyscript*, the stdout and stderr streams of shell commands and processes being piped to are now
+    captured and included in the ``CommandResult`` structure.
 * Potentially breaking changes
     * The following commands now write to stderr instead of stdout when printing an error. This will make catching
     errors easier in pyscript.
@@ -24,6 +26,9 @@
         * Added ``allow_redirection``, ``terminators``, ``multiline_commands``, and ``shortcuts`` as optional arguments 
         to ``cmd.Cmd.__init__()`
         * A few instance attributes were moved inside ``StatementParser`` and properties were created for accessing them
+    * ``self.pipe_proc`` is now called ``self.cur_pipe_proc_reader`` and is a ``ProcReader`` class.
+    * Shell commands and commands being piped to while in a *pyscript* will function as if their output is going
+    to a pipe and not a tty. This was necessary to be able to capture their output.
 
 ## 0.9.11 (March 13, 2019)
 * Bug Fixes

cmd2/cmd2.py

@@ -1,13 +1,14 @@
-#
 # coding=utf-8
 """Shared utility functions"""
 
 import collections
 import os
 import re
+import subprocess
 import sys
+import threading
 import unicodedata
-from typing import Any, Iterable, List, Optional, Union
+from typing import Any, BinaryIO, Iterable, List, Optional, TextIO, Union
 
 from wcwidth import wcswidth
 
@@ -140,7 +141,6 @@ def which(editor: str) -> Optional[str]:
     :param editor: filename of the editor to check, ie 'notepad.exe' or 'vi'
     :return: a full path or None
     """
-    import subprocess
     try:
         editor_path = subprocess.check_output(['which', editor], stderr=subprocess.STDOUT).strip()
         editor_path = editor_path.decode()
@@ -262,6 +262,32 @@ def natural_sort(list_to_sort: Iterable[str]) -> List[str]:
     return sorted(list_to_sort, key=natural_keys)
 
 
+def unquote_redirection_tokens(args: List[str]) -> None:
+    """
+    Unquote redirection tokens in a list of command-line arguments
+    This is used when redirection tokens have to be passed to another command
+    :param args: the command line args
+    """
+    for i, arg in enumerate(args):
+        unquoted_arg = strip_quotes(arg)
+        if unquoted_arg in constants.REDIRECTION_TOKENS:
+            args[i] = unquoted_arg
+
+
+def find_editor() -> str:
+    """Find a reasonable editor to use by default for the system that the cmd2 application is running on."""
+    editor = os.environ.get('EDITOR')
+    if not editor:
+        if sys.platform[:3] == 'win':
+            editor = 'notepad'
+        else:
+            # Favor command-line editors first so we don't leave the terminal to edit
+            for editor in ['vim', 'vi', 'emacs', 'nano', 'pico', 'gedit', 'kate', 'subl', 'geany', 'atom']:
+                if which(editor):
+                    break
+    return editor
+
+
 class StdSim(object):
     """
     Class to simulate behavior of sys.stdout or sys.stderr.
@@ -315,7 +341,14 @@ def readbytes(self) -> bytes:
 
     def clear(self) -> None:
         """Clear the internal contents"""
-        self.buffer.byte_buf = b''
+        self.buffer.byte_buf = bytearray()
+
+    def isatty(self) -> bool:
+        """StdSim only considered an interactive stream if `echo` is True and `inner_stream` is a tty."""
+        if self.echo:
+            return self.inner_stream.isatty()
+        else:
+            return False
 
     def __getattr__(self, item: str):
         if item in self.__dict__:
@@ -329,7 +362,7 @@ class ByteBuf(object):
     Used by StdSim to write binary data and stores the actual bytes written
     """
     def __init__(self, std_sim_instance: StdSim) -> None:
-        self.byte_buf = b''
+        self.byte_buf = bytearray()
         self.std_sim_instance = std_sim_instance
 
     def write(self, b: bytes) -> None:
@@ -342,27 +375,140 @@ def write(self, b: bytes) -> None:
             self.std_sim_instance.inner_stream.buffer.write(b)
 
 
-def unquote_redirection_tokens(args: List[str]) -> None:
+class ProcReader(object):
     """
-    Unquote redirection tokens in a list of command-line arguments
-    This is used when redirection tokens have to be passed to another command
-    :param args: the command line args
+    Used to captured stdout and stderr from a Popen process if any of those were set to subprocess.PIPE.
+    If neither are pipes, then the process will run normally and no output will be captured.
     """
-    for i, arg in enumerate(args):
-        unquoted_arg = strip_quotes(arg)
-        if unquoted_arg in constants.REDIRECTION_TOKENS:
-            args[i] = unquoted_arg
+    def __init__(self, proc: subprocess.Popen, stdout: Union[StdSim, BinaryIO, TextIO],
+                 stderr: Union[StdSim, BinaryIO, TextIO]) -> None:
+        """
+        ProcReader initializer
+        :param proc: the Popen process being read from
+        :param stdout: the stream to write captured stdout
+        :param stderr: the stream to write captured stderr
+        """
+        self._proc = proc
+        self._stdout = stdout
+        self._stderr = stderr
+
+        self._out_thread = threading.Thread(name='out_thread', target=self._reader_thread_func,
+                                            kwargs={'read_stdout': True})
+
+        self._err_thread = threading.Thread(name='out_thread', target=self._reader_thread_func,
+                                            kwargs={'read_stdout': False})
+
+        # Start the reader threads for pipes only
+        if self._proc.stdout is not None:
+            self._out_thread.start()
+        if self._proc.stderr is not None:
+            self._err_thread.start()
+
+    def send_sigint(self) -> None:
+        """Send a SIGINT to the process similar to if <Ctrl>+C were pressed."""
+        import signal
+        if sys.platform.startswith('win'):
+            signal_to_send = signal.CTRL_C_EVENT
+        else:
+            signal_to_send = signal.SIGINT
+        self._proc.send_signal(signal_to_send)
+
+    def terminate(self) -> None:
+        """Terminate the process"""
+        self._proc.terminate()
+
+    def wait(self) -> None:
+        """Wait for the process to finish"""
+        if self._out_thread.is_alive():
+            self._out_thread.join()
+        if self._err_thread.is_alive():
+            self._err_thread.join()
+
+        # Handle case where the process ended before the last read could be done.
+        # This will return None for the streams that weren't pipes.
+        out, err = self._proc.communicate()
+
+        if out:
+            self._write_bytes(self._stdout, out)
+        if err:
+            self._write_bytes(self._stderr, err)
+
+    def _reader_thread_func(self, read_stdout: bool) -> None:
+        """
+        Thread function that reads a stream from the process
+        :param read_stdout: if True, then this thread deals with stdout. Otherwise it deals with stderr.
+        """
+        if read_stdout:
+            read_stream = self._proc.stdout
+            write_stream = self._stdout
+        else:
+            read_stream = self._proc.stderr
+            write_stream = self._stderr
+
+        # The thread should have been started only if this stream was a pipe
+        assert read_stream is not None
+
+        # Run until process completes
+        while self._proc.poll() is None:
+            # noinspection PyUnresolvedReferences
+            available = read_stream.peek()
+            if available:
+                read_stream.read(len(available))
+                self._write_bytes(write_stream, available)
+
+    @staticmethod
+    def _write_bytes(stream: Union[StdSim, BinaryIO, TextIO], to_write: bytes) -> None:
+        """
+        Write bytes to a stream
+        :param stream: the stream being written to
+        :param to_write: the bytes being written
+        """
+        try:
+            if hasattr(stream, 'buffer'):
+                stream.buffer.write(to_write)
+            else:
+                stream.write(to_write)
+        except BrokenPipeError:
+            # This occurs if output is being piped to a process that closed
+            pass
 
 
-def find_editor() -> str:
-    """Find a reasonable editor to use by default for the system that the cmd2 application is running on."""
-    editor = os.environ.get('EDITOR')
-    if not editor:
-        if sys.platform[:3] == 'win':
-            editor = 'notepad'
-        else:
-            # Favor command-line editors first so we don't leave the terminal to edit
-            for editor in ['vim', 'vi', 'emacs', 'nano', 'pico', 'gedit', 'kate', 'subl', 'geany', 'atom']:
-                if which(editor):
-                    break
-    return editor
+class ContextFlag(object):
+    """A context manager which is also used as a boolean flag value within the default sigint handler.
+
+    Its main use is as a flag to prevent the SIGINT handler in cmd2 from raising a KeyboardInterrupt
+    while a critical code section has set the flag to True. Because signal handling is always done on the
+    main thread, this class is not thread-safe since there is no need.
+    """
+    def __init__(self) -> None:
+        # When this flag has a positive value, it is considered set.
+        # When it is 0, it is not set. It should never go below 0.
+        self.__count = 0
+
+    def __bool__(self) -> bool:
+        return self.__count > 0
+
+    def __enter__(self) -> None:
+        self.__count += 1
+
+    def __exit__(self, *args) -> None:
+        self.__count -= 1
+        if self.__count < 0:
+            raise ValueError("count has gone below 0")
+
+
+class RedirectionSavedState(object):
+    """Created by each command to store information about their redirection."""
+
+    def __init__(self, self_stdout: Union[StdSim, BinaryIO, TextIO], sys_stdout: Union[StdSim, BinaryIO, TextIO],
+                 pipe_proc_reader: Optional[ProcReader]) -> None:
+        # Used to restore values after the command ends
+        self.saved_self_stdout = self_stdout
+        self.saved_sys_stdout = sys_stdout
+        self.saved_pipe_proc_reader = pipe_proc_reader
+
+        # Tells if the command is redirecting
+        self.redirecting = False
+
+        # If the command created a process to pipe to, then then is its reader
+        self.pipe_proc_reader = None

cmd2/utils.py

@@ -5,6 +5,7 @@
 Copyright 2016 Federico Ceratto <[email protected]>
 Released under MIT license, see LICENSE file
 """
+import sys
 from typing import Optional
 from unittest import mock
 
@@ -13,6 +14,13 @@
 import cmd2
 from cmd2.utils import StdSim
 
+# Python 3.4 require contextlib2 for temporarily redirecting stderr and stdout
+if sys.version_info < (3, 5):
+    # noinspection PyUnresolvedReferences
+    from contextlib2 import redirect_stdout, redirect_stderr
+else:
+    from contextlib import redirect_stdout, redirect_stderr
+
 # Prefer statically linked gnureadline if available (for macOS compatibility due to issues with libedit)
 try:
     import gnureadline as readline
@@ -126,19 +134,33 @@ def normalize(block):
 
 
 def run_cmd(app, cmd):
-    """ Clear StdSim buffer, run the command, extract the buffer contents, """
-    app.stdout.clear()
-    app.onecmd_plus_hooks(cmd)
-    out = app.stdout.getvalue()
-    app.stdout.clear()
-    return normalize(out)
+    """ Clear out and err StdSim buffers, run the command, and return out and err """
+    saved_sysout = sys.stdout
+    sys.stdout = app.stdout
+
+    # This will be used to capture app.stdout and sys.stdout
+    copy_cmd_stdout = StdSim(app.stdout)
+
+    # This will be used to capture sys.stderr
+    copy_stderr = StdSim(sys.stderr)
+
+    try:
+        app.stdout = copy_cmd_stdout
+        with redirect_stdout(copy_cmd_stdout):
+            with redirect_stderr(copy_stderr):
+                app.onecmd_plus_hooks(cmd)
+    finally:
+        app.stdout = copy_cmd_stdout.inner_stream
+        sys.stdout = saved_sysout
+
+    out = copy_cmd_stdout.getvalue()
+    err = copy_stderr.getvalue()
+    return normalize(out), normalize(err)
 
 
 @fixture
 def base_app():
-    c = cmd2.Cmd()
-    c.stdout = StdSim(c.stdout)
-    return c
+    return cmd2.Cmd()
 
 
 def complete_tester(text: str, line: str, begidx: int, endidx: int, app) -> Optional[str]: