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