Blocks system import inside function

block_import_decorator.py

from functools import wraps


def block_imports(*imports):
    def real_decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            saved = {}
            import sys
            for i in imports:
                saved[i] = sys.modules[i]
                sys.modules[i] = None
            retval = func(*args, **kwargs)
            for i in imports:
                sys.modules[i] = saved[i]
            return retval

        return wrapper

    return real_decorator


@block_imports("os")
def correct():
    try:
        import os
    except ModuleNotFoundError:
        print("os cannot be imported")


correct()
import os

print("import succeded outside function")