导入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一旦import脚本完成,就会被调用到被调用的部分.

我需要的是一种方法:

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

Answer 2

这是AST方法的一个更简单(也更天真)的替代方案:

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

请记住,这是一种非常直接的方式,可以绕过Python的导入机制.