File size: 3,480 Bytes
d1ceb73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
import tempfile
import os
import shutil
import sys
import contextlib
import site
import io
from filelock import FileLock
@contextlib.contextmanager
def tempdir(cd=lambda dir: None, **kwargs):
temp_dir = tempfile.mkdtemp(**kwargs)
orig_dir = os.getcwd()
try:
cd(temp_dir)
yield temp_dir
finally:
cd(orig_dir)
shutil.rmtree(temp_dir)
@contextlib.contextmanager
def environment(**replacements):
"""
In a context, patch the environment with replacements. Pass None values
to clear the values.
"""
saved = dict((key, os.environ[key]) for key in replacements if key in os.environ)
# remove values that are null
remove = (key for (key, value) in replacements.items() if value is None)
for key in list(remove):
os.environ.pop(key, None)
replacements.pop(key)
os.environ.update(replacements)
try:
yield saved
finally:
for key in replacements:
os.environ.pop(key, None)
os.environ.update(saved)
@contextlib.contextmanager
def quiet():
"""
Redirect stdout/stderr to StringIO objects to prevent console output from
distutils commands.
"""
old_stdout = sys.stdout
old_stderr = sys.stderr
new_stdout = sys.stdout = io.StringIO()
new_stderr = sys.stderr = io.StringIO()
try:
yield new_stdout, new_stderr
finally:
new_stdout.seek(0)
new_stderr.seek(0)
sys.stdout = old_stdout
sys.stderr = old_stderr
@contextlib.contextmanager
def save_user_site_setting():
saved = site.ENABLE_USER_SITE
try:
yield saved
finally:
site.ENABLE_USER_SITE = saved
@contextlib.contextmanager
def save_pkg_resources_state():
import pkg_resources
pr_state = pkg_resources.__getstate__()
# also save sys.path
sys_path = sys.path[:]
try:
yield pr_state, sys_path
finally:
sys.path[:] = sys_path
pkg_resources.__setstate__(pr_state)
@contextlib.contextmanager
def suppress_exceptions(*excs):
try:
yield
except excs:
pass
def multiproc(request):
"""
Return True if running under xdist and multiple
workers are used.
"""
try:
worker_id = request.getfixturevalue('worker_id')
except Exception:
return False
return worker_id != 'master'
@contextlib.contextmanager
def session_locked_tmp_dir(request, tmp_path_factory, name):
"""Uses a file lock to guarantee only one worker can access a temp dir"""
# get the temp directory shared by all workers
base = tmp_path_factory.getbasetemp()
shared_dir = base.parent if multiproc(request) else base
locked_dir = shared_dir / name
with FileLock(locked_dir.with_suffix(".lock")):
# ^-- prevent multiple workers to access the directory at once
locked_dir.mkdir(exist_ok=True, parents=True)
yield locked_dir
@contextlib.contextmanager
def save_paths():
"""Make sure ``sys.path``, ``sys.meta_path`` and ``sys.path_hooks`` are preserved"""
prev = sys.path[:], sys.meta_path[:], sys.path_hooks[:]
try:
yield
finally:
sys.path, sys.meta_path, sys.path_hooks = prev
@contextlib.contextmanager
def save_sys_modules():
"""Make sure initial ``sys.modules`` is preserved"""
prev_modules = sys.modules
try:
sys.modules = sys.modules.copy()
yield
finally:
sys.modules = prev_modules
|