__init__.py

#!/usr/bin/env python
#
# __init__.py - Miscellaneous functions used throughout fsl_ci.
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
#


import os.path    as op
import               os
import               sys
import               errno
import               shlex
import               time
import               tempfile
import contextlib as ctxlib
import subprocess as sp

import yaml


__version__ = '0.15.2'
"""Current version of the fsl-ci-rules."""


USERNAME = 'fsl-ci-rules'
"""Username to be used for all git interactions which require one. """


EMAIL = 'fsl-ci-rules@git.fmrib.ox.ac.uk'
"""Password to be used for all git interactions which require one. """


def fprint(*args, **kwargs):
    """Print with flush=True. """
    print(*args, **kwargs, flush=True)


@ctxlib.contextmanager
def tempdir():
    """Context manager to create, and change into, a temporary directory, and
    then afterwards delete it and change back to the original working
    directory.
    """
    with tempfile.TemporaryDirectory() as td:
        prevdir = os.getcwd()
        os.chdir(td)
        try:
            yield td
        finally:
            os.chdir(prevdir)


@ctxlib.contextmanager
def indir(dirname):
    """Context manager to change into a directory, and then afterwards
    change back to the original working directory.
    """
    prevdir = os.getcwd()
    os.chdir(dirname)
    try:
        yield
    finally:
        os.chdir(prevdir)


def sprun(cmd, **kwargs):
    """Runs the given command with subprocess.run. """
    fprint(f'Running {cmd}')
    if not kwargs.get('shell', False):
        cmd = shlex.split(cmd)

    # we set stdout/stderr for compatibility
    # with CaptureStdout (see below)
    if 'check'  not in kwargs: kwargs['check']  = True

    if 'capture_output' not in kwargs:
        if 'stdout' not in kwargs: kwargs['stdout'] = sys.stdout
        if 'stderr' not in kwargs: kwargs['stderr'] = sys.stderr

    return sp.run(cmd, **kwargs)


def spcap(cmd, **kwargs):
    """Call sprun(..., capture_output=True), and return stdout as a string. """
    return sprun(cmd, capture_output=True, **kwargs).stdout.decode()


def which(command):
    """Run "which command". """
    return spcap(f'which {command}').strip()


@ctxlib.contextmanager
def lockdir(dirname):
    """Primitive mechanism by which concurrent access to a directory can be
    prevented. Attempts to create a semaphore file in the directory, but waits
    if that file already exists. Removes the file when finished.
    """

    delay    = 10
    lockfile = op.join(dirname, '.fsl_ci.lockdir')

    while True:
        try:
            fprint(f'Attempting to lock {dirname} for exclusive access.')
            fd = os.open(lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR)
            break
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise e
            fprint(f'{dirname} is already locked - '
                   f'trying again in {delay} seconds ...')
        time.sleep(10)

    fprint(f'Exclusive access acquired for {dirname} ...')
    try:
        yield

    finally:
        fprint(f'Relinquishing lock on {dirname}')
        os.close( fd)
        os.unlink(lockfile)


def loadyaml(s):
    """Loads a YAML string, returning a dict-like. """
    return yaml.load(s, Loader=yaml.Loader)


def dumpyaml(o):
    """Dumps the given YAML to a string. """
    return yaml.dump(o, sort_keys=False)


class CaptureStdout:
    """Context manager which captures stdout and stderr. """

    def __init__(self):
        self.tmpdir        = tempfile.TemporaryDirectory()
        self.__mock_stdout = open(op.join(self.tmpdir.name, 'stdout'), 'wt')
        self.__mock_stderr = open(op.join(self.tmpdir.name, 'stderr'), 'wt')


    def __enter__(self):
        if self.tmpdir is None:
            raise RuntimeError('CaptureStdout can only be used once')

        self.__real_stdout = sys.stdout
        self.__real_stderr = sys.stderr
        sys.stdout         = self.__mock_stdout
        sys.stderr         = self.__mock_stderr

        return self

    def __exit__(self, *args, **kwargs):

        sys.stdout = self.__real_stdout
        sys.stderr = self.__real_stderr

        self.__mock_stdout.close()
        self.__mock_stderr.close()

        with open(self.__mock_stdout.name, 'rt') as f:
            self.__mock_stdout = f.read()
        with open(self.__mock_stderr.name, 'rt') as f:
            self.__mock_stderr = f.read()

        self.tmpdir.cleanup()
        self.tmpdir = None

    @property
    def real_stdout(self):
        return self.__real_stdout

    @property
    def real_stderr(self):
        return self.__real_stderr

    @property
    def stdout(self):
        return self.__mock_stdout

    @property
    def stderr(self):
        return self.__mock_stderr