Function temporary_main_spec

Lib/test/test_pyclbr.py:28–46

A context manager that temporarily sets the `__spec__` attribute of the `__main__` module if it's missing.

()

@contextmanager
def temporary_main_spec():
    """
    A context manager that temporarily sets the `__spec__` attribute
    of the `__main__` module if it's missing.
    """
    main_mod = sys.modules.get("__main__")
    if main_mod is None:
        yield  # Do nothing if __main__ is not present
        return

    original_spec = getattr(main_mod, "__spec__", None)
    if original_spec is None:
        main_mod.__spec__ = importlib.machinery.ModuleSpec(
            name="__main__", loader=None, origin="built-in"
        )
    try:
        yield
    finally:
        main_mod.__spec__ = original_spec


class PyclbrTest(TestCase):