Override `import` for more sophisticated module import

user_205147 0
Viewed 2421

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?

import python
3 Answers

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.

user_1858280
Related