• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import contextlib
2import importlib
3import importlib.util
4import os
5import sys
6import unittest
7import warnings
8
9from .os_helper import unlink
10
11
12@contextlib.contextmanager
13def _ignore_deprecated_imports(ignore=True):
14    """Context manager to suppress package and module deprecation
15    warnings when importing them.
16
17    If ignore is False, this context manager has no effect.
18    """
19    if ignore:
20        with warnings.catch_warnings():
21            warnings.filterwarnings("ignore", ".+ (module|package)",
22                                    DeprecationWarning)
23            yield
24    else:
25        yield
26
27
28def unload(name):
29    try:
30        del sys.modules[name]
31    except KeyError:
32        pass
33
34
35def forget(modname):
36    """'Forget' a module was ever imported.
37
38    This removes the module from sys.modules and deletes any PEP 3147/488 or
39    legacy .pyc files.
40    """
41    unload(modname)
42    for dirname in sys.path:
43        source = os.path.join(dirname, modname + '.py')
44        # It doesn't matter if they exist or not, unlink all possible
45        # combinations of PEP 3147/488 and legacy pyc files.
46        unlink(source + 'c')
47        for opt in ('', 1, 2):
48            unlink(importlib.util.cache_from_source(source, optimization=opt))
49
50
51def make_legacy_pyc(source):
52    """Move a PEP 3147/488 pyc file to its legacy pyc location.
53
54    :param source: The file system path to the source file.  The source file
55        does not need to exist, however the PEP 3147/488 pyc file must exist.
56    :return: The file system path to the legacy pyc file.
57    """
58    pyc_file = importlib.util.cache_from_source(source)
59    up_one = os.path.dirname(os.path.abspath(source))
60    legacy_pyc = os.path.join(up_one, source + 'c')
61    os.rename(pyc_file, legacy_pyc)
62    return legacy_pyc
63
64
65def import_module(name, deprecated=False, *, required_on=()):
66    """Import and return the module to be tested, raising SkipTest if
67    it is not available.
68
69    If deprecated is True, any module or package deprecation messages
70    will be suppressed. If a module is required on a platform but optional for
71    others, set required_on to an iterable of platform prefixes which will be
72    compared against sys.platform.
73    """
74    with _ignore_deprecated_imports(deprecated):
75        try:
76            return importlib.import_module(name)
77        except ImportError as msg:
78            if sys.platform.startswith(tuple(required_on)):
79                raise
80            raise unittest.SkipTest(str(msg))
81
82
83def _save_and_remove_modules(names):
84    orig_modules = {}
85    prefixes = tuple(name + '.' for name in names)
86    for modname in list(sys.modules):
87        if modname in names or modname.startswith(prefixes):
88            orig_modules[modname] = sys.modules.pop(modname)
89    return orig_modules
90
91
92def import_fresh_module(name, fresh=(), blocked=(), deprecated=False):
93    """Import and return a module, deliberately bypassing sys.modules.
94
95    This function imports and returns a fresh copy of the named Python module
96    by removing the named module from sys.modules before doing the import.
97    Note that unlike reload, the original module is not affected by
98    this operation.
99
100    *fresh* is an iterable of additional module names that are also removed
101    from the sys.modules cache before doing the import. If one of these
102    modules can't be imported, None is returned.
103
104    *blocked* is an iterable of module names that are replaced with None
105    in the module cache during the import to ensure that attempts to import
106    them raise ImportError.
107
108    The named module and any modules named in the *fresh* and *blocked*
109    parameters are saved before starting the import and then reinserted into
110    sys.modules when the fresh import is complete.
111
112    Module and package deprecation messages are suppressed during this import
113    if *deprecated* is True.
114
115    This function will raise ImportError if the named module cannot be
116    imported.
117    """
118    # NOTE: test_heapq, test_json and test_warnings include extra sanity checks
119    # to make sure that this utility function is working as expected
120    with _ignore_deprecated_imports(deprecated):
121        # Keep track of modules saved for later restoration as well
122        # as those which just need a blocking entry removed
123        fresh = list(fresh)
124        blocked = list(blocked)
125        names = {name, *fresh, *blocked}
126        orig_modules = _save_and_remove_modules(names)
127        for modname in blocked:
128            sys.modules[modname] = None
129
130        try:
131            # Return None when one of the "fresh" modules can not be imported.
132            try:
133                for modname in fresh:
134                    __import__(modname)
135            except ImportError:
136                return None
137            return importlib.import_module(name)
138        finally:
139            _save_and_remove_modules(names)
140            sys.modules.update(orig_modules)
141
142
143class CleanImport(object):
144    """Context manager to force import to return a new module reference.
145
146    This is useful for testing module-level behaviours, such as
147    the emission of a DeprecationWarning on import.
148
149    Use like this:
150
151        with CleanImport("foo"):
152            importlib.import_module("foo") # new reference
153    """
154
155    def __init__(self, *module_names):
156        self.original_modules = sys.modules.copy()
157        for module_name in module_names:
158            if module_name in sys.modules:
159                module = sys.modules[module_name]
160                # It is possible that module_name is just an alias for
161                # another module (e.g. stub for modules renamed in 3.x).
162                # In that case, we also need delete the real module to clear
163                # the import cache.
164                if module.__name__ != module_name:
165                    del sys.modules[module.__name__]
166                del sys.modules[module_name]
167
168    def __enter__(self):
169        return self
170
171    def __exit__(self, *ignore_exc):
172        sys.modules.update(self.original_modules)
173
174
175class DirsOnSysPath(object):
176    """Context manager to temporarily add directories to sys.path.
177
178    This makes a copy of sys.path, appends any directories given
179    as positional arguments, then reverts sys.path to the copied
180    settings when the context ends.
181
182    Note that *all* sys.path modifications in the body of the
183    context manager, including replacement of the object,
184    will be reverted at the end of the block.
185    """
186
187    def __init__(self, *paths):
188        self.original_value = sys.path[:]
189        self.original_object = sys.path
190        sys.path.extend(paths)
191
192    def __enter__(self):
193        return self
194
195    def __exit__(self, *ignore_exc):
196        sys.path = self.original_object
197        sys.path[:] = self.original_value
198
199
200def modules_setup():
201    return sys.modules.copy(),
202
203
204def modules_cleanup(oldmodules):
205    # Encoders/decoders are registered permanently within the internal
206    # codec cache. If we destroy the corresponding modules their
207    # globals will be set to None which will trip up the cached functions.
208    encodings = [(k, v) for k, v in sys.modules.items()
209                 if k.startswith('encodings.')]
210    sys.modules.clear()
211    sys.modules.update(encodings)
212    # XXX: This kind of problem can affect more than just encodings.
213    # In particular extension modules (such as _ssl) don't cope
214    # with reloading properly. Really, test modules should be cleaning
215    # out the test specific modules they know they added (ala test_runpy)
216    # rather than relying on this function (as test_importhooks and test_pkg
217    # do currently). Implicitly imported *real* modules should be left alone
218    # (see issue 10556).
219    sys.modules.update(oldmodules)
220