1import tempfile
2import os
3import shutil
4import sys
5import contextlib
6import site
7
8from setuptools.extern import six
9import pkg_resources
10
11
12@contextlib.contextmanager
13def 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_dir
19    finally:
20        cd(orig_dir)
21        shutil.rmtree(temp_dir)
22
23
24@contextlib.contextmanager
25def environment(**replacements):
26    """
27    In a context, patch the environment with replacements. Pass None values
28    to clear the values.
29    """
30    saved = dict(
31        (key, os.environ[key])
32        for key in replacements
33        if key in os.environ
34    )
35
36    # remove values that are null
37    remove = (key for (key, value) in replacements.items() if value is None)
38    for key in list(remove):
39        os.environ.pop(key, None)
40        replacements.pop(key)
41
42    os.environ.update(replacements)
43
44    try:
45        yield saved
46    finally:
47        for key in replacements:
48            os.environ.pop(key, None)
49        os.environ.update(saved)
50
51
52@contextlib.contextmanager
53def quiet():
54    """
55    Redirect stdout/stderr to StringIO objects to prevent console output from
56    distutils commands.
57    """
58
59    old_stdout = sys.stdout
60    old_stderr = sys.stderr
61    new_stdout = sys.stdout = six.StringIO()
62    new_stderr = sys.stderr = six.StringIO()
63    try:
64        yield new_stdout, new_stderr
65    finally:
66        new_stdout.seek(0)
67        new_stderr.seek(0)
68        sys.stdout = old_stdout
69        sys.stderr = old_stderr
70
71
72@contextlib.contextmanager
73def save_user_site_setting():
74    saved = site.ENABLE_USER_SITE
75    try:
76        yield saved
77    finally:
78        site.ENABLE_USER_SITE = saved
79
80
81@contextlib.contextmanager
82def save_pkg_resources_state():
83    pr_state = pkg_resources.__getstate__()
84    # also save sys.path
85    sys_path = sys.path[:]
86    try:
87        yield pr_state, sys_path
88    finally:
89        sys.path[:] = sys_path
90        pkg_resources.__setstate__(pr_state)
91
92
93@contextlib.contextmanager
94def suppress_exceptions(*excs):
95    try:
96        yield
97    except excs:
98        pass
99