"""Create a controllable import.

One should be able to do several things to control imports:

    * Whitelisting of modules based on name and type
        + built-ins
        + frozen
        + extensions
    * Block all modules based on type
        + .pyc
    * Allow all modules of a certain type
        + .py
"""
import importlib
import _importlib
import functools
import imp
import sys

# Need to inject 'open' into importlib as it will most likely be removed from
# the built-in namespace.
_importlib.open = open


class Whitelister(object):

    """Whitelist the finding and/or loading of modules."""

    def __init__(self, whitelist, *args, **kwargs):
        """Store whitelist and continue on with instantiation."""
        self._whitelist = frozenset(whitelist)
        super(Whitelister, self).__init__(*args, **kwargs)

    def check_name(self, name):
        """Check a module name against the whitelist, returning True if it
        passes or False otherwise."""
        if name in self._whitelist:
            return True
        return False

    def find_module(self, name, path=None):
        """If the module name passes the whitelist, allow the finding of the
        modules to continue, otherwise return None."""
        if self.check_name(name):
            return super(Whitelister, self).find_module(name, path)
        return None

    def load_module(self, name):
        """If the module names passes the whitelist, allow the loading of the
        module to continue, otherwise raise ImportError."""
        if self.check_name(name):
            return super(Whitelister, self).load_module(name)
        raise ImportError("cannot import module")


class WhitelistBuiltin(Whitelister, importlib.BuiltinImporter):

    """Whitelisting importer/loader for built-in modules."""

    pass


class WhitelistFrozen(Whitelister, importlib.FrozenImporter):

    """Whitelisting importer/loader for frozen modules."""

    pass


class WhitelistExtensionImporter(Whitelister, importlib.ExtensionFileImporter):

    """Whitelisting importer for extension modules."""

    pass


class PyOnlyFileLoader(importlib._PyFileLoader):

    """Only load source files; no bytecode."""

    def __init__(self, *args, **kwargs):
        """Remove the bytecode path."""
        super(PyOnlyFileLoader, self).__init__(*args, **kwargs)
        self._bytecode_path = lambda: None

    def write_bytecode(self, *args, **kwargs):
        """Do not write out any bytecode."""
        return False


class PyOnlyFileImporter(importlib.PyFileImporter):

    """Only find source files; no bytecode."""

    _loader = PyOnlyFileLoader

    def __init__(self, path_entry):
        super(PyOnlyFileImporter, self).__init__(path_entry)
        self._suffixes = importlib.suffix_list(imp.PY_SOURCE)


class ControlledImport(importlib.Import):

    """Represent a controllable version of import that allows for more
    fine-grained control over what can and cannot be imported."""

    def __init__(self, safe_builtins=(), safe_frozen=(), safe_extensions=()):
        """Set up importation where built-in, frozen, and extension modules
        must be whitelisted and .pyc files are completely blocked
        while all. py files are allowed.
        
        Import data structures in sys are directly mutated when an instance is
        created!
        """
        importlib.Import.__init__(self, None, ())
        # Clear out any traces of the previous import machinery.
        sys.path_importer_cache.clear()
        sys.meta_path = []
        sys.path_hooks = []
        # Whitelist built-in and frozen modules on sys.meta_path.
        sys.meta_path.append(WhitelistBuiltin(safe_builtins))
        sys.meta_path.append(WhitelistFrozen(safe_frozen))
        # Whitelist extension modules on sys.path.
        ext_importer = functools.partial(WhitelistExtensionImporter,
                                            safe_extensions)
        # Allow all .py files but not .pyc files on sys.path.
        # XXX
        fs_factory = importlib.chaining_fs_path_hook(ext_importer,
                                                        PyOnlyFileImporter)
        sys.path_hooks.append(fs_factory)

    def module_from_cache(self, name):
        """Override so that any module name starting with a dot raises
        ImportError."""
        if name.startswith('.'):
            raise ImportError("module names starting with '.' cannot be "
                                "imported")
        return importlib.Import.module_from_cache(self, name)

    def post_import(self, module):
        """Strip off the __loader__ attribute so that the object's global
        namespace is not exposed."""
        try:
            del module.__loader__
        except AttributeError:
            pass
        return module