使用 sys.argv argparse 调用 main

Question

我有一个 main 函数，它使用 argparse 从 sys.argv 获取 args：

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-arg1', '--argument1', type=str, required=True)
    parser.add_argument('-arg2', '--argument2', type=str, required=True)
    args = parser.parse_args()

main 函数是从另一个模块调用的 mod.main()。 当我打印传递给主函数的内容时，参数是： ['-arg1', 'val1', '-arg2', 'val2']这是正确的参数列表。

但是，当我运行它时，args 似乎没有通过 argparse 正确传递。

['-arg1', 'val1', '-arg2', 'val2']
usage: -arg1 [-h] -arg2 ARG2 -arg1 ARG1
-arg1: error: the following arguments are required: -arg1

问题是我没有运行从命令行调用 module_to_call.main() 的 run_module.py，而是将其作为 args 传递给 spark-submit 以便在 EMR 上运行它。

"HadoopJarStep": {
                "Args": [
                    "spark-submit",..., "run_module.py", "module_to_call", "-arg1", "val1", "-arg2", "val2",...

Answer 1

这是给你的一个最小的例子

#!/usr/bin/env python
# argparse_example.py
import sys
import argparse

def main():
    print(sys.argv)
    parser = argparse.ArgumentParser()
    parser.add_argument('-arg1', '--argument1', type=str, required=True)
    parser.add_argument('-arg2', '--argument2', type=str, required=True)
    args = parser.parse_args()
    print(args)

if __name__ == "__main__":
    main()

你可以运行它

python argparse_example.py -arg1 val1 -arg2 val2

并得到

['argparse_example.py', '-arg1', 'val1', '-arg2', 'val2']
Namespace(argument1='val1', argument2='val2')

或者不带参数运行它

python argparse_example.py

获取帮助信息：

['argg.py']
usage: argg.py [-h] -arg1 ARGUMENT1 -arg2 ARGUMENT2
argg.py: error: the following arguments are required: -arg1/--argument1, -arg2/--argument2

更新

从 OP 的评论中，我意识到sys.argv在运行main()的过程中被一些封装代码破坏了，第一个参数被切断了。

参数列表['-arg1', 'val1', '-arg2', 'val2']不是正确的参数列表，因为它缺少sys.argv[0] - 从命令行。 当argparse处理sys.argv[1:]时，它无法找到'-arg1' ，因此失败并出现错误。

在这种情况下，您可以向sys.argv添加一个虚拟元素，以确保在您的代码作为模块运行时将-arg1传递给argparse 。

#!/usr/bin/env python
# argparse_example.py
import sys
import argparse

def main():
    print(sys.argv)
    if __name__ != "__main__":
        # we are running as a module so add a dummy argument
        sys.argv = ['argparse_example.py'] + sys.argv
    parser = argparse.ArgumentParser()
    parser.add_argument('-arg1', '--argument1', type=str, required=True)
    parser.add_argument('-arg2', '--argument2', type=str, required=True)
    args = parser.parse_args()
    print(args)

if __name__ == "__main__":
    main()

您可以使用以下封装代码对其进行测试

#!/usr/bin/env python
# runner.py
import sys
import argparse_example

if __name__ == "__main__":
    sys.argv = sys.argv[1:]
    argparse_example.main()

并像这样从命令行调用它：

python runner.py -arg1 val1 -arg2 val2

结果应该与直接从命令行调用argparse_example时相同，即使sys.argv将缺少第一个元素：

['-arg1', 'val1', '-arg2', 'val2']
Namespace(argument1='val1', argument2='val2')