codekingpro/portable-devtools

1import contextlib2import io3import os4import shutil5import site6import sys7import tempfile8 9from filelock import FileLock10 11 12@contextlib.contextmanager13def tempdir(cd=lambda dir: None, **kwargs):14    temp_dir = tempfile.mkdtemp(**kwargs)15    orig_dir = os.getcwd()16    try:17        cd(temp_dir)18        yield temp_dir19    finally:20        cd(orig_dir)21        shutil.rmtree(temp_dir)22 23 24@contextlib.contextmanager25def environment(**replacements):26    """27    In a context, patch the environment with replacements. Pass None values28    to clear the values.29    """30    saved = dict((key, os.environ[key]) for key in replacements if key in os.environ)31 32    # remove values that are null33    remove = (key for (key, value) in replacements.items() if value is None)34    for key in list(remove):35        os.environ.pop(key, None)36        replacements.pop(key)37 38    os.environ.update(replacements)39 40    try:41        yield saved42    finally:43        for key in replacements:44            os.environ.pop(key, None)45        os.environ.update(saved)46 47 48@contextlib.contextmanager49def quiet():50    """51    Redirect stdout/stderr to StringIO objects to prevent console output from52    distutils commands.53    """54 55    old_stdout = sys.stdout56    old_stderr = sys.stderr57    new_stdout = sys.stdout = io.StringIO()58    new_stderr = sys.stderr = io.StringIO()59    try:60        yield new_stdout, new_stderr61    finally:62        new_stdout.seek(0)63        new_stderr.seek(0)64        sys.stdout = old_stdout65        sys.stderr = old_stderr66 67 68@contextlib.contextmanager69def save_user_site_setting():70    saved = site.ENABLE_USER_SITE71    try:72        yield saved73    finally:74        site.ENABLE_USER_SITE = saved75 76 77@contextlib.contextmanager78def suppress_exceptions(*excs):79    try:80        yield81    except excs:82        pass83 84 85def multiproc(request):86    """87    Return True if running under xdist and multiple88    workers are used.89    """90    try:91        worker_id = request.getfixturevalue('worker_id')92    except Exception:93        return False94    return worker_id != 'master'95 96 97@contextlib.contextmanager98def session_locked_tmp_dir(request, tmp_path_factory, name):99    """Uses a file lock to guarantee only one worker can access a temp dir"""100    # get the temp directory shared by all workers101    base = tmp_path_factory.getbasetemp()102    shared_dir = base.parent if multiproc(request) else base103 104    locked_dir = shared_dir / name105    with FileLock(locked_dir.with_suffix(".lock")):106        # ^-- prevent multiple workers to access the directory at once107        locked_dir.mkdir(exist_ok=True, parents=True)108        yield locked_dir109 110 111@contextlib.contextmanager112def save_paths():113    """Make sure ``sys.path``, ``sys.meta_path`` and ``sys.path_hooks`` are preserved"""114    prev = sys.path[:], sys.meta_path[:], sys.path_hooks[:]115 116    try:117        yield118    finally:119        sys.path, sys.meta_path, sys.path_hooks = prev120 121 122@contextlib.contextmanager123def save_sys_modules():124    """Make sure initial ``sys.modules`` is preserved"""125    prev_modules = sys.modules126 127    try:128        sys.modules = sys.modules.copy()129        yield130    finally:131        sys.modules = prev_modules132