导入python模块而不实际执行

小编典典

导入python模块而不实际执行

python

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

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经产生就被调用。

我需要的是一种方法：

首先检查它是否具有init（），execute（）和cleanup（）
如果它们存在，一切都会很好
如果不存在，抱怨
不要运行任何其他代码，或者至少不要运行，直到我知道没有init（）

通常，我会强迫使用相同的旧if __name__ == '__main__'技巧，但是我对用户提供的脚本几乎没有控制权，因此我正在寻找相对轻松的解决方案。我已经看到了各种复杂的技巧，包括解析脚本，但是没有什么真的很简单。我很惊讶它不存在..或者也许我没有得到任何东西。

谢谢。

阅读 215

2020-12-20

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

2020-12-20