<commit_before>import os
from subprocess import Popen, PIPE
import sys
import threading
import codecs
import locale

from .exceptions import Failure, PlatformError
from .platform import WINDOWS

from .vendor import six

def normalize_hide(val):
    hide_vals = (None, False, 'out', 'stdout', 'err', 'stderr', 'both', True)
    if val not in hide_vals:
        err = "'hide' got {0!r} which is not in {1!r}"
        raise ValueError(err.format(val, hide_vals))
    if val in (None, False):
        hide = ()
    elif val in ('both', True):
        hide = ('out', 'err')
    elif val == 'stdout':
        hide = ('out',)
    elif val == 'stderr':
        hide = ('err',)
    else:
        hide = (val,)
    return hide

# TODO: remove 'exception' field in run_* return values, if we don't run into
# situations similar to the one found in pexpect re: spurious IOErrors.

class Runner(object):
    """
    Partially-abstract core command-running API.

    This class is not usable by itself and must be subclassed to implement, at
    minimum, ``run_direct`` and ``run_pty``. An explanation of its methods
    follows:

    ``run`` is the primary API call which handles high level logic (echoing the
    commands locally, constructing a useful `Result` object, etc) and
    wraps/delegates to the below methods for actual execution.

    ``select_method`` takes ``pty`` and ``fallback`` kwargs and returns the
    current object's ``run_direct`` or ``run_pty`` method, depending on those
    kwargs' values and environmental cues/limitations. `Runner` itself has a
    useful default implementation of this method, but overriding may be
    sometimes useful.

    ``run_direct`` and ``run_pty`` are fully abstract in `Runner`; in
    subclasses, they should perform actual command execution, hooking directly
    into a subprocess' stdout/stderr pipes and returning those pipes' eventual
    full contents as distinct strings.

    ``run_direct``/``run_pty`` both have a signature of ``(self, command, warn,
    hide, encoding, out_stream, err_stream)`` (see `run` for semantics of
    these) and must return a 4-tuple of ``(stdout, stderr, exitcode,
    exception)`` (see `Result` for their meaning).

    ``run_pty`` differs from ``run_direct`` in that it should utilize a
    pseudo-terminal, and is expected to only return a useful ``stdout`` (with
    ``stderr`` usually set to the empty string.)

    For a subclass implementation example, see the source code for `.Local`.
    """
    def __init__(self, context):
        """
        Create a new runner with a handle on some `.Context`.

        :param context:
            a `.Context` instance, used to transmit default options and provide
            access to other contextualized information (e.g. a remote-oriented
            `.Runner` might want a `.Context` subclass holding info about
            hostnames and ports.)

            .. note::
                The `.Context` given to `.Runner` instances **must** contain
                default config values for the `.Runner` class in question. At a
                minimum, this means values for each of the default
                `.Runner.run` keyword arguments such as ``echo`` and ``warn``.

        :raises exceptions.ValueError:
            if not all expected default values are found in ``context``.
        """
        #: The `.Context` given to the same-named argument of `__init__`.
        self.context = context

    def run(self, command, **kwargs):
        """
        Execute ``command``, returning a `Result` object.

        .. note::
            All kwargs will default to the values found in this instance's
            `~.Runner.context` attribute, specifically in its configuration's
            ``run`` subtree (e.g. ``run.echo`` provides the default value for
            the ``echo`` keyword, etc). The base default values are described
            in the parameter list below.

        :param str command: The shell command to execute.

        :param bool warn:
            Whether to warn and continue, instead of raising `.Failure`, when
            the executed command exits with a nonzero status. Default:
            ``False``.

        :param hide:
            Allows the caller to disable ``run``'s default behavior of copying
            the subprocess' stdout and stderr to the controlling terminal.
            Specify ``hide='out'`` (or ``'stdout'``) to hide only the stdout
            stream, ``hide='err'`` (or ``'stderr'``) to hide only stderr, or
            ``hide='both'`` (or ``True``) to hide both streams.

            The default value is ``None``, meaning to print everything;
            ``False`` will also disable hiding.

            .. note::
                Stdout and stderr are always captured and stored in the
                ``Result`` object, regardless of ``hide``'s value.

        :param bool pty:
            By default, ``run`` connects directly to the invoked process and
            reads its stdout/stderr streams. Some programs will buffer (or even
            behave) differently in this situation compared to using an actual
            terminal or pty. To use a pty, specify ``pty=True``.

            .. warning::
                Due to their nature, ptys have a single output stream, so the
                ability to tell stdout apart from stderr is **not possible**
                when ``pty=True``. As such, all output will appear on
                ``out_stream`` (see below) and be captured into the ``stdout``
                result attribute. ``err_stream`` and ``stderr`` will always be
                empty when ``pty=True``.

        :param bool fallback:
            Controls auto-fallback behavior re: problems offering a pty when
            ``pty=True``. Whether this has any effect depends on the specific
            `Runner` subclass being invoked. Default: ``True``.

        :param bool echo:
            Controls whether `.run` prints the command string to local stdout
            prior to executing it. Default: ``False``.

        :param str encoding:
            Override auto-detection of which encoding the subprocess is using
            for its stdout/stderr streams. Defaults to the return value of
            ``locale.getpreferredencoding(False)``).

        :param out_stream:
            A file-like stream object to which the subprocess' standard error
            should be written. If ``None`` (the default), ``sys.stdout`` will
            be used.

        :param err_stream:
            Same as ``out_stream``, except for standard error, and defaulting
            to ``sys.stderr``.

        :returns: `Result`

        :raises: `.Failure` (if the command exited nonzero & ``warn=False``)
        """
        exception = False
        # Normalize kwargs w/ config
        opts = {}
        for key, value in six.iteritems(self.context.config.run):
            runtime = kwargs.pop(key, None)
            opts[key] = value if runtime is None else runtime
        # TODO: handle invalid kwarg keys (anything left in kwargs)
        # Normalize 'hide' from one of the various valid input values
        opts['hide'] = normalize_hide(opts['hide'])
        # Derive stream objects
        out_stream = sys.stdout if opts['out_stream'] is None else opts['out_stream']
        err_stream = sys.stderr if opts['err_stream'] is None else opts['err_stream']
        # Do the things
        if opts['echo']:
            print("\033[1;37m{0}\033[0m".format(command))
        func = self.select_method(pty=opts['pty'], fallback=opts['fallback'])
        stdout, stderr, exited, exception = func(
            command=command,
            warn=opts['warn'],
            hide=opts['hide'],
            encoding=opts['encoding'],
            out_stream=out_stream,
            err_stream=err_stream,
        )
        # TODO: make this test less gross? Feels silly to just return a bool in
        # select_method which is tantamount to this, though.
        func_name = getattr(func, 'func_name', getattr(func, '__name__'))
        used_pty = func_name == 'run_pty'
        result = Result(
            stdout=stdout,
            stderr=stderr,
            exited=exited,
            pty=used_pty,
            exception=exception,
        )
        if not (result or opts['warn']):
            raise Failure(result)
        return result

    def select_method(self, pty, fallback):
        # NOTE: fallback not used: no falling back implemented by default.
        return getattr(self, 'run_pty' if pty else 'run_direct')

    def run_direct(self, command, warn, hide, encoding, out_stream, err_stream):
        raise NotImplementedError

    def run_pty(self, command, warn, hide, encoding, out_stream, err_stream):
        raise NotImplementedError

class Local(Runner):
    """
    Execute a command on the local system in a subprocess.

    .. note::
        When Invoke itself is executed without a valid PTY (i.e.
        ``os.isatty(sys.stdin)`` is ``False``), it's not possible to present a
        handle on our PTY to local subprocesses. In such situations, `Local`
        will fallback to behaving as if ``pty=False``, on the theory that
        degraded execution is better than none at all, as well as printing a
        warning to stderr.

        To disable this behavior (i.e. if ``os.isatty`` is causing false
        negatives in your environment), say ``fallback=False``.
    """
    def select_method(self, pty=False, fallback=True):
        func = self.run_direct
        if pty:
            func = self.run_pty
            if not os.isatty(sys.stdin.fileno()) and fallback:
                sys.stderr.write("WARNING: stdin is not a pty; falling back to non-pty execution!\n") # noqa
                func = self.run_direct
        return func

    def run_direct(self, command, warn, hide, encoding, out_stream, err_stream):
        process = Popen(
            command,
            shell=True,
            stdout=PIPE,
            stderr=PIPE,
        )

        if encoding is None:
            encoding = locale.getpreferredencoding(False)

        def display(src, dst, cap, hide):
            def get():
                while True:
                    data = os.read(src.fileno(), 1000)
                    if not data:
                        break
                    yield data
            for data in codecs.iterdecode(get(), encoding, errors='replace'):
                if not hide:
                    dst.write(data)
                    dst.flush()
                cap.append(data)

        stdout = []
        stderr = []
        threads = []

        for args in (
            (process.stdout, out_stream, stdout, 'out' in hide),
            (process.stderr, err_stream, stderr, 'err' in hide),
        ):
            t = threading.Thread(target=display, args=args)
            threads.append(t)
            t.start()

        process.wait()
        for t in threads:
            t.join()

        stdout = ''.join(stdout)
        stderr = ''.join(stderr)
        if WINDOWS:
            # "Universal newlines" - replace all standard forms of
            # newline with \n. This is not technically Windows related
            # (\r as newline is an old Mac convention) but we only apply
            # the translation for Windows as that's the only platform
            # it is likely to matter for these days.
            stdout = stdout.replace("\r\n", "\n").replace("\r", "\n")
            stderr = stderr.replace("\r\n", "\n").replace("\r", "\n")

        return stdout, stderr, process.returncode, None

    def run_pty(self, command, warn, hide, encoding, out_stream, err_stream):
        # TODO: re-insert Windows "lol y u no pty" stuff here
        import pty
        pid, parent_fd = pty.fork()
        # If we're the child process, load up the actual command in a shell,
        # just as subprocess does; this replaces our process - whose pipes are
        # all hooked up to the PTY - with the "real" one.
        if pid == 0:
            # Use execv for bare-minimum "exec w/ variable # args" behavior.
            # No need for the 'p' (use PATH to find executable) or 'e' (define
            # a custom/overridden shell env) variants, for now.
            # TODO: use /bin/sh or whatever subprocess does. Only using bash
            # for now because that's what we have been testing against.
            # TODO: also see if subprocess is using equivalent of execvp...
            # TODO: both pty.spawn() and pexpect.spawn() do a lot of
            # setup/teardown involving tty.*, setwinsize, getrlimit, signal.
            # Ostensibly we'll want some of that eventually, but if possible
            # write tests - integration-level if necessary - before adding it!
            os.execv('/bin/bash', ['/bin/bash', '-c', command])

        if encoding is None:
            encoding = locale.getpreferredencoding(False)

        def display(src_fd, dst, cap, hide):
            def get():
                while True:
                    data = os.read(src_fd, 1000)
                    if not data:
                        break
                    yield data
            for data in codecs.iterdecode(get(), encoding, errors='replace'):
                if not hide:
                    dst.write(data)
                    dst.flush()
                cap.append(data)

        stdout = []
        stderr = []
        threads = []

        # TODO: can we simply run this loop in the main thread now?
        # TODO: or is there some other benefit to threading it, such as
        # eventual stdin control?
        for args in (
            (parent_fd, out_stream, stdout, 'out' in hide),
            #(process.stderr, err_stream, stderr, 'err' in hide),
        ):
            t = threading.Thread(target=display, args=args)
            threads.append(t)
            t.start()

        # Wait in main thread until child appears to have exited.
        while True:
            # TODO: set 2nd value to os.WNOHANG in some situations?
            pid_val, status = os.waitpid(pid, 0)
            # waitpid() sets the 'pid' return val to 0 when no children have
            # exited yet; when it is NOT zero, we know the child's stopped.
            if pid_val != 0:
                break
            # TODO: io sleep?

        # Derive return code from waitpid() result
        returncode = os.WEXITSTATUS(status)

        # Join reader threads & pack things up
        for t in threads:
            t.join()

        stdout = ''.join(stdout)
        stderr = ''.join(stderr)
        if WINDOWS:
            # "Universal newlines" - replace all standard forms of
            # newline with \n. This is not technically Windows related
            # (\r as newline is an old Mac convention) but we only apply
            # the translation for Windows as that's the only platform
            # it is likely to matter for these days.
            stdout = stdout.replace("\r\n", "\n").replace("\r", "\n")
            stderr = stderr.replace("\r\n", "\n").replace("\r", "\n")

        return stdout, stderr, returncode, None

class Result(object):
    """
    A container for information about the result of a command execution.

    `Result` instances have the following attributes:

    * ``stdout``: The subprocess' standard output, as a multiline string.
    * ``stderr``: Same as ``stdout`` but containing standard error (unless
      the process was invoked via a pty; see `.Runner.run`.)
    * ``exited``: An integer representing the subprocess' exit/return code.
    * ``return_code``: An alias to ``exited``.
    * ``ok``: A boolean equivalent to ``exited == 0``.
    * ``failed``: The inverse of ``ok``: ``True`` if the program exited with a
      nonzero return code.
    * ``pty``: A boolean describing whether the subprocess was invoked with a
      pty or not; see `.Runner.run`.
    * ``exception``: Typically ``None``, but may be an exception object if
      ``pty`` was ``True`` and ``run`` had to swallow an apparently-spurious
      ``OSError``. Solely for sanity checking/debugging purposes.

    `Result` objects' truth evaluation is equivalent to their ``ok``
    attribute's value.
    """
    # TODO: inherit from namedtuple instead? heh
    def __init__(self, stdout, stderr, exited, pty, exception=None):
        self.exited = self.return_code = exited
        self.stdout = stdout
        self.stderr = stderr
        self.pty = pty
        self.exception = exception

    def __nonzero__(self):
        # Holy mismatch between name and implementation, Batman!
        return self.exited == 0

    # Python 3 ahoy
    def __bool__(self):
        return self.__nonzero__()

    def __str__(self):
        ret = ["Command exited with status {0}.".format(self.exited)]
        for x in ('stdout', 'stderr'):
            val = getattr(self, x)
            ret.append("""=== {0} ===
{1}
""".format(x, val.rstrip()) if val else "(no {0})".format(x))
        return "\n".join(ret)

    @property
    def ok(self):
        return self.exited == 0

    @property
    def failed(self):
        return not self.ok
<commit_msg>First steps towards unification<commit_after>import os
from subprocess import Popen, PIPE
import sys
import threading
import codecs
import locale

from .exceptions import Failure, PlatformError
from .platform import WINDOWS

from .vendor import six

def normalize_hide(val):
    hide_vals = (None, False, 'out', 'stdout', 'err', 'stderr', 'both', True)
    if val not in hide_vals:
        err = "'hide' got {0!r} which is not in {1!r}"
        raise ValueError(err.format(val, hide_vals))
    if val in (None, False):
        hide = ()
    elif val in ('both', True):
        hide = ('out', 'err')
    elif val == 'stdout':
        hide = ('out',)
    elif val == 'stderr':
        hide = ('err',)
    else:
        hide = (val,)
    return hide

# TODO: remove 'exception' field in run_* return values, if we don't run into
# situations similar to the one found in pexpect re: spurious IOErrors.

class Runner(object):
    """
    Partially-abstract core command-running API.

    This class is not usable by itself and must be subclassed to implement, at
    minimum, ``run_direct`` and ``run_pty``. An explanation of its methods
    follows:

    ``run`` is the primary API call which handles high level logic (echoing the
    commands locally, constructing a useful `Result` object, etc) and
    wraps/delegates to the below methods for actual execution.

    ``select_method`` takes ``pty`` and ``fallback`` kwargs and returns the
    current object's ``run_direct`` or ``run_pty`` method, depending on those
    kwargs' values and environmental cues/limitations. `Runner` itself has a
    useful default implementation of this method, but overriding may be
    sometimes useful.

    ``run_direct`` and ``run_pty`` are fully abstract in `Runner`; in
    subclasses, they should perform actual command execution, hooking directly
    into a subprocess' stdout/stderr pipes and returning those pipes' eventual
    full contents as distinct strings.

    ``run_direct``/``run_pty`` both have a signature of ``(self, command, warn,
    hide, encoding, out_stream, err_stream)`` (see `run` for semantics of
    these) and must return a 4-tuple of ``(stdout, stderr, exitcode,
    exception)`` (see `Result` for their meaning).

    ``run_pty`` differs from ``run_direct`` in that it should utilize a
    pseudo-terminal, and is expected to only return a useful ``stdout`` (with
    ``stderr`` usually set to the empty string.)

    For a subclass implementation example, see the source code for `.Local`.
    """
    def __init__(self, context):
        """
        Create a new runner with a handle on some `.Context`.

        :param context:
            a `.Context` instance, used to transmit default options and provide
            access to other contextualized information (e.g. a remote-oriented
            `.Runner` might want a `.Context` subclass holding info about
            hostnames and ports.)

            .. note::
                The `.Context` given to `.Runner` instances **must** contain
                default config values for the `.Runner` class in question. At a
                minimum, this means values for each of the default
                `.Runner.run` keyword arguments such as ``echo`` and ``warn``.

        :raises exceptions.ValueError:
            if not all expected default values are found in ``context``.
        """
        #: The `.Context` given to the same-named argument of `__init__`.
        self.context = context

    def run(self, command, **kwargs):
        """
        Execute ``command``, returning a `Result` object.

        .. note::
            All kwargs will default to the values found in this instance's
            `~.Runner.context` attribute, specifically in its configuration's
            ``run`` subtree (e.g. ``run.echo`` provides the default value for
            the ``echo`` keyword, etc). The base default values are described
            in the parameter list below.

        :param str command: The shell command to execute.

        :param bool warn:
            Whether to warn and continue, instead of raising `.Failure`, when
            the executed command exits with a nonzero status. Default:
            ``False``.

        :param hide:
            Allows the caller to disable ``run``'s default behavior of copying
            the subprocess' stdout and stderr to the controlling terminal.
            Specify ``hide='out'`` (or ``'stdout'``) to hide only the stdout
            stream, ``hide='err'`` (or ``'stderr'``) to hide only stderr, or
            ``hide='both'`` (or ``True``) to hide both streams.

            The default value is ``None``, meaning to print everything;
            ``False`` will also disable hiding.

            .. note::
                Stdout and stderr are always captured and stored in the
                ``Result`` object, regardless of ``hide``'s value.

        :param bool pty:
            By default, ``run`` connects directly to the invoked process and
            reads its stdout/stderr streams. Some programs will buffer (or even
            behave) differently in this situation compared to using an actual
            terminal or pty. To use a pty, specify ``pty=True``.

            .. warning::
                Due to their nature, ptys have a single output stream, so the
                ability to tell stdout apart from stderr is **not possible**
                when ``pty=True``. As such, all output will appear on
                ``out_stream`` (see below) and be captured into the ``stdout``
                result attribute. ``err_stream`` and ``stderr`` will always be
                empty when ``pty=True``.

        :param bool fallback:
            Controls auto-fallback behavior re: problems offering a pty when
            ``pty=True``. Whether this has any effect depends on the specific
            `Runner` subclass being invoked. Default: ``True``.

        :param bool echo:
            Controls whether `.run` prints the command string to local stdout
            prior to executing it. Default: ``False``.

        :param str encoding:
            Override auto-detection of which encoding the subprocess is using
            for its stdout/stderr streams. Defaults to the return value of
            ``locale.getpreferredencoding(False)``).

        :param out_stream:
            A file-like stream object to which the subprocess' standard error
            should be written. If ``None`` (the default), ``sys.stdout`` will
            be used.

        :param err_stream:
            Same as ``out_stream``, except for standard error, and defaulting
            to ``sys.stderr``.

        :returns: `Result`

        :raises: `.Failure` (if the command exited nonzero & ``warn=False``)
        """
        exception = False
        # Normalize kwargs w/ config
        opts = {}
        for key, value in six.iteritems(self.context.config.run):
            runtime = kwargs.pop(key, None)
            opts[key] = value if runtime is None else runtime
        # TODO: handle invalid kwarg keys (anything left in kwargs)
        # Normalize 'hide' from one of the various valid input values
        opts['hide'] = normalize_hide(opts['hide'])
        # Derive stream objects
        out_stream = sys.stdout if opts['out_stream'] is None else opts['out_stream']
        err_stream = sys.stderr if opts['err_stream'] is None else opts['err_stream']
        # Do the things
        if opts['echo']:
            print("\033[1;37m{0}\033[0m".format(command))
        func = self.select_method(pty=opts['pty'], fallback=opts['fallback'])
        stdout, stderr, exited, exception = func(
            command=command,
            warn=opts['warn'],
            hide=opts['hide'],
            encoding=opts['encoding'],
            out_stream=out_stream,
            err_stream=err_stream,
        )
        # TODO: make this test less gross? Feels silly to just return a bool in
        # select_method which is tantamount to this, though.
        func_name = getattr(func, 'func_name', getattr(func, '__name__'))
        used_pty = func_name == 'run_pty'
        result = Result(
            stdout=stdout,
            stderr=stderr,
            exited=exited,
            pty=used_pty,
            exception=exception,
        )
        if not (result or opts['warn']):
            raise Failure(result)
        return result

    def select_method(self, pty, fallback):
        # NOTE: fallback not used: no falling back implemented by default.
        return getattr(self, 'run_pty' if pty else 'run_direct')

    def run_direct(self, command, warn, hide, encoding, out_stream, err_stream):
        raise NotImplementedError

    def run_pty(self, command, warn, hide, encoding, out_stream, err_stream):
        raise NotImplementedError

class Local(Runner):
    """
    Execute a command on the local system in a subprocess.

    .. note::
        When Invoke itself is executed without a valid PTY (i.e.
        ``os.isatty(sys.stdin)`` is ``False``), it's not possible to present a
        handle on our PTY to local subprocesses. In such situations, `Local`
        will fallback to behaving as if ``pty=False``, on the theory that
        degraded execution is better than none at all, as well as printing a
        warning to stderr.

        To disable this behavior (i.e. if ``os.isatty`` is causing false
        negatives in your environment), say ``fallback=False``.
    """
    def select_method(self, pty=False, fallback=True):
        func = self.run_direct
        if pty:
            func = self.run_pty
            if not os.isatty(sys.stdin.fileno()) and fallback:
                sys.stderr.write("WARNING: stdin is not a pty; falling back to non-pty execution!\n") # noqa
                func = self.run_direct
        return func

    def _normalize_encoding(self, encoding):
        return locale.getpreferredencoding(False)

    def _mux(self, source_fd, dest, buffer_, hide, encoding):
        # Inner generator yielding read data
        def get():
            while True:
                data = os.read(source_fd, 1000)
                if not data:
                    break
                yield data
        # Use generator in iterdecode() to decode stream data, then print/save
        for data in codecs.iterdecode(get(), encoding, errors='replace'):
            if not hide:
                dest.write(data)
                dest.flush()
            buffer_.append(data)

    def run_direct(self, command, warn, hide, encoding, out_stream, err_stream):
        process = Popen(
            command,
            shell=True,
            stdout=PIPE,
            stderr=PIPE,
        )

        encoding = self._normalize_encoding(encoding)

        stdout = []
        stderr = []
        threads = []

        for args in (
            (process.stdout.fileno(), out_stream, stdout, 'out' in hide, encoding),
            (process.stderr.fileno(), err_stream, stderr, 'err' in hide, encoding),
        ):
            t = threading.Thread(target=self._mux, args=args)
            threads.append(t)
            t.start()

        process.wait()
        for t in threads:
            t.join()

        stdout = ''.join(stdout)
        stderr = ''.join(stderr)
        if WINDOWS:
            # "Universal newlines" - replace all standard forms of
            # newline with \n. This is not technically Windows related
            # (\r as newline is an old Mac convention) but we only apply
            # the translation for Windows as that's the only platform
            # it is likely to matter for these days.
            stdout = stdout.replace("\r\n", "\n").replace("\r", "\n")
            stderr = stderr.replace("\r\n", "\n").replace("\r", "\n")

        return stdout, stderr, process.returncode, None

    def run_pty(self, command, warn, hide, encoding, out_stream, err_stream):
        # TODO: re-insert Windows "lol y u no pty" stuff here
        import pty
        pid, parent_fd = pty.fork()
        # If we're the child process, load up the actual command in a shell,
        # just as subprocess does; this replaces our process - whose pipes are
        # all hooked up to the PTY - with the "real" one.
        if pid == 0:
            # Use execv for bare-minimum "exec w/ variable # args" behavior.
            # No need for the 'p' (use PATH to find executable) or 'e' (define
            # a custom/overridden shell env) variants, for now.
            # TODO: use /bin/sh or whatever subprocess does. Only using bash
            # for now because that's what we have been testing against.
            # TODO: also see if subprocess is using equivalent of execvp...
            # TODO: both pty.spawn() and pexpect.spawn() do a lot of
            # setup/teardown involving tty.*, setwinsize, getrlimit, signal.
            # Ostensibly we'll want some of that eventually, but if possible
            # write tests - integration-level if necessary - before adding it!
            os.execv('/bin/bash', ['/bin/bash', '-c', command])

        encoding = self._normalize_encoding(encoding)

        def display(src_fd, dst, cap, hide):
            def get():
                while True:
                    data = os.read(src_fd, 1000)
                    if not data:
                        break
                    yield data
            for data in codecs.iterdecode(get(), encoding, errors='replace'):
                if not hide:
                    dst.write(data)
                    dst.flush()
                cap.append(data)

        stdout = []
        stderr = []
        threads = []

        # TODO: can we simply run this loop in the main thread now?
        # TODO: or is there some other benefit to threading it, such as
        # eventual stdin control?
        for args in (
            (parent_fd, out_stream, stdout, 'out' in hide),
            #(process.stderr, err_stream, stderr, 'err' in hide),
        ):
            t = threading.Thread(target=display, args=args)
            threads.append(t)
            t.start()

        # Wait in main thread until child appears to have exited.
        while True:
            # TODO: set 2nd value to os.WNOHANG in some situations?
            pid_val, status = os.waitpid(pid, 0)
            # waitpid() sets the 'pid' return val to 0 when no children have
            # exited yet; when it is NOT zero, we know the child's stopped.
            if pid_val != 0:
                break
            # TODO: io sleep?

        # Derive return code from waitpid() result
        returncode = os.WEXITSTATUS(status)

        # Join reader threads & pack things up
        for t in threads:
            t.join()

        stdout = ''.join(stdout)
        stderr = ''.join(stderr)
        if WINDOWS:
            # "Universal newlines" - replace all standard forms of
            # newline with \n. This is not technically Windows related
            # (\r as newline is an old Mac convention) but we only apply
            # the translation for Windows as that's the only platform
            # it is likely to matter for these days.
            stdout = stdout.replace("\r\n", "\n").replace("\r", "\n")
            stderr = stderr.replace("\r\n", "\n").replace("\r", "\n")

        return stdout, stderr, returncode, None

class Result(object):
    """
    A container for information about the result of a command execution.

    `Result` instances have the following attributes:

    * ``stdout``: The subprocess' standard output, as a multiline string.
    * ``stderr``: Same as ``stdout`` but containing standard error (unless
      the process was invoked via a pty; see `.Runner.run`.)
    * ``exited``: An integer representing the subprocess' exit/return code.
    * ``return_code``: An alias to ``exited``.
    * ``ok``: A boolean equivalent to ``exited == 0``.
    * ``failed``: The inverse of ``ok``: ``True`` if the program exited with a
      nonzero return code.
    * ``pty``: A boolean describing whether the subprocess was invoked with a
      pty or not; see `.Runner.run`.
    * ``exception``: Typically ``None``, but may be an exception object if
      ``pty`` was ``True`` and ``run`` had to swallow an apparently-spurious
      ``OSError``. Solely for sanity checking/debugging purposes.

    `Result` objects' truth evaluation is equivalent to their ``ok``
    attribute's value.
    """
    # TODO: inherit from namedtuple instead? heh
    def __init__(self, stdout, stderr, exited, pty, exception=None):
        self.exited = self.return_code = exited
        self.stdout = stdout
        self.stderr = stderr
        self.pty = pty
        self.exception = exception

    def __nonzero__(self):
        # Holy mismatch between name and implementation, Batman!
        return self.exited == 0

    # Python 3 ahoy
    def __bool__(self):
        return self.__nonzero__()

    def __str__(self):
        ret = ["Command exited with status {0}.".format(self.exited)]
        for x in ('stdout', 'stderr'):
            val = getattr(self, x)
            ret.append("""=== {0} ===
{1}
""".format(x, val.rstrip()) if val else "(no {0})".format(x))
        return "\n".join(ret)

    @property
    def ok(self):
        return self.exited == 0

    @property
    def failed(self):
        return not self.ok
<|endoftext|>


<commit_before>import os
from subprocess import Popen, PIPE
import sys
import threading
import codecs
import locale

from .exceptions import Failure, PlatformError
from .platform import WINDOWS

from .vendor import six

def normalize_hide(val):
    hide_vals = (None, False, 'out', 'stdout', 'err', 'stderr', 'both', True)
    if val not in hide_vals:
        err = "'hide' got {0!r} which is not in {1!r}"
        raise ValueError(err.format(val, hide_vals))
    if val in (None, False):
        hide = ()
    elif val in ('both', True):
        hide = ('out', 'err')
    elif val == 'stdout':
        hide = ('out',)
    elif val == 'stderr':
        hide = ('err',)
    else:
        hide = (val,)
    return hide

class Runner(object):
    """
    Partially-abstract core command-running API.

    This class is not usable by itself and must be subclassed to implement, at
    minimum, ``run_direct`` and ``run_pty``. An explanation of its methods
    follows:

    ``run`` is the primary API call which handles high level logic (echoing the
    commands locally, constructing a useful `Result` object, etc) and
    wraps/delegates to the below methods for actual execution.

    ``select_method`` takes ``pty`` and ``fallback`` kwargs and returns the
    current object's ``run_direct`` or ``run_pty`` method, depending on those
    kwargs' values and environmental cues/limitations. `Runner` itself has a
    useful default implementation of this method, but overriding may be
    sometimes useful.

    ``run_direct`` and ``run_pty`` are fully abstract in `Runner`; in
    subclasses, they should perform actual command execution, hooking directly
    into a subprocess' stdout/stderr pipes and returning those pipes' eventual
    full contents as distinct strings.

    ``run_direct``/``run_pty`` both have a signature of ``(self, command, warn,
    hide, encoding, out_stream, err_stream)`` (see `run` for semantics of
    these) and must return a 4-tuple of ``(stdout, stderr, exitcode,
    exception)`` (see `Result` for their meaning).

    ``run_pty`` differs from ``run_direct`` in that it should utilize a
    pseudo-terminal, and is expected to only return a useful ``stdout`` (with
    ``stderr`` usually set to the empty string.)

    For a subclass implementation example, see the source code for `.Local`.
    """
    def __init__(self, context):
        """
        Create a new runner with a handle on some `.Context`.

        :param context:
            a `.Context` instance, used to transmit default options and provide
            access to other contextualized information (e.g. a remote-oriented
            `.Runner` might want a `.Context` subclass holding info about
            hostnames and ports.)

            .. note::
                The `.Context` given to `.Runner` instances **must** contain
                default config values for the `.Runner` class in question. At a
                minimum, this means values for each of the default
                `.Runner.run` keyword arguments such as ``echo`` and ``warn``.

        :raises exceptions.ValueError:
            if not all expected default values are found in ``context``.
        """
        #: The `.Context` given to the same-named argument of `__init__`.
        self.context = context

    def run(self, command, **kwargs):
        """
        Execute ``command``, returning a `Result` object.

        .. note::
            All kwargs will default to the values found in this instance's
            `~.Runner.context` attribute, specifically in its configuration's
            ``run`` subtree (e.g. ``run.echo`` provides the default value for
            the ``echo`` keyword, etc). The base default values are described
            in the parameter list below.

        :param str command: The shell command to execute.

        :param bool warn:
            Whether to warn and continue, instead of raising `.Failure`, when
            the executed command exits with a nonzero status. Default:
            ``False``.

        :param hide:
            Allows the caller to disable ``run``'s default behavior of copying
            the subprocess' stdout and stderr to the controlling terminal.
            Specify ``hide='out'`` (or ``'stdout'``) to hide only the stdout
            stream, ``hide='err'`` (or ``'stderr'``) to hide only stderr, or
            ``hide='both'`` (or ``True``) to hide both streams.

            The default value is ``None``, meaning to print everything;
            ``False`` will also disable hiding.

            .. note::
                Stdout and stderr are always captured and stored in the
                ``Result`` object, regardless of ``hide``'s value.

        :param bool pty:
            By default, ``run`` connects directly to the invoked process and
            reads its stdout/stderr streams. Some programs will buffer (or even
            behave) differently in this situation compared to using an actual
            terminal or pty. To use a pty, specify ``pty=True``.

            .. warning::
                Due to their nature, ptys have a single output stream, so the
                ability to tell stdout apart from stderr is **not possible**
                when ``pty=True``. As such, all output will appear on
                ``out_stream`` (see below) and be captured into the ``stdout``
                result attribute. ``err_stream`` and ``stderr`` will always be
                empty when ``pty=True``.

        :param bool fallback:
            Controls auto-fallback behavior re: problems offering a pty when
            ``pty=True``. Whether this has any effect depends on the specific
            `Runner` subclass being invoked. Default: ``True``.

        :param bool echo:
            Controls whether `.run` prints the command string to local stdout
            prior to executing it. Default: ``False``.

        :param str encoding:
            Override auto-detection of which encoding the subprocess is using
            for its stdout/stderr streams. Defaults to the return value of
            ``locale.getpreferredencoding(False)``).

        :param out_stream:
            A file-like stream object to which the subprocess' standard error
            should be written. If ``None`` (the default), ``sys.stdout`` will
            be used.

        :param err_stream:
            Same as ``out_stream``, except for standard error, and defaulting
            to ``sys.stderr``.

        :returns: `Result`

        :raises: `.Failure` (if the command exited nonzero & ``warn=False``)
        """
        exception = False
        # Normalize kwargs w/ config
        opts = {}
        for key, value in six.iteritems(self.context.config.run):
            runtime = kwargs.pop(key, None)
            opts[key] = value if runtime is None else runtime
        # TODO: handle invalid kwarg keys (anything left in kwargs)
        # Normalize 'hide' from one of the various valid input values
        opts['hide'] = normalize_hide(opts['hide'])
        # Derive stream objects
        out_stream = sys.stdout if opts['out_stream'] is None else opts['out_stream']
        err_stream = sys.stderr if opts['err_stream'] is None else opts['err_stream']
        # Do the things
        if opts['echo']:
            print("\033[1;37m{0}\033[0m".format(command))
        func = self.select_method(pty=opts['pty'], fallback=opts['fallback'])
        stdout, stderr, exited, exception = func(
            command=command,
            warn=opts['warn'],
            hide=opts['hide'],
            encoding=opts['encoding'],
            out_stream=out_stream,
            err_stream=err_stream,
        )
        # TODO: make this test less gross? Feels silly to just return a bool in
        # select_method which is tantamount to this, though.
        func_name = getattr(func, 'func_name', getattr(func, '__name__'))
        used_pty = func_name == 'run_pty'
        result = Result(
            stdout=stdout,
            stderr=stderr,
            exited=exited,
            pty=used_pty,
            exception=exception,
        )
        if not (result or opts['warn']):
            raise Failure(result)
        return result

    def select_method(self, pty, fallback):
        # NOTE: fallback not used: no falling back implemented by default.
        return getattr(self, 'run_pty' if pty else 'run_direct')

    def run_direct(self, command, warn, hide, encoding, out_stream, err_stream):
        raise NotImplementedError

    def run_pty(self, command, warn, hide, encoding, out_stream, err_stream):
        raise NotImplementedError

class Local(Runner):
    """
    Execute a command on the local system in a subprocess.

    .. note::
        When Invoke itself is executed without a valid PTY (i.e.
        ``os.isatty(sys.stdin)`` is ``False``), it's not possible to present a
        handle on our PTY to local subprocesses. In such situations, `Local`
        will fallback to behaving as if ``pty=False``, on the theory that
        degraded execution is better than none at all, as well as printing a
        warning to stderr.

        To disable this behavior (i.e. if ``os.isatty`` is causing false
        negatives in your environment), say ``fallback=False``.
    """
    def select_method(self, pty=False, fallback=True):
        func = self.run_direct
        if pty:
            func = self.run_pty
            if not os.isatty(sys.stdin.fileno()) and fallback:
                sys.stderr.write("WARNING: stdin is not a pty; falling back to non-pty execution!\n") # noqa
                func = self.run_direct
        return func

    def run_direct(self, command, warn, hide, encoding, out_stream, err_stream):
        process = Popen(
            command,
            shell=True,
            stdout=PIPE,
            stderr=PIPE,
        )

        if encoding is None:
            encoding = locale.getpreferredencoding(False)

        def display(src, dst, cap, hide):
            def get():
                while True:
                    data = os.read(src.fileno(), 1000)
                    if not data:
                        break
                    yield data
            for data in codecs.iterdecode(get(), encoding, errors='replace'):
                if not hide:
                    dst.write(data)
                    dst.flush()
                cap.append(data)

        stdout = []
        stderr = []
        threads = []

        for args in (
            (process.stdout, out_stream, stdout, 'out' in hide),
            (process.stderr, err_stream, stderr, 'err' in hide),
        ):
            t = threading.Thread(target=display, args=args)
            threads.append(t)
            t.start()

        process.wait()
        for t in threads:
            t.join()

        stdout = ''.join(stdout)
        stderr = ''.join(stderr)
        if WINDOWS:
            # "Universal newlines" - replace all standard forms of
            # newline with \n. This is not technically Windows related
            # (\r as newline is an old Mac convention) but we only apply
            # the translation for Windows as that's the only platform
            # it is likely to matter for these days.
            stdout = stdout.replace("\r\n", "\n").replace("\r", "\n")
            stderr = stderr.replace("\r\n", "\n").replace("\r", "\n")

        return stdout, stderr, process.returncode, None

    def run_pty(self, command, warn, hide, encoding, out_stream, err_stream):
        # Sanity check: platforms that can't pexpect should explode usefully
        # here. (Without this, the pexpect import throws an inner
        # ImportException trying to 'import pty' which is unavailable on
        # Windows. Better to do this here than truly-fork pexpect.)
        if WINDOWS:
            err = "You seem to be on Windows, which doesn't support ptys!"
            raise PlatformError(err)
        # Proceed as normal for POSIX/etc platforms, with a runtime import
        from .vendor import pexpect

        out = []
        def out_filter(text):
            # Maybe use encoding here as in run() above...
            out.append(text.decode("utf-8", 'replace'))
            if 'out' not in hide:
                return text
            else:
                return b""
        # Hand pexpect useful args for os.execve(binary, args).
        # TODO: try to use same default binary value as Popen does (is usually
        # /bin/sh but IIRC there's caveats)
        # TODO: allow control over this
        p = pexpect.spawn("/bin/bash", ["-c", command])
        # Ensure pexpect doesn't barf with OSError if we fall off the end of
        # the child's input on some platforms (e.g. Linux).
        exception = None
        try:
            p.interact(output_filter=out_filter)
        except OSError as e:
            # Only capture the OSError we expect
            if "Input/output error" not in str(e):
                raise
            # Ensure it ties off the child, sets exitstatus, etc
            p.close()
            # Capture the exception in case it's NOT the OSError we think it
            # is and folks need to debug
            exception = e
        return "".join(out), "", p.exitstatus, exception

class Result(object):
    """
    A container for information about the result of a command execution.

    `Result` instances have the following attributes:

    * ``stdout``: The subprocess' standard output, as a multiline string.
    * ``stderr``: Same as ``stdout`` but containing standard error (unless
      the process was invoked via a pty; see `.Runner.run`.)
    * ``exited``: An integer representing the subprocess' exit/return code.
    * ``return_code``: An alias to ``exited``.
    * ``ok``: A boolean equivalent to ``exited == 0``.
    * ``failed``: The inverse of ``ok``: ``True`` if the program exited with a
      nonzero return code.
    * ``pty``: A boolean describing whether the subprocess was invoked with a
      pty or not; see `.Runner.run`.
    * ``exception``: Typically ``None``, but may be an exception object if
      ``pty`` was ``True`` and ``run`` had to swallow an apparently-spurious
      ``OSError``. Solely for sanity checking/debugging purposes.

    `Result` objects' truth evaluation is equivalent to their ``ok``
    attribute's value.
    """
    # TODO: inherit from namedtuple instead? heh
    def __init__(self, stdout, stderr, exited, pty, exception=None):
        self.exited = self.return_code = exited
        self.stdout = stdout
        self.stderr = stderr
        self.pty = pty
        self.exception = exception

    def __nonzero__(self):
        # Holy mismatch between name and implementation, Batman!
        return self.exited == 0

    # Python 3 ahoy
    def __bool__(self):
        return self.__nonzero__()

    def __str__(self):
        ret = ["Command exited with status {0}.".format(self.exited)]
        for x in ('stdout', 'stderr'):
            val = getattr(self, x)
            ret.append("""=== {0} ===
{1}
""".format(x, val.rstrip()) if val else "(no {0})".format(x))
        return "\n".join(ret)

    @property
    def ok(self):
        return self.exited == 0

    @property
    def failed(self):
        return not self.ok
<commit_msg>First draft of non-pexpect-using run_pty. Works in trivial case already!<commit_after>import os
from subprocess import Popen, PIPE
import sys
import threading
import codecs
import locale

from .exceptions import Failure, PlatformError
from .platform import WINDOWS

from .vendor import six

def normalize_hide(val):
    hide_vals = (None, False, 'out', 'stdout', 'err', 'stderr', 'both', True)
    if val not in hide_vals:
        err = "'hide' got {0!r} which is not in {1!r}"
        raise ValueError(err.format(val, hide_vals))
    if val in (None, False):
        hide = ()
    elif val in ('both', True):
        hide = ('out', 'err')
    elif val == 'stdout':
        hide = ('out',)
    elif val == 'stderr':
        hide = ('err',)
    else:
        hide = (val,)
    return hide

class Runner(object):
    """
    Partially-abstract core command-running API.

    This class is not usable by itself and must be subclassed to implement, at
    minimum, ``run_direct`` and ``run_pty``. An explanation of its methods
    follows:

    ``run`` is the primary API call which handles high level logic (echoing the
    commands locally, constructing a useful `Result` object, etc) and
    wraps/delegates to the below methods for actual execution.

    ``select_method`` takes ``pty`` and ``fallback`` kwargs and returns the
    current object's ``run_direct`` or ``run_pty`` method, depending on those
    kwargs' values and environmental cues/limitations. `Runner` itself has a
    useful default implementation of this method, but overriding may be
    sometimes useful.

    ``run_direct`` and ``run_pty`` are fully abstract in `Runner`; in
    subclasses, they should perform actual command execution, hooking directly
    into a subprocess' stdout/stderr pipes and returning those pipes' eventual
    full contents as distinct strings.

    ``run_direct``/``run_pty`` both have a signature of ``(self, command, warn,
    hide, encoding, out_stream, err_stream)`` (see `run` for semantics of
    these) and must return a 4-tuple of ``(stdout, stderr, exitcode,
    exception)`` (see `Result` for their meaning).

    ``run_pty`` differs from ``run_direct`` in that it should utilize a
    pseudo-terminal, and is expected to only return a useful ``stdout`` (with
    ``stderr`` usually set to the empty string.)

    For a subclass implementation example, see the source code for `.Local`.
    """
    def __init__(self, context):
        """
        Create a new runner with a handle on some `.Context`.

        :param context:
            a `.Context` instance, used to transmit default options and provide
            access to other contextualized information (e.g. a remote-oriented
            `.Runner` might want a `.Context` subclass holding info about
            hostnames and ports.)

            .. note::
                The `.Context` given to `.Runner` instances **must** contain
                default config values for the `.Runner` class in question. At a
                minimum, this means values for each of the default
                `.Runner.run` keyword arguments such as ``echo`` and ``warn``.

        :raises exceptions.ValueError:
            if not all expected default values are found in ``context``.
        """
        #: The `.Context` given to the same-named argument of `__init__`.
        self.context = context

    def run(self, command, **kwargs):
        """
        Execute ``command``, returning a `Result` object.

        .. note::
            All kwargs will default to the values found in this instance's
            `~.Runner.context` attribute, specifically in its configuration's
            ``run`` subtree (e.g. ``run.echo`` provides the default value for
            the ``echo`` keyword, etc). The base default values are described
            in the parameter list below.

        :param str command: The shell command to execute.

        :param bool warn:
            Whether to warn and continue, instead of raising `.Failure`, when
            the executed command exits with a nonzero status. Default:
            ``False``.

        :param hide:
            Allows the caller to disable ``run``'s default behavior of copying
            the subprocess' stdout and stderr to the controlling terminal.
            Specify ``hide='out'`` (or ``'stdout'``) to hide only the stdout
            stream, ``hide='err'`` (or ``'stderr'``) to hide only stderr, or
            ``hide='both'`` (or ``True``) to hide both streams.

            The default value is ``None``, meaning to print everything;
            ``False`` will also disable hiding.

            .. note::
                Stdout and stderr are always captured and stored in the
                ``Result`` object, regardless of ``hide``'s value.

        :param bool pty:
            By default, ``run`` connects directly to the invoked process and
            reads its stdout/stderr streams. Some programs will buffer (or even
            behave) differently in this situation compared to using an actual
            terminal or pty. To use a pty, specify ``pty=True``.

            .. warning::
                Due to their nature, ptys have a single output stream, so the
                ability to tell stdout apart from stderr is **not possible**
                when ``pty=True``. As such, all output will appear on
                ``out_stream`` (see below) and be captured into the ``stdout``
                result attribute. ``err_stream`` and ``stderr`` will always be
                empty when ``pty=True``.

        :param bool fallback:
            Controls auto-fallback behavior re: problems offering a pty when
            ``pty=True``. Whether this has any effect depends on the specific
            `Runner` subclass being invoked. Default: ``True``.

        :param bool echo:
            Controls whether `.run` prints the command string to local stdout
            prior to executing it. Default: ``False``.

        :param str encoding:
            Override auto-detection of which encoding the subprocess is using
            for its stdout/stderr streams. Defaults to the return value of
            ``locale.getpreferredencoding(False)``).

        :param out_stream:
            A file-like stream object to which the subprocess' standard error
            should be written. If ``None`` (the default), ``sys.stdout`` will
            be used.

        :param err_stream:
            Same as ``out_stream``, except for standard error, and defaulting
            to ``sys.stderr``.

        :returns: `Result`

        :raises: `.Failure` (if the command exited nonzero & ``warn=False``)
        """
        exception = False
        # Normalize kwargs w/ config
        opts = {}
        for key, value in six.iteritems(self.context.config.run):
            runtime = kwargs.pop(key, None)
            opts[key] = value if runtime is None else runtime
        # TODO: handle invalid kwarg keys (anything left in kwargs)
        # Normalize 'hide' from one of the various valid input values
        opts['hide'] = normalize_hide(opts['hide'])
        # Derive stream objects
        out_stream = sys.stdout if opts['out_stream'] is None else opts['out_stream']
        err_stream = sys.stderr if opts['err_stream'] is None else opts['err_stream']
        # Do the things
        if opts['echo']:
            print("\033[1;37m{0}\033[0m".format(command))
        func = self.select_method(pty=opts['pty'], fallback=opts['fallback'])
        stdout, stderr, exited, exception = func(
            command=command,
            warn=opts['warn'],
            hide=opts['hide'],
            encoding=opts['encoding'],
            out_stream=out_stream,
            err_stream=err_stream,
        )
        # TODO: make this test less gross? Feels silly to just return a bool in
        # select_method which is tantamount to this, though.
        func_name = getattr(func, 'func_name', getattr(func, '__name__'))
        used_pty = func_name == 'run_pty'
        result = Result(
            stdout=stdout,
            stderr=stderr,
            exited=exited,
            pty=used_pty,
            exception=exception,
        )
        if not (result or opts['warn']):
            raise Failure(result)
        return result

    def select_method(self, pty, fallback):
        # NOTE: fallback not used: no falling back implemented by default.
        return getattr(self, 'run_pty' if pty else 'run_direct')

    def run_direct(self, command, warn, hide, encoding, out_stream, err_stream):
        raise NotImplementedError

    def run_pty(self, command, warn, hide, encoding, out_stream, err_stream):
        raise NotImplementedError

class Local(Runner):
    """
    Execute a command on the local system in a subprocess.

    .. note::
        When Invoke itself is executed without a valid PTY (i.e.
        ``os.isatty(sys.stdin)`` is ``False``), it's not possible to present a
        handle on our PTY to local subprocesses. In such situations, `Local`
        will fallback to behaving as if ``pty=False``, on the theory that
        degraded execution is better than none at all, as well as printing a
        warning to stderr.

        To disable this behavior (i.e. if ``os.isatty`` is causing false
        negatives in your environment), say ``fallback=False``.
    """
    def select_method(self, pty=False, fallback=True):
        func = self.run_direct
        if pty:
            func = self.run_pty
            if not os.isatty(sys.stdin.fileno()) and fallback:
                sys.stderr.write("WARNING: stdin is not a pty; falling back to non-pty execution!\n") # noqa
                func = self.run_direct
        return func

    def run_direct(self, command, warn, hide, encoding, out_stream, err_stream):
        process = Popen(
            command,
            shell=True,
            stdout=PIPE,
            stderr=PIPE,
        )

        if encoding is None:
            encoding = locale.getpreferredencoding(False)

        def display(src, dst, cap, hide):
            def get():
                while True:
                    data = os.read(src.fileno(), 1000)
                    if not data:
                        break
                    yield data
            for data in codecs.iterdecode(get(), encoding, errors='replace'):
                if not hide:
                    dst.write(data)
                    dst.flush()
                cap.append(data)

        stdout = []
        stderr = []
        threads = []

        for args in (
            (process.stdout, out_stream, stdout, 'out' in hide),
            (process.stderr, err_stream, stderr, 'err' in hide),
        ):
            t = threading.Thread(target=display, args=args)
            threads.append(t)
            t.start()

        process.wait()
        for t in threads:
            t.join()

        stdout = ''.join(stdout)
        stderr = ''.join(stderr)
        if WINDOWS:
            # "Universal newlines" - replace all standard forms of
            # newline with \n. This is not technically Windows related
            # (\r as newline is an old Mac convention) but we only apply
            # the translation for Windows as that's the only platform
            # it is likely to matter for these days.
            stdout = stdout.replace("\r\n", "\n").replace("\r", "\n")
            stderr = stderr.replace("\r\n", "\n").replace("\r", "\n")

        return stdout, stderr, process.returncode, None

    def run_pty(self, command, warn, hide, encoding, out_stream, err_stream):
        # TODO: re-insert Windows "lol y u no pty" stuff here
        import pty
        pid, parent_fd = pty.fork()
        # If we're the child process, load up the actual command in a shell,
        # just as subprocess does; this replaces our process - whose pipes are
        # all hooked up to the PTY - with the "real" one.
        if pid == 0:
            # Use execv for bare-minimum "exec w/ variable # args" behavior.
            # No need for the 'p' (use PATH to find executable) or 'e' (define
            # a custom/overridden shell env) variants, for now.
            # TODO: use /bin/sh or whatever subprocess does. Only using bash
            # for now because that's what we have been testing against.
            # TODO: also see if subprocess is using equivalent of execvp...
            # TODO: both pty.spawn() and pexpect.spawn() do a lot of
            # setup/teardown involving tty.*, setwinsize, getrlimit, signal.
            # Ostensibly we'll want some of that eventually, but if possible
            # write tests - integration-level if necessary - before adding it!
            os.execv('/bin/bash', ['/bin/bash', '-c', command])

        if encoding is None:
            encoding = locale.getpreferredencoding(False)

        def display(src_fd, dst, cap, hide):
            def get():
                while True:
                    data = os.read(src_fd, 1000)
                    if not data:
                        break
                    yield data
            for data in codecs.iterdecode(get(), encoding, errors='replace'):
                if not hide:
                    dst.write(data)
                    dst.flush()
                cap.append(data)

        stdout = []
        stderr = []
        threads = []

        # TODO: can we simply run this loop in the main thread now?
        # TODO: or is there some other benefit to threading it, such as
        # eventual stdin control?
        for args in (
            (parent_fd, out_stream, stdout, 'out' in hide),
            #(process.stderr, err_stream, stderr, 'err' in hide),
        ):
            t = threading.Thread(target=display, args=args)
            threads.append(t)
            t.start()

        # Wait in main thread until child appears to have exited.
        while True:
            # TODO: set 2nd value to os.WNOHANG in some situations?
            pid_val, status = os.waitpid(pid, 0)
            # waitpid() sets the 'pid' return val to 0 when no children have
            # exited yet; when it is NOT zero, we know the child's stopped.
            if pid_val != 0:
                break
            # TODO: io sleep?

        # Derive return code from waitpid() result
        returncode = os.WEXITSTATUS(status)

        # Join reader threads & pack things up
        for t in threads:
            t.join()

        stdout = ''.join(stdout)
        stderr = ''.join(stderr)
        if WINDOWS:
            # "Universal newlines" - replace all standard forms of
            # newline with \n. This is not technically Windows related
            # (\r as newline is an old Mac convention) but we only apply
            # the translation for Windows as that's the only platform
            # it is likely to matter for these days.
            stdout = stdout.replace("\r\n", "\n").replace("\r", "\n")
            stderr = stderr.replace("\r\n", "\n").replace("\r", "\n")

        return stdout, stderr, returncode, None

    def _run_pty(self, command, warn, hide, encoding, out_stream, err_stream):
        from invoke import pty
        e = pty.spawn(['/bin/bash', '-c', command])
        print e

        return "uh", "wat", 1, None

        # Sanity check: platforms that can't pexpect should explode usefully
        # here. (Without this, the pexpect import throws an inner
        # ImportException trying to 'import pty' which is unavailable on
        # Windows. Better to do this here than truly-fork pexpect.)
        if WINDOWS:
            err = "You seem to be on Windows, which doesn't support ptys!"
            raise PlatformError(err)
        # Proceed as normal for POSIX/etc platforms, with a runtime import
        from .vendor import pexpect

        out = []
        def out_filter(text):
            # Maybe use encoding here as in run() above...
            out.append(text.decode("utf-8", 'replace'))
            if 'out' not in hide:
                return text
            else:
                return b""
        # Hand pexpect useful args for os.execve(binary, args).
        # TODO: try to use same default binary value as Popen does (is usually
        # /bin/sh but IIRC there's caveats)
        # TODO: allow control over this
        p = pexpect.spawn("/bin/bash", ["-c", command])
        # Ensure pexpect doesn't barf with OSError if we fall off the end of
        # the child's input on some platforms (e.g. Linux).
        exception = None
        try:
            p.interact(output_filter=out_filter)
        except OSError as e:
            # Only capture the OSError we expect
            if "Input/output error" not in str(e):
                raise
            # Ensure it ties off the child, sets exitstatus, etc
            p.close()
            # Capture the exception in case it's NOT the OSError we think it
            # is and folks need to debug
            exception = e
        return "".join(out), "", p.exitstatus, exception

class Result(object):
    """
    A container for information about the result of a command execution.

    `Result` instances have the following attributes:

    * ``stdout``: The subprocess' standard output, as a multiline string.
    * ``stderr``: Same as ``stdout`` but containing standard error (unless
      the process was invoked via a pty; see `.Runner.run`.)
    * ``exited``: An integer representing the subprocess' exit/return code.
    * ``return_code``: An alias to ``exited``.
    * ``ok``: A boolean equivalent to ``exited == 0``.
    * ``failed``: The inverse of ``ok``: ``True`` if the program exited with a
      nonzero return code.
    * ``pty``: A boolean describing whether the subprocess was invoked with a
      pty or not; see `.Runner.run`.
    * ``exception``: Typically ``None``, but may be an exception object if
      ``pty`` was ``True`` and ``run`` had to swallow an apparently-spurious
      ``OSError``. Solely for sanity checking/debugging purposes.

    `Result` objects' truth evaluation is equivalent to their ``ok``
    attribute's value.
    """
    # TODO: inherit from namedtuple instead? heh
    def __init__(self, stdout, stderr, exited, pty, exception=None):
        self.exited = self.return_code = exited
        self.stdout = stdout
        self.stderr = stderr
        self.pty = pty
        self.exception = exception

    def __nonzero__(self):
        # Holy mismatch between name and implementation, Batman!
        return self.exited == 0

    # Python 3 ahoy
    def __bool__(self):
        return self.__nonzero__()

    def __str__(self):
        ret = ["Command exited with status {0}.".format(self.exited)]
        for x in ('stdout', 'stderr'):
            val = getattr(self, x)
            ret.append("""=== {0} ===
{1}
""".format(x, val.rstrip()) if val else "(no {0})".format(x))
        return "\n".join(ret)

    @property
    def ok(self):
        return self.exited == 0

    @property
    def failed(self):
        return not self.ok
<|endoftext|>


<commit_before># Copyright 2016 VMware, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Implements the interface for side car based implementation
of a KV store for vmdks.
"""

from ctypes import \
        CDLL, POINTER, byref, Structure,\
        c_char_p, c_int32, c_bool, c_uint32, c_uint64
import json
import logging
import sys
import errno
import time
import threadutils

# Python version 3.5.1
PYTHON64_VERSION = 50659824

# Conversion constants
KB = 1024
MB = (1024 * 1024)
GB = (MB * 1024)

# Side car create/open options
KV_SIDECAR_CREATE = 0

# Start size for a side car
KV_CREATE_SIZE = 0

# VSphere lib to access ESX proprietary APIs.
DISK_LIB64 = "/lib64/libvmsnapshot.so"
DISK_LIB = "/lib/libvmsnapshot.so"
lib = None
use_sidecar_create = False
DVOL_KEY = "docker-volume-vsphere"

# Volume attributes
VOL_SIZE = 'size'
VOL_ALLOC = 'allocated'

# Results in a buffered, locked, filter-less open,
# all vmdks are opened with these flags
VMDK_OPEN_DEFAULT = 524312
VMDK_OPEN_NOIO = 524289

# Default kv side car alignment
KV_ALIGN = 4096

# Flag to track the version of Python on the platform
is_64bits = False

# Number of times and sleep time to retry on IOError EBUSY
EBUSY_RETRY_COUNT = 8
EBUSY_RETRY_SLEEP = 0.5

# Get a Reentrant locking decorator
diskLibLock = threadutils.get_lock_decorator(reentrant=True)

class disk_info(Structure):
    _fields_ = [('size', c_uint64),
                ('allocated', c_uint64),
                ('fill1', c_uint64),
                ('fill2', c_uint64),
                ('fill3', c_uint32),
                ('fill4', c_uint32)]

@diskLibLock
def load_disk_lib(lib_name):
    """
    Load the disk lib API library
    """
    global lib
    if not lib:
        lib = CDLL(lib_name)
        lib.DiskLib_Init.argtypes = []
        lib.DiskLib_Init.restype = c_bool
        lib.DiskLib_Init()
    return

@diskLibLock
def disk_lib_init():
    """
    Initialize disk lib interfaces
    """
    global is_64bits
    global use_sidecar_create

    # Define arg types for disk lib apis.
    if sys.hexversion >= PYTHON64_VERSION:
        is_64bits = True
        load_disk_lib(DISK_LIB64)
        lib.DiskLib_OpenWithInfo.argtypes = [c_char_p, c_int32,
                                             POINTER(c_uint32),
                                             POINTER(c_uint64),
                                             POINTER(c_uint64)]
        lib.DiskLib_Close.argtypes = [c_uint64]

        lib.DiskLib_SidecarOpen.argtypes = [c_uint64, c_char_p, c_int32,
                                            POINTER(c_uint64)]
        lib.DiskLib_SidecarClose.argtypes = [c_uint64, c_char_p,
                                             POINTER(c_uint64)]
        lib.DiskLib_DBGet.argtypes = [c_uint64, c_char_p, POINTER(c_char_p)]
        lib.DiskLib_DBSet.argtypes = [c_uint64, c_char_p, c_char_p]
        lib.DiskLib_GetSize.argtypes = [c_uint64, c_uint32,
                                        c_uint32, POINTER(disk_info)]
        # Check if this library supports create API
        try:
            lib.DiskLib_SidecarCreate.argtypes = [c_uint64, c_char_p, c_uint64,
                                                  c_int32, POINTER(c_uint64)]
            lib.DiskLib_SidecarCreate.restype = int
            use_sidecar_create = True
        except:
            logging.debug(
                "ESX version doesn't support create API, using open instead.")
            pass
    else:
        load_disk_lib(DISK_LIB)
        lib.DiskLib_OpenWithInfo.argtypes = [c_char_p, c_int32,
                                             POINTER(c_uint32),
                                             POINTER(c_uint32),
                                             POINTER(c_uint32)]
        lib.DiskLib_Close.argtypes = [c_uint32]
        lib.DiskLib_SidecarOpen.argtypes = [c_uint32, c_char_p, c_int32,
                                            POINTER(c_uint32)]
        lib.DiskLib_SidecarClose.argtypes = [c_uint32, c_char_p,
                                             POINTER(c_uint32)]
        lib.DiskLib_DBGet.argtypes = [c_uint32, c_char_p, POINTER(c_char_p)]
        lib.DiskLib_DBSet.argtypes = [c_uint32, c_char_p, c_char_p]
        lib.DiskLib_GetSize.argtypes = [c_uint32, c_uint32,
                                        c_uint32, POINTER(disk_info)]

       # Check if this library supports create API
        try:
            lib.DiskLib_SidecarCreate.argtypes = [c_uint32, c_char_p, c_uint64,
                                                  c_int32, POINTER(c_uint32)]
            lib.DiskLib_SidecarCreate.restype = int
            use_sidecar_create = True
        except:
            logging.debug(
                "ESX version doesn't support create API, using open instead.")

    lib.DiskLib_SidecarMakeFileName.argtypes = [c_char_p, c_char_p]

    # Define result types for disk lib apis.
    lib.DiskLib_OpenWithInfo.restype = int
    lib.DiskLib_Close.restype = int
    lib.DiskLib_SidecarOpen.restype = int
    lib.DiskLib_SidecarClose.restype = int
    lib.DiskLib_SidecarMakeFileName.restype = c_char_p
    lib.DiskLib_DBGet.restype = c_uint32
    lib.DiskLib_DBSet.restype = c_uint32
    lib.DiskLib_GetSize.restype = c_uint32

    return

def kv_esx_init():
    """
    Load and init disklib
    """
    disk_lib_init()

@diskLibLock
def vol_open_path(volpath, open_flags=VMDK_OPEN_DEFAULT):
    """
    Open a VMDK given its path, the VMDK is opened locked just to
    ensure we have exclusive access and its not already in use.
    """
    dhandle = get_uint(0)
    ihandle = get_uint(0)
    key = c_uint32(0)

    res = lib.DiskLib_OpenWithInfo(volpath.encode(), open_flags,
                                   byref(key), byref(dhandle),
                                   byref(ihandle))
    if res != 0:
        logging.warning("Open %s failed - %x", volpath, res)
    return dhandle

@diskLibLock
def create(volpath, kv_dict):
    """
    Create the side car for the volume identified by volpath.
    """
    obj_handle = get_uint(0)
    dhandle = vol_open_path(volpath)

    if not disk_is_valid(dhandle):
        return False
    if use_sidecar_create:
        res = lib.DiskLib_SidecarCreate(dhandle, DVOL_KEY.encode(),
                                        KV_CREATE_SIZE, KV_SIDECAR_CREATE,
                                        byref(obj_handle))
    else:
        res = lib.DiskLib_SidecarOpen(dhandle, DVOL_KEY.encode(),
                                      KV_SIDECAR_CREATE,
                                      byref(obj_handle))
    if res != 0:
        logging.warning("Side car create for %s failed - %x", volpath, res)
        lib.DiskLib_Close(dhandle)
        return False

    lib.DiskLib_SidecarClose(dhandle, DVOL_KEY.encode(), byref(obj_handle))
    lib.DiskLib_Close(dhandle)

    return save(volpath, kv_dict)

@diskLibLock
def delete(volpath):
    """
    Delete the the side car for the given volume
    """
    dhandle = vol_open_path(volpath)
    if not disk_is_valid(dhandle):
        return False
    res = lib.DiskLib_SidecarDelete(dhandle, DVOL_KEY.encode())
    if res != 0:
        logging.warning("Side car delete for %s failed - %x", volpath, res)
        lib.DiskLib_Close(dhandle)
        return False

    lib.DiskLib_Close(dhandle)
    return True

def align_str(kv_str, block):
    """
    Align a given string to the specified block boundary.
    """
    # Align string to the next block boundary. The -1 is to accommodate
    # a newline at the end of the string.
    aligned_len = int((len(kv_str) + block - 1) / block) * block - 1
    return '{:<{width}}\n'.format(kv_str, width=aligned_len)

@diskLibLock
def load(volpath):
    """
    Load and return dictionary from the sidecar
    """
    meta_file = lib.DiskLib_SidecarMakeFileName(volpath.encode(),
                                                DVOL_KEY.encode())
    retry_count = 0
    while True:
        try:
            with open(meta_file, "r") as fh:
                kv_str = fh.read()
            break
        except IOError as open_error:
            # This is a workaround to the timing/locking with metadata files issue #626
            if (hasattr(open_error, "errno") and
                    (open_error.errno == errno.EBUSY and retry_count <= EBUSY_RETRY_COUNT)):
                logging.warning("Meta file %s busy for load(), retrying...", meta_file)
                retry_count += 1
                time.sleep(EBUSY_RETRY_SLEEP)
            else:
                logging.exception("Failed to access %s", meta_file)
                return None

    try:
        return json.loads(kv_str)
    except ValueError:
        logging.exception("Failed to decode meta-data for %s", volpath)
        return None

@diskLibLock
def save(volpath, kv_dict):
    """
    Save the dictionary to side car.
    """
    meta_file = lib.DiskLib_SidecarMakeFileName(volpath.encode(),
                                                DVOL_KEY.encode())
    kv_str = json.dumps(kv_dict)

    retry_count = 0
    while True:
        try:
            with open(meta_file, "w") as fh:
                fh.write(align_str(kv_str, KV_ALIGN))
            break
        except IOError as open_error:
            # This is a workaround to the timing/locking with metadata files issue #626
            if (hasattr(open_error, "errno") and
                    (open_error.errno == errno.EBUSY and retry_count <= EBUSY_RETRY_COUNT)):
                logging.warning("Meta file %s busy for save(), retrying...", meta_file)
                retry_count += 1
                time.sleep(EBUSY_RETRY_SLEEP)
            else:
                logging.exception("Failed to save meta-data for %s", volpath)
                return False

    return True

@diskLibLock
def get_info(volpath):
    """
    Return disk stats for the volume
    """
    dhandle = vol_open_path(volpath, VMDK_OPEN_NOIO)

    if not disk_is_valid(dhandle):
        logging.warning("Failed to open disk - %x", volpath)
        return None

    sinfo = disk_info()
    res = lib.DiskLib_GetSize(dhandle, 0, 1, byref(sinfo))

    lib.DiskLib_Close(dhandle)
    if res != 0:
        logging.warning("Failed to get size of disk %s - %x", volpath, res)
        return None

    return {VOL_SIZE: convert(sinfo.size), VOL_ALLOC: convert(sinfo.allocated)}

def get_uint(val):
    """
    Return a C unsigned int
    """
    if is_64bits:
        return c_uint64(val)
    else:
        return c_uint32(val)

def disk_is_valid(dhandle):
    """
    Check if disk is valid
    """
    if is_64bits:
        return dhandle.value != c_uint64(0).value
    else:
        return dhandle.value != c_uint32(0).value

def convert(size):
    """
    Size conversions
    """
    if size < KB:
        return size
    elif size < MB:
        return '{0}{1}'.format(size / KB, 'KB')
    elif size < GB:
        return '{0}{1}'.format(size / MB, 'MB')
    else:
        return '{0}{1}'.format(size / GB, 'GB')
<commit_msg>Calculate allocated capacity and size on Get() across snaps of a volume (#723)<commit_after># Copyright 2016 VMware, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Implements the interface for side car based implementation
of a KV store for vmdks.
"""

from ctypes import \
        CDLL, POINTER, byref, Structure,\
        c_char_p, c_int32, c_bool, c_uint32, c_uint64
import json
import logging
import sys
import errno
import time
import threadutils

# Python version 3.5.1
PYTHON64_VERSION = 50659824

# Conversion constants
KB = 1024
MB = (1024 * 1024)
GB = (MB * 1024)

# Side car create/open options
KV_SIDECAR_CREATE = 0

# Start size for a side car
KV_CREATE_SIZE = 0

# VSphere lib to access ESX proprietary APIs.
DISK_LIB64 = "/lib64/libvmsnapshot.so"
DISK_LIB = "/lib/libvmsnapshot.so"
lib = None
use_sidecar_create = False
DVOL_KEY = "docker-volume-vsphere"

# Volume attributes
VOL_SIZE = 'size'
VOL_ALLOC = 'allocated'

# Results in a buffered, locked, filter-less open,
# all vmdks are opened with these flags
VMDK_OPEN_DEFAULT = 524312
VMDK_OPEN_DISKCHAIN_NOIO = 524291
VMDK_MAX_SNAPS = 32

# Default kv side car alignment
KV_ALIGN = 4096

# Flag to track the version of Python on the platform
is_64bits = False

# Number of times and sleep time to retry on IOError EBUSY
EBUSY_RETRY_COUNT = 8
EBUSY_RETRY_SLEEP = 0.5

# Get a Reentrant locking decorator
diskLibLock = threadutils.get_lock_decorator(reentrant=True)

class disk_info(Structure):
    _fields_ = [('size', c_uint64),
                ('allocated', c_uint64),
                ('fill1', c_uint64),
                ('fill2', c_uint64),
                ('fill3', c_uint32),
                ('fill4', c_uint32)]

@diskLibLock
def load_disk_lib(lib_name):
    """
    Load the disk lib API library
    """
    global lib
    if not lib:
        lib = CDLL(lib_name)
        lib.DiskLib_Init.argtypes = []
        lib.DiskLib_Init.restype = c_bool
        lib.DiskLib_Init()
    return

@diskLibLock
def disk_lib_init():
    """
    Initialize disk lib interfaces
    """
    global is_64bits
    global use_sidecar_create

    # Define arg types for disk lib apis.
    if sys.hexversion >= PYTHON64_VERSION:
        is_64bits = True
        load_disk_lib(DISK_LIB64)
        lib.DiskLib_OpenWithInfo.argtypes = [c_char_p, c_int32,
                                             POINTER(c_uint32),
                                             POINTER(c_uint64),
                                             POINTER(c_uint64)]
        lib.DiskLib_Close.argtypes = [c_uint64]

        lib.DiskLib_SidecarOpen.argtypes = [c_uint64, c_char_p, c_int32,
                                            POINTER(c_uint64)]
        lib.DiskLib_SidecarClose.argtypes = [c_uint64, c_char_p,
                                             POINTER(c_uint64)]
        lib.DiskLib_DBGet.argtypes = [c_uint64, c_char_p, POINTER(c_char_p)]
        lib.DiskLib_DBSet.argtypes = [c_uint64, c_char_p, c_char_p]
        lib.DiskLib_GetSize.argtypes = [c_uint64, c_uint32,
                                        c_uint32, POINTER(disk_info)]
        # Check if this library supports create API
        try:
            lib.DiskLib_SidecarCreate.argtypes = [c_uint64, c_char_p, c_uint64,
                                                  c_int32, POINTER(c_uint64)]
            lib.DiskLib_SidecarCreate.restype = int
            use_sidecar_create = True
        except:
            logging.debug(
                "ESX version doesn't support create API, using open instead.")
            pass
    else:
        load_disk_lib(DISK_LIB)
        lib.DiskLib_OpenWithInfo.argtypes = [c_char_p, c_int32,
                                             POINTER(c_uint32),
                                             POINTER(c_uint32),
                                             POINTER(c_uint32)]
        lib.DiskLib_Close.argtypes = [c_uint32]
        lib.DiskLib_SidecarOpen.argtypes = [c_uint32, c_char_p, c_int32,
                                            POINTER(c_uint32)]
        lib.DiskLib_SidecarClose.argtypes = [c_uint32, c_char_p,
                                             POINTER(c_uint32)]
        lib.DiskLib_DBGet.argtypes = [c_uint32, c_char_p, POINTER(c_char_p)]
        lib.DiskLib_DBSet.argtypes = [c_uint32, c_char_p, c_char_p]
        lib.DiskLib_GetSize.argtypes = [c_uint32, c_uint32,
                                        c_uint32, POINTER(disk_info)]

       # Check if this library supports create API
        try:
            lib.DiskLib_SidecarCreate.argtypes = [c_uint32, c_char_p, c_uint64,
                                                  c_int32, POINTER(c_uint32)]
            lib.DiskLib_SidecarCreate.restype = int
            use_sidecar_create = True
        except:
            logging.debug(
                "ESX version doesn't support create API, using open instead.")

    lib.DiskLib_SidecarMakeFileName.argtypes = [c_char_p, c_char_p]

    # Define result types for disk lib apis.
    lib.DiskLib_OpenWithInfo.restype = int
    lib.DiskLib_Close.restype = int
    lib.DiskLib_SidecarOpen.restype = int
    lib.DiskLib_SidecarClose.restype = int
    lib.DiskLib_SidecarMakeFileName.restype = c_char_p
    lib.DiskLib_DBGet.restype = c_uint32
    lib.DiskLib_DBSet.restype = c_uint32
    lib.DiskLib_GetSize.restype = c_uint32

    return

def kv_esx_init():
    """
    Load and init disklib
    """
    disk_lib_init()

@diskLibLock
def vol_open_path(volpath, open_flags=VMDK_OPEN_DEFAULT):
    """
    Open a VMDK given its path, the VMDK is opened locked just to
    ensure we have exclusive access and its not already in use.
    """
    dhandle = get_uint(0)
    ihandle = get_uint(0)
    key = c_uint32(0)

    res = lib.DiskLib_OpenWithInfo(volpath.encode(), open_flags,
                                   byref(key), byref(dhandle),
                                   byref(ihandle))
    if res != 0:
        logging.warning("Open %s failed - %x", volpath, res)
    return dhandle

@diskLibLock
def create(volpath, kv_dict):
    """
    Create the side car for the volume identified by volpath.
    """
    obj_handle = get_uint(0)
    dhandle = vol_open_path(volpath)

    if not disk_is_valid(dhandle):
        return False
    if use_sidecar_create:
        res = lib.DiskLib_SidecarCreate(dhandle, DVOL_KEY.encode(),
                                        KV_CREATE_SIZE, KV_SIDECAR_CREATE,
                                        byref(obj_handle))
    else:
        res = lib.DiskLib_SidecarOpen(dhandle, DVOL_KEY.encode(),
                                      KV_SIDECAR_CREATE,
                                      byref(obj_handle))
    if res != 0:
        logging.warning("Side car create for %s failed - %x", volpath, res)
        lib.DiskLib_Close(dhandle)
        return False

    lib.DiskLib_SidecarClose(dhandle, DVOL_KEY.encode(), byref(obj_handle))
    lib.DiskLib_Close(dhandle)

    return save(volpath, kv_dict)

@diskLibLock
def delete(volpath):
    """
    Delete the the side car for the given volume
    """
    dhandle = vol_open_path(volpath)
    if not disk_is_valid(dhandle):
        return False
    res = lib.DiskLib_SidecarDelete(dhandle, DVOL_KEY.encode())
    if res != 0:
        logging.warning("Side car delete for %s failed - %x", volpath, res)
        lib.DiskLib_Close(dhandle)
        return False

    lib.DiskLib_Close(dhandle)
    return True

def align_str(kv_str, block):
    """
    Align a given string to the specified block boundary.
    """
    # Align string to the next block boundary. The -1 is to accommodate
    # a newline at the end of the string.
    aligned_len = int((len(kv_str) + block - 1) / block) * block - 1
    return '{:<{width}}\n'.format(kv_str, width=aligned_len)

@diskLibLock
def load(volpath):
    """
    Load and return dictionary from the sidecar
    """
    meta_file = lib.DiskLib_SidecarMakeFileName(volpath.encode(),
                                                DVOL_KEY.encode())
    retry_count = 0
    while True:
        try:
            with open(meta_file, "r") as fh:
                kv_str = fh.read()
            break
        except IOError as open_error:
            # This is a workaround to the timing/locking with metadata files issue #626
            if (hasattr(open_error, "errno") and
                    (open_error.errno == errno.EBUSY and retry_count <= EBUSY_RETRY_COUNT)):
                logging.warning("Meta file %s busy for load(), retrying...", meta_file)
                retry_count += 1
                time.sleep(EBUSY_RETRY_SLEEP)
            else:
                logging.exception("Failed to access %s", meta_file)
                return None

    try:
        return json.loads(kv_str)
    except ValueError:
        logging.exception("Failed to decode meta-data for %s", volpath)
        return None

@diskLibLock
def save(volpath, kv_dict):
    """
    Save the dictionary to side car.
    """
    meta_file = lib.DiskLib_SidecarMakeFileName(volpath.encode(),
                                                DVOL_KEY.encode())
    kv_str = json.dumps(kv_dict)

    retry_count = 0
    while True:
        try:
            with open(meta_file, "w") as fh:
                fh.write(align_str(kv_str, KV_ALIGN))
            break
        except IOError as open_error:
            # This is a workaround to the timing/locking with metadata files issue #626
            if (hasattr(open_error, "errno") and
                    (open_error.errno == errno.EBUSY and retry_count <= EBUSY_RETRY_COUNT)):
                logging.warning("Meta file %s busy for save(), retrying...", meta_file)
                retry_count += 1
                time.sleep(EBUSY_RETRY_SLEEP)
            else:
                logging.exception("Failed to save meta-data for %s", volpath)
                return False

    return True

@diskLibLock
def get_info(volpath):
    """
    Return disk stats for the volume
    """
    dhandle = vol_open_path(volpath, VMDK_OPEN_DISKCHAIN_NOIO)

    if not disk_is_valid(dhandle):
        logging.warning("Failed to open disk - %x", volpath)
        return None

    sinfo = disk_info()
    res = lib.DiskLib_GetSize(dhandle, 0, VMDK_MAX_SNAPS, byref(sinfo))

    lib.DiskLib_Close(dhandle)
    if res != 0:
        logging.warning("Failed to get size of disk %s - %x", volpath, res)
        return None

    return {VOL_SIZE: convert(sinfo.size), VOL_ALLOC: convert(sinfo.allocated)}

def get_uint(val):
    """
    Return a C unsigned int
    """
    if is_64bits:
        return c_uint64(val)
    else:
        return c_uint32(val)

def disk_is_valid(dhandle):
    """
    Check if disk is valid
    """
    if is_64bits:
        return dhandle.value != c_uint64(0).value
    else:
        return dhandle.value != c_uint32(0).value

def convert(size):
    """
    Size conversions
    """
    if size < KB:
        return size
    elif size < MB:
        return '{0}{1}'.format(size / KB, 'KB')
    elif size < GB:
        return '{0}{1}'.format(size / MB, 'MB')
    else:
        return '{0}{1}'.format(size / GB, 'GB')
<|endoftext|>


<commit_before>#!/usr/bin/env python
'''
CREATED:2014-01-17 16:30:07 by Brian McFee <brm2132@columbia.edu>

Compute segmentation evaluation metrics

Usage:

./segment_eval.py TRUTH.TXT PREDICTION.TXT
'''

import argparse
import sys
import os
from collections import OrderedDict

import mir_eval

def evaluate(truth_file=None, prediction_file=None):
    '''Load data and perform the evaluation'''

    # load the data
    truth_segs, truth_labels    = mir_eval.io.load_annotation(truth_file)
    pred_segs, pred_labels      = mir_eval.io.load_annotation(prediction_file)

    # Convert to boundaries
    annotated_boundaries        = mir_eval.util.segments_to_boundaries(truth_segs)[0]
    predicted_boundaries        = mir_eval.util.segments_to_boundaries(pred_segs)[0]

    # Adjust the predictions to match the annotation time span (0 to end-of-track)

    # The [0] at the end is to just get the boundaries without the labels
    predicted_boundaries        = mir_eval.util.adjust_times(predicted_boundaries,
                                                        t_min=annotated_boundaries.min(),
                                                        t_max=annotated_boundaries.max())[0]

    # Now compute all the metrics

    M = OrderedDict()
    # Boundary detection
    M['P@0.5'], M['R@0.5'], M['F@0.5']  = mir_eval.segment.boundary_detection(annotated_boundaries,
                                                                  predicted_boundaries,
                                                                  window=0.5)

    M['P@3.0'], M['R@3.0'], M['F@3.0']  = mir_eval.segment.boundary_detection(annotated_boundaries,
                                                                 predicted_boundaries,
                                                                 window=3.0)
    # Boundary deviation
    M['True-to-Pred'], M['Pred-to-True'] = mir_eval.segment.boundary_deviation(annotated_boundaries,
                                                                  predicted_boundaries)

    # Pairwise clustering
    M['Pair-P'], M['Pair-R'], M['Pair-F'] = mir_eval.segment.frame_clustering_pairwise(annotated_boundaries,
                                                                         predicted_boundaries)

    # Adjusted rand index
    M['ARI']                     = mir_eval.segment.frame_clustering_ari(annotated_boundaries,
                                                                    predicted_boundaries)

    # Mutual information metrics
    M['MI'], M['AMI'], M['NMI']            = mir_eval.segment.frame_clustering_mi(annotated_boundaries,
                                                                   predicted_boundaries)

    # Conditional entropy metrics
    M['S_Over'], M['S_Under'], M['S_F'] = mir_eval.segment.frame_clustering_nce(annotated_boundaries,
                                                                    predicted_boundaries)

    # And print them
    print os.path.basename(prediction_file)
    for key, value in M.iteritems():
        print '%20s:\t%0.3f' % (key, value)

    pass

def process_arguments():
    '''Argparse function to get the program parameters'''

    parser = argparse.ArgumentParser(description='mir_eval segmentation evaluation')

    parser.add_argument(    'truth_file',
                            action      =   'store',
                            help        =   'path to the ground truth annotation')

    parser.add_argument(    'prediction_file',
                            action      =   'store',
                            help        =   'path to the prediction file')

    return vars(parser.parse_args(sys.argv[1:]))

if __name__ == '__main__':
    # Get the parameters
    parameters = process_arguments()

    # Run the beat tracker
    evaluate(**parameters)
<commit_msg>format cleanup<commit_after>#!/usr/bin/env python
'''
CREATED:2014-01-17 16:30:07 by Brian McFee <brm2132@columbia.edu>

Compute segmentation evaluation metrics

Usage:

./segment_eval.py TRUTH.TXT PREDICTION.TXT
'''

import argparse
import sys
import os
from collections import OrderedDict

import mir_eval

def evaluate(truth_file=None, prediction_file=None):
    '''Load data and perform the evaluation'''

    # load the data
    truth_segs, truth_labels    = mir_eval.io.load_annotation(truth_file)
    pred_segs, pred_labels      = mir_eval.io.load_annotation(prediction_file)

    # Convert to boundaries
    annotated_boundaries        = mir_eval.util.segments_to_boundaries(truth_segs)[0]
    predicted_boundaries        = mir_eval.util.segments_to_boundaries(pred_segs)[0]

    # Adjust the predictions to match the annotation time span (0 to end-of-track)

    # The [0] at the end is to just get the boundaries without the labels
    predicted_boundaries        = mir_eval.util.adjust_times(predicted_boundaries,
                                                        t_min=annotated_boundaries.min(),
                                                        t_max=annotated_boundaries.max())[0]

    # Now compute all the metrics

    M = OrderedDict()
    # Boundary detection
    M['P@0.5'], M['R@0.5'], M['F@0.5']  = mir_eval.segment.boundary_detection(annotated_boundaries,
                                                                  predicted_boundaries,
                                                                  window=0.5)

    M['P@3.0'], M['R@3.0'], M['F@3.0']  = mir_eval.segment.boundary_detection(annotated_boundaries,
                                                                 predicted_boundaries,
                                                                 window=3.0)
    # Boundary deviation
    M['True-to-Pred'], M['Pred-to-True'] = mir_eval.segment.boundary_deviation(annotated_boundaries,
                                                                  predicted_boundaries)

    # Pairwise clustering
    M['Pair-P'], M['Pair-R'], M['Pair-F'] = mir_eval.segment.frame_clustering_pairwise(annotated_boundaries,
                                                                         predicted_boundaries)

    # Adjusted rand index
    M['ARI']                     = mir_eval.segment.frame_clustering_ari(annotated_boundaries,
                                                                    predicted_boundaries)

    # Mutual information metrics
    M['MI'], M['AMI'], M['NMI']            = mir_eval.segment.frame_clustering_mi(annotated_boundaries,
                                                                   predicted_boundaries)

    # Conditional entropy metrics
    M['S_Over'], M['S_Under'], M['S_F'] = mir_eval.segment.frame_clustering_nce(annotated_boundaries,
                                                                    predicted_boundaries)

    # And print them
    print os.path.basename(prediction_file)
    for key, value in M.iteritems():
        print '\t%12s:\t%0.3f' % (key, value)

    pass

def process_arguments():
    '''Argparse function to get the program parameters'''

    parser = argparse.ArgumentParser(description='mir_eval segmentation evaluation')

    parser.add_argument(    'truth_file',
                            action      =   'store',
                            help        =   'path to the ground truth annotation')

    parser.add_argument(    'prediction_file',
                            action      =   'store',
                            help        =   'path to the prediction file')

    return vars(parser.parse_args(sys.argv[1:]))

if __name__ == '__main__':
    # Get the parameters
    parameters = process_arguments()

    # Run the beat tracker
    evaluate(**parameters)
<|endoftext|>