Importing a python module without actually executing it

Question

In the context of a complex application, I need to import user-supplied 'scripts'. Ideally, a script would have

def init():
    blah

def execute():
    more blah

def cleanup():
    yadda

so I'd just

import imp
fname, path, desc = imp.find_module(userscript)
foo = imp.load_module(userscript, fname, path, desc)
foo.init()

However, as we all know, the user's script is executed as soon as load_module runs. Which means, a script can be something like this:

def init():
    blah

yadda

yielding to the yadda part being called as soon as I import the script.

What I need is a way to:

1. check first whether it has init(), execute() and cleanup()
2. if they exist, all is well
3. if they don't exist, complain
4. don't run any other code, or at least not until I know there's no init()

Normally I'd force the use the same old if __name__ == '__main__' trick, but I have little control on the user-supplied script, so I'm looking for a relatively painless solution. I have seen all sorts of complicated tricks, including parsing the script, but nothing really simple. I'm surprised it does not exist.. or maybe I'm not getting something.

Thanks.

Answer 1

My attempt using the ast module:

import ast

# which syntax elements are allowed at module level?
whitelist = [
  # docstring
  lambda x: isinstance(x, ast.Expr) \
             and isinstance(x.value, ast.Str),
  # import
  lambda x: isinstance(x, ast.Import),
  # class
  lambda x: isinstance(x, ast.ClassDef),
  # function
  lambda x: isinstance(x, ast.FunctionDef),
]

def validate(source, required_functions):
  tree = ast.parse(source)

  functions = set()
  required_functions = set(required_functions)

  for item in tree.body:
    if isinstance(item, ast.FunctionDef):
      functions.add(item.name)
      continue

    if all(not checker(item) for checker in whitelist):
      return False

  # at least the required functions must be there
  return len(required_functions - functions) == 0


if __name__ == "__main__":
  required_funcs = [ "init", "execute", "cleanup" ]
  with open("/tmp/test.py", "rb") as f:
    print("yay!" if validate(f.read(), required_funcs) else "d'oh!")

Answer 2

Here's a simpler (and more naive) alternative to the AST approach:

import sys
from imp import find_module, new_module, PY_SOURCE


EXPECTED = ("init", "execute", "cleanup")

def import_script(name):
    fileobj, path, description = find_module(name)

    if description[2] != PY_SOURCE:
        raise ImportError("no source file found")

    code = compile(fileobj.read(), path, "exec")

    expected = list(EXPECTED)
    for const in code.co_consts:
        if isinstance(const, type(code)) and const.co_name in expected:
            expected.remove(const.co_name)
    if expected:
        raise ImportError("missing expected function: {}".format(expected))

    module = new_module(name)
    exec(code, module.__dict__)
    sys.modules[name] = module
    return module

Keep in mind, this is a very direct way of doing it and circumvents extensions to Python's import machinery.