Post import hooks in Python 3

Question

I would like to have some callback run whenever a particular module is imported. For example (using a fake @imp.when_imported function that does not really exist):

@imp.when_imported('numpy')
def set_linewidth(numpy):
    import shutil
    numpy.set_printoptions(linewidth=shutil.get_terminal_size()[0])

This feature was designed in PEP 369: Post import hooks but was withdrawn with the reason:

This PEP has been withdrawn by its author, as much of the detailed design is no longer valid following the migration to importlib in Python 3.3.

But importlib has no clear solution. How does one use importlib to implement a post-import hook?

Answer 1

I would be shocked to find out that this is the best way to do this ... However, since early python2.x versions, monkey patching __import__ has been supported. We can take advantage of that here:

try:
    import builtins  # python3.x
except ImportError:
    import __builtin__ as builtins  # python2.x
import sys
import collections

_builtin_import = builtins.__import__

def _my_import(name, globals=None, locals=None, fromlist=(), level=0):
    already_imported = name in sys.modules

    mod = _builtin_import(
        name,
        globals=globals,
        locals=locals,
        fromlist=fromlist,
        level=level)

    if not already_imported and name in _post_import_hooks:
        for hook in _post_import_hooks[name]:
            hook()
    return mod

builtins.__import__ = _my_import

_post_import_hooks = collections.defaultdict(list)

def on_import(name):
    def decorator(func):
        _post_import_hooks[name].append(func)
        return func
    return decorator

@on_import('numpy')
def print_hi():
    print('Hello Numpy')

print('before numpy')
import numpy
print('after numpy')

This answer makes a super simple registry for registering callbacks. The decorator just registers the function and then returns it. It doesn't do any fancy checking (for whether the module is already loaded, for example), but could easily be extended to do that.

Obviously the downside is if some other module decides to monkey patch __import__, then you're out of luck -- Either this module or the other one is likely to end up broken.

I've tested this and it seems to work on both python2.x and python3.x.

Answer 2

The wrapt module provides an implementation of this.

• https://pypi.python.org/pypi/wrapt

Watch this video about wrapt, including this feature:

• https://www.youtube.com/watch?v=u7oj-ghfhUk

Don't think the documentation for wrapt mentions it yet.

Some of the blogs posts at end of:

• http://blog.dscpl.com.au/p/decorators-and-monkey-patching.html

talk about it though.

There is a companion module for wrapt called autowrapt which allows you to do monkey patching using this mechanism without needing to change the application code itself to trigger it.

• https://pypi.python.org/pypi/autowrapt

Answer 3

Does this work?

import importlib

class ImportHook:

    def __init__(self, func):
        self.func = func
        self.module = None

    def __enter__(self):
        return self

    def get_module(self, module_name):
        self.module = importlib.import_module(module_name)
        return self.module

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.module is not None:
            self.func(self.module)

def set_linewidth(module):
    import shutil
    module.set_printoptions(linewidth=shutil.get_terminal_size()[0])

with ImportHook(set_linewidth) as hook:
    numpy = hook.get_module('numpy')