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!")