Friday, November 27, 2009

Injecting attributes into Python modules

One of the things I don't like about frameworks like Django and Pyramid is the amount of boilerplate imports you end up having at the beginning of modules, especially those modules that are more like configuration files. It spoils the DSL-like nature of them, and I was interested in finding a way to be able to import a module with certain attributes already defined. The __import__ function doesn't allow you to do this, because the locals argument is ignored, so I looked for another way. Before I describe the method I came up with, here's how you might use it:
import elixir
from inject import module_inject
module_inject('myapp.models', elixir)
import myapp.models
Easy!
PEP 302 describes the import hooks that have been available since Python 2.3, and defines an import protocol. By adding an object with find_module and load_module methods to sys.meta_path, you can get hooked into the import process. find_module is called with the module name to see if an object knows how to load it. load_module is then called to do the actual loading. The class below implements both of those methods.
class InjectionLoader(object):
    def __init__(self, name, dicts):
        self.name = name
        self.dicts = dicts
    
    def find_module(self, fullname, path=None):
        if fullname == self.name:
            return self
    
    def load_module(self, fullname):
        # Get the leaf module name and the directory it should be found in
        if '.' in fullname:
            package, leaf = fullname.rsplit('.', 1)
            path = sys.modules[package].__path__   
        else:
            leaf = fullname
            path = None

        # Open the module file
        file, filename, description = imp.find_module(leaf, path)

        # Get the existing module or create a new one (for reload to work)
        module = sys.modules.setdefault(fullname, imp.new_module(fullname))
        module.__file__ = filename
        module.__loader__ = self  
        
        code = compile(file.read(), filename, 'exec')
        
        # Populate the module namespace with the injected attributes
        for d in self.dicts:
            module.__dict__.update(d)
            
        # Finally execute the module with its injected attributes
        eval(code, module.__dict__)
        return module
It's instantiated with the module name it's injecting to, and the dicts it is injecting. To make it easier to use, I wrote a helper function, module_inject. It takes a module name, and one or more dicts or modules. Dicts are injected as-is. Modules have their __dict__s injected, but only those attributes listed in the module's __all__ attribute, or if that isn't present then only those that don't begin with a double underscore, are used. This is like doing a from module import * at the beginning of the imported module. Here is its implementation:
def module_inject(name, *args):
    """Set a hook so that when module 'name' is imported, it is executed with
    the attributes in 'args' already in module scope. The arguments can be
    dictionaries or modules (see 'normalize_dict')."""
    args = map(normalize_dict, args)
    sys.meta_path.append(InjectionLoader(name, args))

def normalize_dict(d):
    """If the argument is a module, return the module's dictionary filtered
    by the module's __all__ attribute, otherwise return the argument as-is.
    If the module doesn't have an __all__ attribute, use all the attributes
    that don't begin with a double underscore."""
    if isinstance(d, types.ModuleType):
        keys = getattr(
            d,
            '__all__',
            filter(lambda k: not k.startswith('__'), d.__dict__.keys())
        )
        d = dict([(key, d.__dict__[key]) for key in keys])
    return d
It's something to be used with caution, though. In general, the Python mantra of *explicit is better than implicit* is a good guideline to follow.
Update: somebody asked me about the use of file as a local variable. I'm actually torn on the issue. Yes, it does shadow the built-in file function, but on the other hand it's concise, and it's the same name used in the Python documentation.

No comments:

Post a Comment