Willi Ballenthin
Willi Ballenthin

Reputation: 6604

How can I exec a file and provide hooked imports in Python 3?

I'm using Python 3 to emulate the Python scripting interface provided by another tool. (When scripting this tool, you run Python scripts, which can do import mytool to access the scripting API.)

I've implemented the interfaces exposed by this tool, and would like to write a loader that you'd invoke such as:

python run_mytool_script.py dump_menu_items.py

This loader would allow you to interact with some portion of the tool's functionality without the tool actually being installed. Ideally, this should allow me to run existing scripts designed for the tool without modification.

Within run_mytool_script.py I expect to:

  1. initialize the emulated scripting interface
  2. prepare import hooks
  3. exec the script dump_menu_items.py

However, I can't quite figure out how to create the import hooks. How can I install a hook so that my emulated scripting interface is exposed as mytool once the script does import mytool?

Note, the emulated scripting interface must be initialized at runtime, so installing a package named mytool doesn't do the trick.

Upvotes: 0

Views: 601

Answers (2)

Willi Ballenthin
Willi Ballenthin

Reputation: 6604

Following the example from the cPython tests here, I've come up with the following potential solution. I'll update this post as I realize its pros and cons.

class HookedImporter(importlib.abc.MetaPathFinder, importlib.abc.Loader):
    def __init__(self, hooks=None):
        self.hooks = hooks

    def find_spec(self, name, path, target=None):
        if name not in self.hooks:
            return None

        spec = importlib.util.spec_from_loader(name, self)
        return spec

    def create_module(self, spec):
        # req'd in 3.6
        logger.info('hooking import: %s', spec.name)
        module = importlib.util._Module(spec.name)
        mod = self.hooks[spec.name]
        for attr in dir(mod):
            if attr.startswith('__'):
                continue
            module.__dict__[attr] = getattr(mod, attr)
        return module

    def exec_module(self, module):
        # module is already loaded (imported by line `import idb` above),
        # so no need to re-execute.
        #
        # req'd in 3.6.
        return

    def install(self):
        sys.meta_path.insert(0, self)

... somewhere later on ...

api = mytool.from_config(...)
hooks = {
    'mytool': api.mytool,
}

importer = HookedImporter(hooks=hooks)
importer.install()


with open(args.script_path, 'rb') as f:
    g = {
        '__name__': '__main__',
    }
    g.update(hooks)
    exec(f.read(), g)

Upvotes: 0

zwer
zwer

Reputation: 25799

Well, there are a couple of ways to do that. Let's start with the most complex one - a fully dynamic creation of mytool. You can use the imp module to create a new module, then define its structure and, finally, add it to the global module list so that everything running in the same interpreter stack can import it:

run_mytool_script.py:

import imp
import sys

# lets first deal with loading of our external script so we don't waste cycles if a script
# was not provided as an argument
pending_script = None  # hold a script path we should execute
if __name__ == "__main__":  # make sure we're running this script directly, not imported
    if len(sys.argv) > 1:  # we need at least one argument to tell us which script to run
        pending_script = sys.argv[1]  # set the script to run as the first argument
    else:
        print("Please provide a path for a script to run")  # a script was not provided
        exit(1)

# now lets create the `mytool` module dynamically
mytool = imp.new_module("mytool")  # create a new module

def hello(name):  # define a local function
    print("Hello {}!".format(name))
mytool.__dict__["hello"] = hello  # add the function to the `mytool` module

sys.modules["mytool"] = mytool  # add 'mytool' to the global list so it can be imported

# now everything running in the same Python interpreter stack is able to import mytool
# lets run our `pending_script` if we're running the main script
if pending_script:
    try:
        execfile(pending_script)  # run the script
    except NameError:  # Python 3.x doesn't have execfile, use the exec function instead
        with open(pending_script, "r") as f:
            exec(f.read())  # read and run the script

Now you can create another script that trusts you that there is a mytool module when being run through your proxy script, e.g.:

dump_menu_items.py:

import mytool
mytool.hello("Dump Menu Items")

And now when you run it:

$ python run_mytool_script.py dump_menu_items.py
Hello Dump Menu Items!

While neat, this is a rather tedious task - I'd suggest you to create a proper mytool module in the same folder where run_mytool_script.py is, have run_mytool_script.py initialize everything that's needed and then just use the last portion of the script to run your external script - that's much easier to manage and generally a much nicer approach.

Upvotes: 2

Related Questions