# 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()
