Override `import` for more sophisticated module import

Question

Is it possible to somehow override import so that I can do some more sophisticated operations on a module before it gets imported?

As an example: I have a larger application that uses matplotlib for secondary features that are not vital for the overall functionality of the application. In case that matplotlib is not installed I just want to mock the functionality so that the import and all calls to matplotlib functions appear to be working, just without actually doing anything. A simple warning should then just indicate that the module is not installed, though that fact would not impair the core functionality of the application. I already have an import function that, in the case that matplotlib is not installed, returns a MagicMock object instead of the actual module which just mimics the behavior of the matplotlib API.

So, all import matplotlib... or from matplotlib import... should then be automatically overridden by the corresponding function call. I could replace all import and from ... import expressions by hand but I'd like to make but there are a lot of them. I'd rather like to have this functionality automatically by overriding import.

Is that possible?

Answer 1

The easiest way I've found is to replace the __import__ function with your own implementation (described here). Then if someone tries to import matplotlib, you just import a different module instead:

def _import(name, *args, **kwargs):
    if name == 'matplotlib': # if someone tries to import matplotlib...
        name = 'my_mocked_matplotlib' # ...import the mocked version instead
    return original_import(name, *args, **kwargs)

import builtins
original_import = builtins.__import__
builtins.__import__ = _import

---

To restrict the custom import behavior to only a few modules, you can use introspection (with the inspect module) to find out from which module the import has been performed:

import inspect

def _import(name, *args, **kwargs):
    if name == 'matplotlib': # if someone tries to import matplotlib...
        # find out which module is performing the import
        frame = inspect.currentframe().f_back
        module_path = frame.f_globals['__file__']

        # if the import is happening in module1 or module2, redirect it
        if module_path in ('/path/to/module1.py','/path/to/module2.py'):
            name = 'my_mocked_matplotlib' # ...import the mocked version instead

    return original_import(name, *args, **kwargs)

Answer 2

For a truly sophisticated import (adding functionality, hooking, etc.) you can define a meta hook and override whatever portion of the import behavior that you need. Example override script:

from os.path import (dirname, abspath, join, splitext)
from os import (listdir)
import importlib.util
import imp
import sys
_scripts_path = dirname(abspath(__file__))
_script_list = [splitext(f)[0] for f in listdir(_scripts_path) if splitext(f)[1] == '.py']

class custom_import_hook(object):
    def find_module(self, name, path):
        if name not in _script_list: return None
        return self
    
    class _cmds(object):
        @staticmethod
        def my_print(string):
            print(string)
    
    def load_module(self, name):
        sys.modules.setdefault(name, imp.new_module(name))
        spec = importlib.util.spec_from_file_location('module.name', join(_scripts_path, f'{name}.py'))
        foo = importlib.util.module_from_spec(spec)
        for cmd in self._cmds.__dict__.keys():
            if cmd[0] == '_': continue
            setattr(foo, cmd, getattr(self._cmds, cmd))
        sys.meta_path.append(self)
        spec.loader.exec_module(foo)
        sys.meta_path.remove(self)
        return foo

meta_hook = custom_import_hook()
sys.meta_path.insert(0, meta_hook)
import test_script

then place in the same folder a script called 'test_script.py':

my_print('test')

and run the override script, you'll see that it defines 'my_print' such that when test_script is executed that your custom print command is used. You can customize essentially anything you want at this point (e.g. wrapping a library with custom command replacement), since the hook gives you full control over the import.