text: # Easily capture stdout/stderr of the current process and subprocesses.
  #
  # Author: Peter Odding <peter@peterodding.com>
  # Last Change: March 7, 2020
  # URL: https://capturer.readthedocs.io
  
  """Easily capture stdout/stderr of the current process and subprocesses."""
  
  # Standard library modules.
  import multiprocessing
  import os
  import pty
  import shutil
  import signal
  import sys
  import tempfile
  import time
  
  # External dependencies.
  from humanfriendly.deprecation import define_aliases
  from humanfriendly.text import compact, dedent
  from humanfriendly.terminal import clean_terminal_output
  
  # Semi-standard module versioning.
  __version__ = '3.0'
  
  # Define aliases for backwards compatibility.
  define_aliases(module_name=__name__, interpret_carriage_returns='humanfriendly.terminal.clean_terminal_output')
  
  DEFAULT_TEXT_ENCODING = 'UTF-8'
  """
  The name of the default character encoding used to convert captured output to
  Unicode text (a string).
  """
  
  GRACEFUL_SHUTDOWN_SIGNAL = signal.SIGUSR1
  """
  The number of the UNIX signal used to communicate graceful shutdown requests
  from the main process to the output relay process (an integer). See also
  :func:`~MultiProcessHelper.enable_graceful_shutdown()`.
  """
  
  TERMINATION_DELAY = 0.01
  """
  The number of seconds to wait before terminating the output relay process (a
  floating point number).
  """
  
  PARTIAL_DEFAULT = False
  """Whether partial reads are enabled or disabled by default (a boolean)."""
  
  STDOUT_FD = 1
  """
  The number of the file descriptor that refers to the standard output stream (an
  integer).
  """
  
  STDERR_FD = 2
  """
  The number of the file descriptor that refers to the standard error stream (an
  integer).
  """
  
  
  def enable_old_api():
      """
      Enable backwards compatibility with the old API.
  
      This function is called when the :mod:`capturer` module is imported. It
      modifies the :class:`CaptureOutput` class to install method proxies for
      :func:`~PseudoTerminal.get_handle()`, :func:`~PseudoTerminal.get_bytes()`,
      :func:`~PseudoTerminal.get_lines()`, :func:`~PseudoTerminal.get_text()`,
      :func:`~PseudoTerminal.save_to_handle()` and
      :func:`~PseudoTerminal.save_to_path()`.
      """
      for name in ('get_handle', 'get_bytes', 'get_lines', 'get_text', 'save_to_handle', 'save_to_path'):
          setattr(CaptureOutput, name, create_proxy_method(name))
  
  
  def create_proxy_method(name):
      """
      Create a proxy method for use by :func:`enable_old_api()`.
  
      :param name: The name of the :class:`PseudoTerminal` method to call when
                   the proxy method is called.
      :returns: A proxy method (a callable) to be installed on the
                :class:`CaptureOutput` class.
      """
      # Define the proxy method.
      def proxy_method(self, *args, **kw):
          if not hasattr(self, 'output'):
              raise TypeError(compact("""
                  The old calling interface is only supported when
                  merged=True and start_capture() has been called!
              """))
          real_method = getattr(self.output, name)
          return real_method(*args, **kw)
      # Get the docstring of the real method.
      docstring = getattr(PseudoTerminal, name).__doc__
      # Change the docstring to explain that this concerns a proxy method,
      # but only when Sphinx is active (to avoid wasting time generating a
      # docstring that no one is going to look at).
      if 'sphinx' in sys.modules:
          # Remove the signature from the docstring to make it possible to
          # remove leading indentation from the remainder of the docstring.
          lines = docstring.splitlines()
          signature = lines.pop(0)
          # Recompose the docstring from the signature, the remainder of the
          # original docstring and the note about proxy methods.
          docstring = '\n\n'.join([
              signature,
              dedent('\n'.join(lines)),
              dedent("""
                  .. note:: This method is a proxy for the :func:`~PseudoTerminal.{name}()`
                            method of the :class:`PseudoTerminal` class. It requires
                            `merged` to be :data:`True` and it expects that
                            :func:`start_capture()` has been called. If this is not
                            the case then :exc:`~exceptions.TypeError` is raised.
              """, name=name),
          ])
      # Copy the (possible modified) docstring.
      proxy_method.__doc__ = docstring
      return proxy_method
  
  
  class MultiProcessHelper(object):
  
      """
      Helper to spawn and manipulate child processes using :mod:`multiprocessing`.
  
      This class serves as a base class for :class:`CaptureOutput` and
      :class:`PseudoTerminal` because both classes need the same child process
      handling logic.
      """
  
      def __init__(self):
          """Initialize a :class:`MultiProcessHelper` object."""
          self.processes = []
  
      def start_child(self, target):
          """
          Start a child process using :class:`multiprocessing.Process`.
  
          :param target: The callable to run in the child process. Expected to
                         take a single argument which is a
                         :class:`multiprocessing.Event` to be set when the child
                         process has finished initialization.
          """
          started_event = multiprocessing.Event()
          child_process = multiprocessing.Process(target=target, args=(started_event,))
          self.processes.append(child_process)
          child_process.daemon = True
          child_process.start()
          started_event.wait()
  
      def stop_children(self):
          """
          Gracefully shut down all child processes.
  
          Child processes are expected to call :func:`enable_graceful_shutdown()`
          during initialization.
          """
          while self.processes:
              child_process = self.processes.pop()
              if child_process.is_alive():
                  os.kill(child_process.pid, GRACEFUL_SHUTDOWN_SIGNAL)
              child_process.join()
  
      def wait_for_children(self):
          """Wait for all child processes to terminate."""
          for child_process in self.processes:
              child_process.join()
  
      def enable_graceful_shutdown(self):
          """
          Register a signal handler that converts :data:`GRACEFUL_SHUTDOWN_SIGNAL` to an exception.
  
          Used by :func:`~PseudoTerminal.capture_loop()` to gracefully interrupt
          the blocking :func:`os.read()` call when the capture loop needs to be
          terminated (this is required for coverage collection).
          """
          signal.signal(GRACEFUL_SHUTDOWN_SIGNAL, self.raise_shutdown_request)
  
      def raise_shutdown_request(self, signum, frame):
          """Raise :exc:`ShutdownRequested` when :data:`GRACEFUL_SHUTDOWN_SIGNAL` is received."""
          raise ShutdownRequested
  
  
  class CaptureOutput(MultiProcessHelper):
  
      """Context manager to capture the standard output and error streams."""
  
      def __init__(self, merged=True, encoding=DEFAULT_TEXT_ENCODING,
                   termination_delay=TERMINATION_DELAY, chunk_size=1024,
                   relay=True):
          """
          Initialize a :class:`CaptureOutput` object.
  
          :param merged: Whether to capture and relay the standard output and
                         standard error streams as one stream (a boolean,
                         defaults to :data:`True`). When this is :data:`False`
                         the ``stdout`` and ``stderr`` attributes of the
                         :class:`CaptureOutput` object are
                         :class:`PseudoTerminal` objects that can be used to
                         get at the output captured from each stream separately.
          :param encoding: The name of the character encoding used to decode the
                           captured output (a string, defaults to
                           :data:`DEFAULT_TEXT_ENCODING`).
          :param termination_delay: The number of seconds to wait before
                                    terminating the output relay process (a
                                    floating point number, defaults to
                                    :data:`TERMINATION_DELAY`).
          :param chunk_size: The maximum number of bytes to read from the
                             captured streams on each call to :func:`os.read()`
                             (an integer).
          :param relay: If this is :data:`True` (the default) then captured
                        output is relayed to the terminal or parent process,
                        if it's :data:`False` the captured output is hidden
                        (swallowed).
          """
          # Initialize the superclass.
          super(CaptureOutput, self).__init__()
          # Store constructor arguments.
          self.chunk_size = chunk_size
          self.encoding = encoding
          self.merged = merged
          self.relay = relay
          self.termination_delay = termination_delay
          # Initialize instance variables.
          self.pseudo_terminals = []
          self.streams = []
          # Initialize stdout/stderr stream containers.
          self.stdout_stream = self.initialize_stream(sys.stdout, STDOUT_FD)
          self.stderr_stream = self.initialize_stream(sys.stderr, STDERR_FD)
  
      def initialize_stream(self, file_obj, expected_fd):
          """
          Initialize one or more :class:`Stream` objects to capture a standard stream.
  
          :param file_obj: A file-like object with a ``fileno()`` method.
          :param expected_fd: The expected file descriptor of the file-like object.
          :returns: The :class:`Stream` connected to the file descriptor of the
                    file-like object.
  
          By default this method just initializes a :class:`Stream` object
          connected to the given file-like object and its underlying file
          descriptor (a simple one-liner).
  
          If however the file descriptor of the file-like object doesn't have the
          expected value (``expected_fd``) two :class:`Stream` objects will be
          created instead: One of the stream objects will be connected to the
          file descriptor of the file-like object and the other stream object
          will be connected to the file descriptor that was expected
          (``expected_fd``).
  
          This approach is intended to make sure that "nested" output capturing
          works as expected: Output from the current Python process is captured
          from the file descriptor of the file-like object while output from
          subprocesses is captured from the file descriptor given by
          ``expected_fd`` (because the operating system defines special semantics
          for the file descriptors with the numbers one and two that we can't
          just ignore).
  
          For more details refer to `issue 2 on GitHub
          <https://github.com/xolox/python-capturer/issues/2>`_.
          """
          real_fd = file_obj.fileno()
          stream_obj = Stream(real_fd)
          self.streams.append((expected_fd, stream_obj))
          if real_fd != expected_fd:
              self.streams.append((expected_fd, Stream(expected_fd)))
          return stream_obj
  
      def __enter__(self):
          """Automatically call :func:`start_capture()` when entering a :keyword:`with` block."""
          self.start_capture()
          return self
  
      def __exit__(self, exc_type=None, exc_value=None, traceback=None):
          """Automatically call :func:`finish_capture()` when leaving a :keyword:`with` block."""
          self.finish_capture()
  
      @property
      def is_capturing(self):
          """:data:`True` if output is being captured, :data:`False` otherwise."""
          return any(stream.is_redirected for kind, stream in self.streams)
  
      def start_capture(self):
          """
          Start capturing the standard output and error streams.
  
          :raises: :exc:`~exceptions.TypeError` when output is already being
                   captured.
  
          This method is called automatically when using the capture object as a
          context manager. It's provided under a separate name in case someone
          wants to extend :class:`CaptureOutput` and build their own context
          manager on top of it.
          """
          if self.is_capturing:
              raise TypeError("Output capturing is already enabled!")
          if self.merged:
              # Capture (and most likely relay) stdout/stderr as one stream.
              fd = self.stderr_stream.original_fd if self.relay else None
              self.output = self.allocate_pty(relay_fd=fd)
              for kind, stream in self.streams:
                  self.output.attach(stream)
          else:
              # Capture (and most likely relay) stdout/stderr as separate streams.
              if self.relay:
                  # Start the subprocess to relay output.
                  self.output_queue = multiprocessing.Queue()
                  self.start_child(self.merge_loop)
              else:
                  # Disable relaying of output.
                  self.output_queue = None
              self.stdout = self.allocate_pty(output_queue=self.output_queue, queue_token=STDOUT_FD)
              self.stderr = self.allocate_pty(output_queue=self.output_queue, queue_token=STDERR_FD)
              for kind, stream in self.streams:
                  if kind == STDOUT_FD:
                      self.stdout.attach(stream)
                  elif kind == STDERR_FD:
                      self.stderr.attach(stream)
                  else:
                      raise Exception("Programming error: Unrecognized stream type!")
          # Start capturing and relaying of output (in one or two subprocesses).
          for pseudo_terminal in self.pseudo_terminals:
              pseudo_terminal.start_capture()
  
      def finish_capture(self):
          """
          Stop capturing the standard output and error streams.
  
          This method is called automatically when using the capture object as a
          context manager. It's provided under a separate name in case someone
          wants to extend :class:`CaptureOutput` and build their own context
          manager on top of it.
          """
          for pseudo_terminal in self.pseudo_terminals:
              pseudo_terminal.finish_capture()
          self.wait_for_children()
  
      def allocate_pty(self, relay_fd=None, output_queue=None, queue_token=None):
          """
          Allocate a pseudo terminal.
  
          Internal shortcut for :func:`start_capture()` to allocate multiple
          pseudo terminals without code duplication.
          """
          obj = PseudoTerminal(
              self.encoding, self.termination_delay, self.chunk_size,
              relay_fd=relay_fd, output_queue=output_queue,
              queue_token=queue_token,
          )
          self.pseudo_terminals.append(obj)
          return obj
  
      def merge_loop(self, started_event):
          """
          Merge and relay output in a child process.
  
          This internal method is used when standard output and standard error
          are being captured separately. It's responsible for emitting each
          captured line on the appropriate stream without interleaving text
          within lines.
          """
          buffers = {
              STDOUT_FD: OutputBuffer(self.stdout_stream.original_fd),
              STDERR_FD: OutputBuffer(self.stderr_stream.original_fd),
          }
          started_event.set()
          while buffers:
              captured_from, output = self.output_queue.get()
              if output:
                  buffers[captured_from].add(output)
              else:
                  buffers[captured_from].flush()
                  buffers.pop(captured_from)
  
  
  class OutputBuffer(object):
  
      """
      Helper for :func:`CaptureOutput.merge_loop()`.
  
      Buffers captured output and flushes to the appropriate stream after each
      line break.
      """
  
      def __init__(self, fd):
          """
          Initialize an :class:`OutputBuffer` object.
  
          :param fd: The number of the file descriptor where output should be
                     flushed (an integer).
          """
          self.fd = fd
          self.buffer = b''
  
      def add(self, output):
          """
          Add output to the buffer and flush appropriately.
  
          :param output: The output to add to the buffer (a string).
          """
          self.buffer += output
          if b'\n' in self.buffer:
              before, _, self.buffer = self.buffer.rpartition(b'\n')
              os.write(self.fd, before + b'\n')
  
      def flush(self):
          """Flush any remaining buffered output to the stream."""
          os.write(self.fd, self.buffer)
          self.buffer = b''
  
  
  class PseudoTerminal(MultiProcessHelper):
  
      """
      Helper for :class:`CaptureOutput`.
  
      Manages capturing of output and exposing the captured output.
      """
  
      def __init__(self, encoding, termination_delay, chunk_size, relay_fd, output_queue, queue_token):
          """
          Initialize a :class:`PseudoTerminal` object.
  
          :param encoding: The name of the character encoding used to decode the
                           captured output (a string, defaults to
                           :data:`DEFAULT_TEXT_ENCODING`).
          :param termination_delay: The number of seconds to wait before
                                    terminating the output relay process (a
                                    floating point number, defaults to
                                    :data:`TERMINATION_DELAY`).
          :param chunk_size: The maximum number of bytes to read from the
                             captured stream(s) on each call to :func:`os.read()`
                             (an integer).
          :param relay_fd: The number of the file descriptor where captured
                           output should be relayed to (an integer or
                           :data:`None` if ``output_queue`` and ``queue_token``
                           are given).
          :param output_queue: The multiprocessing queue where captured output
                               chunks should be written to (a
                               :class:`multiprocessing.Queue` object or
                               :data:`None` if ``relay_fd`` is given).
          :param queue_token: A unique identifier added to each output chunk
                              written to the queue (any value or :data:`None` if
                              ``relay_fd`` is given).
          """
          # Initialize the superclass.
          super(PseudoTerminal, self).__init__()
          # Store constructor arguments.
          self.encoding = encoding
          self.termination_delay = termination_delay
          self.chunk_size = chunk_size
          self.relay_fd = relay_fd
          self.output_queue = output_queue
          self.queue_token = queue_token
          # Initialize instance variables.
          self.streams = []
          # Allocate a pseudo terminal so we can fake subprocesses into
          # thinking that they are connected to a real terminal (this will
          # trigger them to use e.g. ANSI escape sequences).
          self.master_fd, self.slave_fd = pty.openpty()
          # Create a temporary file in which we'll store the output received on
          # the master end of the pseudo terminal.
          self.output_fd, output_file = tempfile.mkstemp()
          self.output_handle = open(output_file, 'rb')
          # Unlink the temporary file because we have a readable file descriptor
          # and a writable file descriptor and that's all we need! If this
          # surprises you I suggest you investigate why unlink() was named the
          # way it was in UNIX :-).
          os.unlink(output_file)
  
      def attach(self, stream):
          """
          Attach a stream to the pseudo terminal.
  
          :param stream: A :class:`Stream` object.
          """
          stream.redirect(self.slave_fd)
          self.streams.append(stream)
  
      def start_capture(self):
          """Start the child process(es) responsible for capturing and relaying output."""
          self.start_child(self.capture_loop)
  
      def finish_capture(self):
          """Stop the process of capturing output and destroy the pseudo terminal."""
          time.sleep(self.termination_delay)
          self.stop_children()
          self.close_pseudo_terminal()
          self.restore_streams()
  
      def close_pseudo_terminal(self):
          """Close the pseudo terminal's master/slave file descriptors."""
          for name in ('master_fd', 'slave_fd'):
              fd = getattr(self, name)
              if fd is not None:
                  os.close(fd)
                  setattr(self, name, None)
  
      def restore_streams(self):
          """Restore the stream(s) attached to the pseudo terminal."""
          for stream in self.streams:
              stream.restore()
  
      # The CaptureOutput class contains proxy methods for the get_handle(),
      # get_bytes(), get_lines(), get_text(), save_to_handle() and save_to_path()
      # methods defined below. By default Sphinx generates method signatures of
      # the form f(proxy, *args, **kw) for these proxy methods, with the result
      # that the online documentation is rather confusing. As a workaround I've
      # included explicit method signatures in the first line of each of the
      # docstrings. This works because of the following Sphinx option:
      # http://www.sphinx-doc.org/en/latest/ext/autodoc.html#confval-autodoc_docstring_signature
  
      def get_handle(self, partial=PARTIAL_DEFAULT):
          """get_handle(partial=False)
          Get the captured output as a Python file object.
  
          :param partial: If :data:`True` (*not the default*) the partial output
                          captured so far is returned, otherwise (*so by
                          default*) the relay process is terminated and output
                          capturing is disabled before returning the captured
                          output (the default is intended to protect unsuspecting
                          users against partial reads).
          :returns: The captured output as a Python file object. The file
                    object's current position is reset to zero before this
                    function returns.
  
          This method is useful when you're dealing with arbitrary amounts of
          captured data that you don't want to load into memory just so you can
          save it to a file again. In fact, in that case you might want to take a
          look at :func:`save_to_path()` and/or :func:`save_to_handle()` :-).
  
          .. warning:: Two caveats about the use of this method:
  
                       1. If partial is :data:`True` (not the default) the output
                          can end in a partial line, possibly in the middle of an
                          ANSI escape sequence or a multi byte character.
  
                       2. If you close this file handle you just lost your last
                          chance to get at the captured output! (calling this
                          method again will not give you a new file handle)
          """
          if not partial:
              self.finish_capture()
          self.output_handle.seek(0)
          return self.output_handle
  
      def get_bytes(self, partial=PARTIAL_DEFAULT):
          """get_bytes(partial=False)
          Get the captured output as binary data.
  
          :param partial: Refer to :func:`get_handle()` for details.
          :returns: The captured output as a binary string.
          """
          return self.get_handle(partial).read()
  
      def get_lines(self, interpreted=True, partial=PARTIAL_DEFAULT):
          """get_lines(interpreted=True, partial=False)
          Get the captured output split into lines.
  
          :param interpreted: If :data:`True` (the default) captured output is
                              processed using :func:`.clean_terminal_output()`.
          :param partial: Refer to :func:`get_handle()` for details.
          :returns: The captured output as a list of Unicode strings.
  
          .. warning:: If partial is :data:`True` (not the default) the output
                       can end in a partial line, possibly in the middle of a
                       multi byte character (this may cause decoding errors).
          """
          output = self.get_bytes(partial)
          output = output.decode(self.encoding)
          if interpreted:
              return clean_terminal_output(output)
          else:
              return output.splitlines()
  
      def get_text(self, interpreted=True, partial=PARTIAL_DEFAULT):
          """get_text(interpreted=True, partial=False)
          Get the captured output as a single string.
  
          :param interpreted: If :data:`True` (the default) captured output is
                              processed using :func:`clean_terminal_output()`.
          :param partial: Refer to :func:`get_handle()` for details.
          :returns: The captured output as a Unicode string.
  
          .. warning:: If partial is :data:`True` (not the default) the output
                       can end in a partial line, possibly in the middle of a
                       multi byte character (this may cause decoding errors).
          """
          output = self.get_bytes(partial)
          output = output.decode(self.encoding)
          if interpreted:
              output = u'\n'.join(clean_terminal_output(output))
          return output
  
      def save_to_handle(self, handle, partial=PARTIAL_DEFAULT):
          """save_to_handle(handle, partial=False)
          Save the captured output to an open file handle.
  
          :param handle: A writable file-like object.
          :param partial: Refer to :func:`get_handle()` for details.
          """
          shutil.copyfileobj(self.get_handle(partial), handle)
  
      def save_to_path(self, filename, partial=PARTIAL_DEFAULT):
          """save_to_path(filename, partial=False)
          Save the captured output to a file.
  
          :param filename: The pathname of the file where the captured output
                           should be written to (a string).
          :param partial: Refer to :func:`get_handle()` for details.
          """
          with open(filename, 'wb') as handle:
              self.save_to_handle(handle, partial)
  
      def capture_loop(self, started_event):
          """
          Continuously read from the master end of the pseudo terminal and relay the output.
  
          This function is run in the background by :func:`start_capture()`
          using the :mod:`multiprocessing` module. It's role is to read output
          emitted on the master end of the pseudo terminal and relay this output
          to the real terminal (so the operator can see what's happening in real
          time) as well as a temporary file (for additional processing by the
          caller).
          """
          self.enable_graceful_shutdown()
          started_event.set()
          try:
              while True:
                  # Read from the master end of the pseudo terminal.
                  output = os.read(self.master_fd, self.chunk_size)
                  if output:
                      # Store the output in the temporary file.
                      os.write(self.output_fd, output)
                      # Relay the output to the real terminal?
                      if self.relay_fd is not None:
                          os.write(self.relay_fd, output)
                      # Relay the output to the master process?
                      if self.output_queue is not None:
                          self.output_queue.put((self.queue_token, output))
                  else:
                      # Relinquish our time slice, or in other words: try to be
                      # friendly to other processes when os.read() calls don't
                      # block. Just for the record, all of my experiments have
                      # shown that os.read() on the master file descriptor
                      # returned by pty.openpty() does in fact block.
                      time.sleep(0)
          except ShutdownRequested:
              # Let the master process know that we're shutting down.
              if self.output_queue is not None:
                  self.output_queue.put((self.queue_token, ''))
  
  
  class Stream(object):
  
      """
      Container for standard stream redirection logic.
  
      Used by :class:`CaptureOutput` to temporarily redirect the standard output
      and standard error streams.
  
      .. attribute:: is_redirected
  
         :data:`True` once :func:`redirect()` has been called, :data:`False` when
         :func:`redirect()` hasn't been called yet or :func:`restore()` has since
         been called.
      """
  
      def __init__(self, fd):
          """
          Initialize a :class:`Stream` object.
  
          :param fd: The file descriptor to be redirected (an integer).
          """
          self.fd = fd
          self.original_fd = os.dup(self.fd)
          self.is_redirected = False
  
      def redirect(self, target_fd):
          """
          Redirect output written to the file descriptor to another file descriptor.
  
          :param target_fd: The file descriptor that should receive the output
                            written to the file descriptor given to the
                            :class:`Stream` constructor (an integer).
          :raises: :exc:`~exceptions.TypeError` when the file descriptor is
                   already being redirected.
          """
          if self.is_redirected:
              msg = "File descriptor %s is already being redirected!"
              raise TypeError(msg % self.fd)
          os.dup2(target_fd, self.fd)
          self.is_redirected = True
  
      def restore(self):
          """Stop redirecting output written to the file descriptor."""
          if self.is_redirected:
              os.dup2(self.original_fd, self.fd)
              self.is_redirected = False
  
  
  class ShutdownRequested(Exception):
  
      """
      Raised by :func:`~MultiProcessHelper.raise_shutdown_request()` to signal
      graceful termination requests (in :func:`~PseudoTerminal.capture_loop()`).
      """
  
  
  enable_old_api()
  