|
1 | 1 | import contextlib
|
2 | 2 | import _imp
|
3 | 3 | import importlib
|
| 4 | +import importlib.machinery |
4 | 5 | import importlib.util
|
5 | 6 | import os
|
6 | 7 | import shutil
|
@@ -332,3 +333,110 @@ def ensure_lazy_imports(imported_module, modules_to_block):
|
332 | 333 | )
|
333 | 334 | from .script_helper import assert_python_ok
|
334 | 335 | assert_python_ok("-S", "-c", script)
|
| 336 | + |
| 337 | + |
| 338 | +@contextlib.contextmanager |
| 339 | +def module_restored(name): |
| 340 | + """A context manager that restores a module to the original state.""" |
| 341 | + missing = object() |
| 342 | + orig = sys.modules.get(name, missing) |
| 343 | + if orig is None: |
| 344 | + mod = importlib.import_module(name) |
| 345 | + else: |
| 346 | + mod = type(sys)(name) |
| 347 | + mod.__dict__.update(orig.__dict__) |
| 348 | + sys.modules[name] = mod |
| 349 | + try: |
| 350 | + yield mod |
| 351 | + finally: |
| 352 | + if orig is missing: |
| 353 | + sys.modules.pop(name, None) |
| 354 | + else: |
| 355 | + sys.modules[name] = orig |
| 356 | + |
| 357 | + |
| 358 | +def create_module(name, loader=None, *, ispkg=False): |
| 359 | + """Return a new, empty module.""" |
| 360 | + spec = importlib.machinery.ModuleSpec( |
| 361 | + name, |
| 362 | + loader, |
| 363 | + origin='<import_helper>', |
| 364 | + is_package=ispkg, |
| 365 | + ) |
| 366 | + return importlib.util.module_from_spec(spec) |
| 367 | + |
| 368 | + |
| 369 | +def _ensure_module(name, ispkg, addparent, clearnone): |
| 370 | + try: |
| 371 | + mod = orig = sys.modules[name] |
| 372 | + except KeyError: |
| 373 | + mod = orig = None |
| 374 | + missing = True |
| 375 | + else: |
| 376 | + missing = False |
| 377 | + if mod is not None: |
| 378 | + # It was already imported. |
| 379 | + return mod, orig, missing |
| 380 | + # Otherwise, None means it was explicitly disabled. |
| 381 | + |
| 382 | + assert name != '__main__' |
| 383 | + if not missing: |
| 384 | + assert orig is None, (name, sys.modules[name]) |
| 385 | + if not clearnone: |
| 386 | + raise ModuleNotFoundError(name) |
| 387 | + del sys.modules[name] |
| 388 | + # Try normal import, then fall back to adding the module. |
| 389 | + try: |
| 390 | + mod = importlib.import_module(name) |
| 391 | + except ModuleNotFoundError: |
| 392 | + if addparent and not clearnone: |
| 393 | + addparent = None |
| 394 | + mod = _add_module(name, ispkg, addparent) |
| 395 | + return mod, orig, missing |
| 396 | + |
| 397 | + |
| 398 | +def _add_module(spec, ispkg, addparent): |
| 399 | + if isinstance(spec, str): |
| 400 | + name = spec |
| 401 | + mod = create_module(name, ispkg=ispkg) |
| 402 | + spec = mod.__spec__ |
| 403 | + else: |
| 404 | + name = spec.name |
| 405 | + mod = importlib.util.module_from_spec(spec) |
| 406 | + sys.modules[name] = mod |
| 407 | + if addparent is not False and spec.parent: |
| 408 | + _ensure_module(spec.parent, True, addparent, bool(addparent)) |
| 409 | + return mod |
| 410 | + |
| 411 | + |
| 412 | +def add_module(spec, *, parents=True): |
| 413 | + """Return the module after creating it and adding it to sys.modules. |
| 414 | +
|
| 415 | + If parents is True then also create any missing parents. |
| 416 | + """ |
| 417 | + return _add_module(spec, False, parents) |
| 418 | + |
| 419 | + |
| 420 | +def add_package(spec, *, parents=True): |
| 421 | + """Return the module after creating it and adding it to sys.modules. |
| 422 | +
|
| 423 | + If parents is True then also create any missing parents. |
| 424 | + """ |
| 425 | + return _add_module(spec, True, parents) |
| 426 | + |
| 427 | + |
| 428 | +def ensure_module_imported(name, *, clearnone=True): |
| 429 | + """Return the corresponding module. |
| 430 | +
|
| 431 | + If it was already imported then return that. Otherwise, try |
| 432 | + importing it (optionally clear it first if None). If that fails |
| 433 | + then create a new empty module. |
| 434 | +
|
| 435 | + It can be helpful to combine this with ready_to_import() and/or |
| 436 | + isolated_modules(). |
| 437 | + """ |
| 438 | + if sys.modules.get(name) is not None: |
| 439 | + mod = sys.modules[name] |
| 440 | + else: |
| 441 | + mod, _, _ = _force_import(name, False, True, clearnone) |
| 442 | + return mod |
0 commit comments