CODEWITHSUNDEEP
pythonpython-3.xpython-import
Mark Lodato
Mark Lodato

Reputation: 53274

Post import hooks in Python 3

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?

Upvotes: 10

Views: 2237

Answers (4)

JanKanis
JanKanis

Reputation: 6702

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)

Upvotes: 0

mgilson
mgilson

Reputation: 310117

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.

Upvotes: 7

BallpointBen
BallpointBen

Reputation: 13867

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

Upvotes: -3

Graham Dumpleton
Graham Dumpleton

Reputation: 58563

The wrapt module provides an implementation of this.

Watch this video about wrapt, including this feature:

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

Some of the blogs posts at end of:

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.

Upvotes: 7

Related Questions