missing_modules.py

"""Ban some modules from being importable inside the context

For example:

with missing_modules("sys"):
    try:
        import sys
    except ImportError:
        print "sys not found"

import sys
print sys.version

results in:

sys not found
(2, 5, 2, 'final', 0)

"""

from __future__ import with_statement
from contextlib import contextmanager
import __builtin__

@contextmanager
def missing_modules(*modnames):
    def myimp(name, *args, **kwargs):
        if name in modnames:
            raise ImportError("No module named %s" % name)
        else:
            return realimport(name, *args, **kwargs)

    realimport = __builtin__.__import__
    __builtin__.__import__ = myimp
    yield
    __builtin__.__import__ = realimport