1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798 |
- import tempfile
- import os
- import shutil
- import sys
- import contextlib
- import site
- import io
- import pkg_resources
- @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():
- 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
|