coalib.misc.ContextManagers — coala 0.9.0.dev99999999999999 documentation

# Source code for coalib.misc.ContextManagers

    import builtins
    import os
    import platform
    import signal
    import sys
    import tempfile
    import threading
    from contextlib import closing, contextmanager
    from io import StringIO

    from coalib.misc.MutableValue import MutableValue

    @contextmanager
    def subprocess_timeout(sub_process, seconds, kill_pg=False):
        """
        Kill subprocess if the sub process takes more the than the timeout.

        :param sub_process: The sub process to run.
        :param seconds:     The number of seconds to allow the test to run for. If
                            set to 0 or a negative value, it waits indefinitely.
                            Floats can be used to specify units smaller than
                            seconds.
        :param kill_pg:     Boolean whether to kill the process group or only this
                            process. (not applicable for windows)
        """
        timedout = MutableValue(False)

        if seconds <= 0:
            yield timedout
            return

        finished = threading.Event()

        if platform.system() == 'Windows':  # pragma: no cover
            kill_pg = False

        def kill_it():
            finished.wait(seconds)
            if not finished.is_set():
                timedout.value = True
                if kill_pg:
                    pgid = os.getpgid(sub_process.pid)
                os.kill(sub_process.pid, signal.SIGINT)
                if kill_pg:
                    os.killpg(pgid, signal.SIGINT)

        thread = threading.Thread(name='timeout-killer', target=kill_it)
        try:
            thread.start()
            yield timedout
        finally:
            finished.set()
            thread.join()

    @contextmanager
    def replace_stdout(replacement):
        """
        Replaces stdout with the replacement, yields back to the caller and then
        reverts everything back.
        """
        _stdout = sys.stdout
        sys.stdout = replacement
        try:
            yield
        finally:
            sys.stdout = _stdout

    @contextmanager
    def replace_stderr(replacement):
        """
        Replaces stderr with the replacement, yields back to the caller and then
        reverts everything back.
        """
        _stderr = sys.stderr
        sys.stderr = replacement
        try:
            yield
        finally:
            sys.stderr = _stderr

    @contextmanager
    def suppress_stdout():
        """
        Suppresses everything going to stdout.
        """
        with open(os.devnull, 'w') as devnull, replace_stdout(devnull):
            yield

    @contextmanager
    def retrieve_stdout():
        """
        Yields a StringIO object from which one can read everything that was
        printed to stdout. (It won't be printed to the real stdout!)

        Example usage:

        with retrieve_stdout() as stdout:
            print("something")  # Won't print to the console
            what_was_printed = stdout.getvalue()  # Save the value
        """
        with closing(StringIO()) as sio, replace_stdout(sio):
            oldprint = builtins.print
            try:
                # Overriding stdout doesn't work with libraries, this ensures even
                # cached variables take this up. Well... it works.
                def newprint(*args, **kwargs):
                    kwargs['file'] = sio
                    oldprint(*args, **kwargs)

                builtins.print = newprint
                yield sio
            finally:
                builtins.print = oldprint

    @contextmanager
    def retrieve_stderr():
        """
        Yields a StringIO object from which one can read everything that was
        printed to stderr. (It won't be printed to the real stderr!)

        Example usage:

        with retrieve_stderr() as stderr:
            print("something")  # Won't print to the console
            what_was_printed = stderr.getvalue()  # Save the value
        """
        with closing(StringIO()) as sio, replace_stderr(sio):
            oldprint = builtins.print
            try:
                # Overriding stderr doesn't work with libraries, this ensures even
                # cached variables take this up. Well... it works.
                def newprint(*args, **kwargs):
                    kwargs['file'] = sio
                    oldprint(*args, **kwargs)

                builtins.print = newprint
                yield sio
            finally:
                builtins.print = oldprint

    @contextmanager
    def simulate_console_inputs(*inputs):
        """
        Does some magic to simulate the given inputs to any calls to the ``input``
        builtin. This yields back an InputGenerator object so you can check
        which input was already used and append any additional inputs you want.
        Example:

            with simulate_console_inputs(0, 1, 2) as generator:
                assert(input() == 0)
                assert(generator.last_input == 0)
                generator.inputs.append(3)
                assert(input() == 1)
                assert(input() == 2)
                assert(input() == 3)
                assert(generator.last_input == 3)

        :param inputs:      Any inputs to simulate.
        :raises ValueError: Raised when was asked for more input but there's no
                            more provided.
        """
        class InputGenerator:

            def __init__(self, inputs):
                self.last_input = -1
                self.inputs = inputs

            def generate_input(self, prompt=''):
                print(prompt, end='')
                self.last_input += 1
                try:
                    return self.inputs[self.last_input]
                except IndexError:
                    raise ValueError('Asked for more input, but no more was '
                                     'provided from `simulate_console_inputs`.')

        input_generator = InputGenerator(list(inputs))
        _input = builtins.input
        builtins.input = input_generator.generate_input
        try:
            yield input_generator
        finally:
            builtins.input = _input

    @contextmanager
    def make_temp(suffix='', prefix='tmp', dir=None):
        """
        Creates a temporary file with a closed stream and deletes it when done.

        :return: A contextmanager retrieving the file path.
        """
        temporary = tempfile.mkstemp(suffix=suffix, prefix=prefix, dir=dir)
        os.close(temporary[0])
        try:
            yield temporary[1]
        finally:
            os.remove(temporary[1])

    @contextmanager
    def prepare_file(lines,
                     filename,
                     force_linebreaks=True,
                     create_tempfile=True,
                     tempfile_kwargs={}):
        """
        Can create a temporary file (if filename is None) with the lines.
        Can also add a trailing newline to each line specified if needed.

        :param lines:            The lines from the file. (list or tuple of strings)
        :param filename:         The filename to be prepared.
        :param force_linebreaks: Whether to append newlines at each line if needed.
        :param create_tempfile:  Whether to save lines in tempfile if needed.
        :param tempfile_kwargs:  Kwargs passed to tempfile.mkstemp().
        """
        if force_linebreaks:
            lines = type(lines)(line if line.endswith('\n') else line+'\n'
                                for line in lines)

        if not create_tempfile and filename is None:
            filename = 'dummy_file_name'

        if not isinstance(filename, str) and create_tempfile:
            with make_temp(**tempfile_kwargs) as filename:
                with open(filename, 'w', encoding='utf-8') as file:
                    file.writelines(lines)
                yield lines, filename
        else:
            yield lines, filename

    @contextmanager
    def change_directory(path):
        old_dir = os.getcwd()
        os.chdir(path)
        try:
            yield
        finally:
            os.chdir(old_dir)