如何让模块认为它的名字是__main__？

Question

我有一个python脚本（实际上是很多脚本），其代码仅在此模块作为主脚本运行时执行：

if __name__ == '__main__':
    print("I am the main script")

但是现在我想要一个测试脚本将它们作为模块加载，这样它就可以进入内部状态。重写（将代码块转换为函数）不是一种选择。如何以认为其名称为__main__的方式导入模块？我确定之前已经在某些导入库或其他导入库的帮助下看到了这一点，但我的搜索中没有出现。

Answer 1

您必须绕过导入机制，而是使用exec代替：

import imp

main = imp.new_module('__main__')
with open(module_filename, 'r') as source:
    exec(source.read(), vars(main))

演示：

>>> source = '''\
... if __name__ == '__main__':
...     print("I am the main script")
... '''
>>> import imp
>>> main = imp.new_module('__main__')
>>> exec(source, vars(main))
I am the main script

考虑创建一个你从__main__后卫调用的函数，而不是走这条路，所以你可以导入该函数进行测试：

def main():
    print("I am the main script")

if __name__ == '__main__':
    main()