将用于命令行的 Python 脚本中的 __main__ 函数转换为可调用函数

Question

我正在使用某人的 GitHub 代码，该代码旨在从命令行调用，如下所示：

> python this_script.py -u <username> -p <password> -i <id_num> ...

这会在this_script.py 中以这种方式解析输入后生成一个输出文本文件：

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='get data', add_help=False, usage='this_script.py -u username -p password [options]')
    parser.add_argument('-u', metavar='<STR>', type=str, help='username')
    parser.add_argument('-p', metavar='<STR>', type=str, help='password')
    parser.add_argument('-i', metavar='<STR>', nargs='+', type=str, help='List of IDs')
    ...

我想在我自己的代码集中使用其中的许多文本文件，因此将其从命令行脚本转换为可调用函数会更方便，语法如下：

def this_script(password, username, *args):
    ...

有没有一种简单的方法可以做到这一点，而不必过多地处理当前定义的 __main__ 块的内部工作？

Answer 1

怎么样：

def do_stuff(args):
    print(args.u)
    print(args.p)


def main(args_list=None):
    parser = argparse.ArgumentParser(description='get data', add_help=False, usage='this_script.py -u username -p password [options]')
    parser.add_argument('-u', metavar='<STR>', type=str, help='username')
    parser.add_argument('-p', metavar='<STR>', type=str, help='password')
    parser.add_argument('-i', metavar='<STR>', nargs='+', type=str, help='List of IDs')
    ...

    if args_list:
        args= parser.parse_args(args_list)
    else:
        args = parser.parse_args()

    do_stuff(args)


if __name__ == "__main__":
    main()

通过将以下行添加到您的包中 __init__.py

from .__main__ import main

您可以从命令行调用脚本作为函数，也可以从另一个脚本中调用，只要您导入包即可，例如

import <package name> as pkg
pkg.main(['-u', 'my_username', '-p', 'my_password'])

Answer 2

是的，简单的方法是使用stdlib runpy。

import runpy

runpy.run_module(mod_name, init_globals=None, run_name=None, alter_sys=False)

并且，请记住不要这样编写自己的脚本。条件下应该只有一行，并且应该是对不带参数的函数的调用，如下所示：

if __name__ == "__main__":
    main()

任何东西都太多了。