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

在我 import 脚本后立即调用 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!")