python - 导入 python 模块而不实际执行它

Question

在复杂应用程序的上下文中，我需要导入用户提供的“脚本”。理想情况下，脚本应该具有

def init():
    blah

def execute():
    more blah

def cleanup():
    yadda

所以我就这样

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

但是，众所周知，一旦 load_module 运行，用户的脚本就会执行。这意味着，脚本可以是这样的:

def init():
    blah

yadda

一旦我导入脚本，就会调用yadda部分。

我需要的是一种方法:

1. 首先检查是否有init()、execute()和cleanup()
2. 如果它们存在，那就一切顺利
3. 如果它们不存在，请投诉
4. 不要运行任何其他代码，或者至少在我知道没有 init() 之前不要运行

通常情况下，我会强制使用相同的旧 if __name__ == '__main__' 技巧，但我对用户提供的脚本几乎没有控制权，所以我正在寻找一种相对轻松的方法解决方案。我见过各种复杂的技巧，包括解析脚本，但没有什么是真正简单的。我很惊讶它不存在......或者也许我没有得到什么。

谢谢。

Answer 1

我尝试使用ast模块:

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