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?

Bohol

Could you just add the code at the very end of the module? – Lacrimator

No, I want to do this for modules I don't own (particularly numpy). – Bohol

Does it have to be right after the import or can it be before the main code? – Lighter

If you have the on_import hook in the module you provide, you might as well import numpy and run that function right away (since modules are global across the entire interpreter). The other consideration is that this "feature" is implicit and hidden to users of the module you provide, which they may not desire. – Caroylncarp

There is no on_import hook - that is what I'm asking. And I don't want to always import numpy because it is slow or may be unavailable. – Bohol

@StevenSummers I am intending this for an interactive shell. So at startup nothing would happen, but then whenever I do import numpy as np in the shell, my function would get called. – Bohol

I switched to the name imp.when_imported to match PEP 369. Hopefully that should remove confusion. – Bohol

This is an interesting use-case, but, asserting that importing numpy "may be slow" seems a bit much. On my system it takes less than .1s (which is a 1-time cost) which is completely negligible -- especially if this is for use in interactive sessions. The unavailable case is more interesting, but can be handled by catching the ImportError. – Barbera

@Barbera It is slow on my machine if the OS has not cached the files, and matplotlib is even slower. Even if it were fast, you're not answering my question, which is how to create a post-import hook. – Bohol

@MarkLodato -- You're right. If I was answering your question, I would have posted an answer :-). I'm just saying that based on the constraints that you've mentioned, it seems like there are less involved solutions. – Barbera

Answer 1

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

Michelemichelina

Thanks for the tip! I'd like to avoid relying on third-party modules so I was looking for a simple solution using importlib, since PEP 369 implied that one was easy. However, looking at wrapt's implementation it seems that the answer to my question is not so simple. – Bohol

I suggest updating the solution to link to github.com/GrahamDumpleton/wrapt/blob/master/src/wrapt/…, which is full implementation of PEP 369. I suppose I could copy the file and strip it down to the bare essentials, but for now I'll just use wrapt. Thanks! – Bohol

Answer 2

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.

Barbera

Replacing __import__ has been frowned upon almost as long as it existed. One of the reasons that the import machinery was changed and sys.meta_path introduced was to stop people doing that. – Michelemichelina

@GrahamDumpleton -- Sure. I even state at least one reason why it's bad in my answer (if multiple modules try to replace it, frequently only one will succeed). I also state that I'd be shocked if this is the best way. I'm still waiting to figure out what the best way is though. Perhaps I need to spend some time trying to figure out the source code of wrapt as you propose above... – Barbera

Good way and it works!!!!!!!!!! It's a very wonderful solution and thank you very much~~~ – Aiguillette

What about TPT, how could it be optimized? Also I think it might be better to check 'and list(filter(lambda hook_name: name in hook_name,_post_import_hooks))' so it will also work for sub packages 'x.y.z' or to check 'and name.split('.')[0] in _post_import_hooks' – Osber

Answer 3

Another much simpler but limited approach is to check if numpy has been imported. That only works if at some point in your program all the imports are done but a specific module like numpy may or may not have been imported depending on configuration or the environment. If you need to handle a module being imported at any point later in the program then this won't work.

import sys
if 'numpy' in sys.modules:
    import numpy  # has already been done, so this is free now
    set_linewidth(numpy)

Milch

Answer 4

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')

Lacrimator

No. I don't see how this would work. AFAICT, all this is do is creating a numpy variable. I'm pretty sure the solution involves setting either sys.path_hooks or sys.meta_path. – Bohol

The with as block imports numpy and calls the hook when finished. – Lacrimator

@BallpointBen: But you don't need a hook for that. If you're importing numpy anyway, you can just run the "hook" code after. The goal is to register a hook in case someone imports numpy later, but not to import numpy at all unless someone else imports it. And the code that imports numpy doesn't know about your hook, and can't be changed to add support for it. – Kidskin