4
\$\begingroup\$

I'm working on an open-source test framework which needs to dynamically import Python source modules from a given directory.

I can't just use importlib.import_module() or __import__(), because it's a black box and I need to do some source-code-rewriting. On the other hand, I feel the PEP 302 Import Hooks to be overkill for what I'm trying to do (they seem to be designed with importing totally foreign file formats in mind), so I rolled my own:

import os
import sys
import types

def import_module(dir_path, module_name):
    if module_name in sys.modules:
        return sys.modules[module_name]

    filename = resolve_filename(dir_path, module_name)

    with open(filename, 'r', encoding='utf-8') as f:
        source = f.read()
    # Here I do some AST munging    
    code = compile(source, filename, 'exec')

    module = create_module_object(module_name, filename)
    exec(code, module.__dict__)
    sys.modules[module_name] = module
    return module

def resolve_filename(dir_path, module_name):
    filename = os.path.join(dir_path, *module_name.split('.'))

    # I happen to know that the calling code will already have
    # determined whether it's a package or not
    if os.path.isdir(filename):
        filename = os.path.join(filename, '__init__.py')
    else:
        filename += '.py'
    return filename    

def create_module_object(module_name, filename):
    module = types.ModuleType(module_name)
    module.__file__ = filename

    if '__init__.py' in filename:
        module.__package__ = module_name
        module.__path__ = [os.path.dirname(filename)]
    else:
        if '.' in module_name:
            module.__package__, _ = module_name.rsplit('.', 1)
        else:
            module.__package__ = ''

    return module

To paraphrase Greenspun's Tenth Rule:

Any sufficiently complicated Python program contains an ad hoc, informally-specified, bug-ridden, slow implementation of half of the import system.

So, does my code do everything it's supposed to? I wrote my tests against an implementation that used importlib.import_module() and then rewrote it as this one, which gives me some confidence, but I'm not sure whether I've missed any important steps.

\$\endgroup\$
  • \$\begingroup\$ Did you consider subclassing importlib.SourceLoader and providing an implementation of the get_code method? This looks to be exactly what you need. \$\endgroup\$ – Gareth Rees Dec 10 '13 at 22:44
  • \$\begingroup\$ @GarethRees Could you elaborate on the advantages of that? \$\endgroup\$ – Benjamin Hodgson Dec 10 '13 at 22:48
  • \$\begingroup\$ You'd only have to implement the bit of the import process that you want to change (not the whole of it as at present) and this would give you more confidence that what you had done was correct. \$\endgroup\$ – Gareth Rees Dec 10 '13 at 23:00
  • \$\begingroup\$ @GarethRees But wouldn't I also have to write a 'finder' and install it onto sys.meta_path? Seems like I'd be trading simplicity in one part for complexity in another. \$\endgroup\$ – Benjamin Hodgson Dec 10 '13 at 23:17
  • 1
    \$\begingroup\$ That's what you have to do if you want import foo to call your code. But if you're prepared to call your code directly, you can shortcut some of the steps. \$\endgroup\$ – Gareth Rees Dec 10 '13 at 23:40
4
\$\begingroup\$

I took @GarethRees's advice and subclassed importlib.abc.SourceLoader. After reading importlib's source code, I arrived at the following implementation:

import importlib.abc

def import_module(dir_path, module_name):
    filename = resolve_filename(dir_path, module_name)

    # I got test failures when I left this out, even though I thought it
    # was a responsibility of the loader.
    # If you know why this is, please enlighten me!
    if module_name in sys.modules:
        return sys.modules[module_name]

    return ASTFrobnicatingLoader(module_name, filename).load_module(module_name)


# in importlib, this function would be the job of the 'finder'
def resolve_filename(dir_path, module_name):
    filename = os.path.join(dir_path, *module_name.split('.'))
    if os.path.isdir(filename):
        filename = os.path.join(filename, '__init__.py')
    else:
        filename += '.py'
    return filename    


class ASTFrobnicatingLoader(importlib.abc.FileLoader, importlib.abc.SourceLoader):
    def get_code(self, fullname):
        source = self.get_source(fullname)
        path = self.get_filename(fullname)

        parsed = ast.parse(source)
        self.frobnicate(parsed)

        return compile(parsed, path, 'exec', dont_inherit=True, optimize=0)

    def module_repr(self, module):
        return '<module {!r} from {!r}>'.format(module.__name__, module.__file__)

import_module now creates an instance of my custom loader and calls its load_module template method, which is provided by the abstract base class. By inheriting from both FileLoader and SourceLoader, I get a 'free' implementation of the whole protocol and I only need to override get_code.

In Python 3.4, SourceLoader provides a source_to_code method which you can override. It would've been ideal for my purposes (because the only thing I'm customising is the generation of the code object) but sadly I'm stuck with having to override the whole of get_code and calling get_source and get_filename manually.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.