# Copyright (c) 2013, Daniel Lindsley.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
#     1. Redistributions of source code must retain the above copyright notice,
#        this list of conditions and the following disclaimer.
#
#     2. Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#
#     3. Neither the name of shell nor the names of its contributors may be used
#        to endorse or promote products derived from this software without
#        specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
shell
=====

A better way to run shell commands in Python.

If you just need to quickly run a command, you can use the ``shell`` shortcut
function::

    >>> from shell import shell
    >>> ls = shell('ls')
    >>> for file in ls.output():
    ...     print file
    'another.txt'

If you need to extend the behavior, you can also use the ``Shell`` object::

    >>> from shell import Shell
    >>> sh = Shell(has_input=True)
    >>> cat = sh.run('cat -u')
    >>> cat.write('Hello, world!')
    >>> cat.output()
    ['Hello, world!']

"""
import shlex
import subprocess


__author__ = 'Daniel Lindsley'
__license__ = 'New BSD'
__version__ = (1, 0, 1)


class ShellException(Exception):
    """The base exception for all shell-related errors."""
    pass


class MissingCommandException(ShellException):
    """Thrown when no command was setup."""
    pass


class CommandError(ShellException):
    """Thrown when a command fails."""
    error_code = 1


class Shell(object):
    """
    Handles executing commands & recording output.

    Optionally accepts a ``has_input`` parameter, which should be a boolean.
    If set to ``True``, the command will wait to execute until you call the
    ``Shell.write`` method & send input. (Default: ``False``)

    Optionally accepts a ``record_output`` parameter, which should be a
    boolean.
    If set to ``True``, the stdout from the command will be recorded.
    (Default: ``True``)

    Optionally accepts a ``record_errors`` parameter, which should be a
    boolean.
    If set to ``True``, the stderr from the command will be recorded.
    (Default: ``True``)

    Optionally accepts a ``strip_empty`` parameter, which should be a boolean.
    If set to ``True``, only non-empty lines from ``Shell.output`` or
    ``Shell.errors`` will be returned. (Default: ``True``)
    """
    def __init__(self, has_input=False, record_output=True, record_errors=True,
                 strip_empty=True):
        self.has_input = has_input
        self.record_output = record_output
        self.record_errors = record_errors
        self.strip_empty = strip_empty

        self.last_command = ''
        self.line_breaks = '\n'
        self.pid = None
        self.code = 0
        self._popen = None
        self._stdout = ''
        self._stderr = ''

    def _split_command(self, command):
        """
        Splits a string command into the individual args to pass to ``Popen``.

        If the ``command`` is an array, it is left untouched.
        """
        if isinstance(command, (tuple, list)):
            return command

        return shlex.split(command)

    def _handle_output(self, stdout, stderr):
        """
        Takes stdout/stderr & optionally retains them internally.

        Requires a ``stdout`` parameter, which should either be the output as
        a string or ``None``.

        Requires a ``stderr`` parameter, which should either be the errors as
        a string or ``None``.

        Records nothing if the ``record_*`` options have been set to ``False``.
        """
        if self.record_output:
            if stdout is not None:
                self._stdout += stdout

        if self.record_errors:
            if stderr is not None:
                self._stderr += stderr

    def _communicate(self, the_input=None):
        """
        Handles communicating with the process & optionally sends input.

        Optionally accepts a ``the_input`` parameter, which can be a string
        to send to the process. (Default: ``None``)
        """
        stdout, stderr = self._popen.communicate(input=the_input)
        self._handle_output(stdout, stderr)

        if self._popen.returncode is not None:
            self.code = self._popen.returncode

    def run(self, command):
        """
        Runs a given command.

        Requires a ``command`` parameter should be either a string command
        (easier) or an array of arguments to send as the command (if you know
        what you're doing).

        Returns the ``Shell`` instance.

        Example::

            >>> from shell import Shell
            >>> sh = Shell()
            >>> sh.run('ls- alh')

        """
        self.last_command = command
        command_bits = self._split_command(command)
        kwargs = {
            'stdout': subprocess.PIPE,
            'stderr': subprocess.PIPE,
            'universal_newlines': True,
        }

        if self.has_input:
            kwargs['stdin'] = subprocess.PIPE

        self._popen = subprocess.Popen(
            command_bits,
            **kwargs
        )
        self.pid = self._popen.pid

        if not self.has_input:
            self._communicate()

        return self

    def write(self, the_input):
        """
        If you're working with an interactive process, sends that input to
        the process.

        This needs to be used in conjunction with the ``has_input=True``
        parameter.

        Requires a ``the_input`` parameter, which should be a string of the
        input to send to the command.

        Returns the ``Shell`` instance.

        Example::

            >>> from shell import Shell
            >>> sh = Shell(has_input=True)
            >>> sh.run('cat -u')
            >>> sh.write('Hello world!')

        """
        if not self._popen:
            raise MissingCommandException(
                "No command has been provided. Please call ``run`` first."
            )

        self._communicate(the_input)
        return self

    def kill(self):
        """
        Kills a given process.

        Example::

            >>> from shell import Shell
            >>> sh = Shell()
            >>> sh.run('some_long_running_thing')
            >>> sh.kill()

        """
        if not self._popen:
            raise MissingCommandException(
                "No command has been provided. Please call ``run`` first."
            )

        self._popen.kill()
        return self

    def output(self, raw=False):
        """
        Returns the output from running a command.

        Optionally accepts a ``raw`` parameter, which should be a boolean. If
        ``raw`` is set to ``False``, you get an array of lines of output. If
        ``raw`` is set to ``True``, the raw string of output is returned.
        (Default: ``False``)

        Example::

            >>> from shell import Shell
            >>> sh = Shell()
            >>> sh.run('ls ~')
            >>> sh.output()
            [
                'hello.txt',
                'world.txt',
            ]

        """
        if raw:
            return self._stdout

        lines = self._stdout.split(self.line_breaks)

        if self.strip_empty:
            lines = [line for line in lines if line]

        return lines

    def errors(self, raw=False):
        """
        Returns the errors from running a command.

        Optionally accepts a ``raw`` parameter, which should be a boolean. If
        ``raw`` is set to ``False``, you get an array of lines of errors. If
        ``raw`` is set to ``True``, the raw string of errors is returned.
        (Default: ``False``)

        Example::

            >>> from shell import Shell
            >>> sh = Shell()
            >>> sh.run('ls /there-s-no-way-anyone/has/this/directory/please')
            >>> sh.errors()
            [
                'ls /there-s-no-way-anyone/has/this/directory/please: No such
                 file or directory'
            ]

        """
        if raw:
            return self._stderr

        lines = self._stderr.split(self.line_breaks)

        if self.strip_empty:
            lines = [line for line in lines if line]

        return lines


def shell(command, has_input=False, record_output=True, record_errors=True,
          strip_empty=True):
    """
    A convenient shortcut for running commands.

    Requires a ``command`` parameter should be either a string command
    (easier) or an array of arguments to send as the command (if you know
    what you're doing).

    Optionally accepts a ``has_input`` parameter, which should be a boolean.
    If set to ``True``, the command will wait to execute until you call the
    ``Shell.write`` method & send input. (Default: ``False``)

    Optionally accepts a ``record_output`` parameter, which should be a
    boolean.
    If set to ``True``, the stdout from the command will be recorded.
    (Default: ``True``)

    Optionally accepts a ``record_errors`` parameter, which should be a
    boolean.
    If set to ``True``, the stderr from the command will be recorded.
    (Default: ``True``)

    Optionally accepts a ``strip_empty`` parameter, which should be a boolean.
    If set to ``True``, only non-empty lines from ``Shell.output`` or
    ``Shell.errors`` will be returned. (Default: ``True``)

    Returns the ``Shell`` instance, which has been run with the given command.

    Example::

        >>> from shell import shell
        >>> sh = shell('ls -alh *py')
        >>> sh.output()
        ['hello.py', 'world.py']

    """
    sh = Shell(
        has_input=has_input,
        record_output=record_output,
        record_errors=record_errors,
        strip_empty=strip_empty
    )
    return sh.run(command)
