如何使用 Python argparse 延迟加载子命令

Question

我有一个Python包，里面有很多子命令和子子命令。它的组织方式有点像这样：

main.py

import argparse

from sum import prepare_arg_parser as prepare_sum_parser
from sub import prepare_arg_parser as prepare_sub_parser

parser = argparse.ArgumentParser()
sub_parsers = parser.add_subparsers(dest="command", required=True)

prepare_sum_parser(sub_parsers.add_parser("sum"))
prepare_sub_parser(sub_parsers.add_parser("sub"))

args = parser.parse_args()
args.func(args)

sum.py

import argparse


def prepare_arg_parser(parser):
    parser.set_defaults(func=do_sum)
    parser.add_argument("a", type=int)
    parser.add_argument("b", type=int)


def do_sum(args):
    print(f"{args.a} + {args.b} = {args.a + args.b}")

等等。

但是，随着模块列表的增长，启动时间也会增加。 我想延迟加载每个子命令，因此它仅加载所选命令所需的内容。

我尝试过这样的事情：

import argparse

def do_sum(args):
    from sum import prepare_arg_parser
    parser = argparse.ArgumentParser()
    prepare_arg_parser(parser)
    sum_args = parser.parse_args(args)
    sum_args.func(sum_args)

parser = argparse.ArgumentParser()
sub_parsers = parser.add_subparsers(dest="command", required=True)

sub_parsers.add_parser("sum").set_defaults(func=do_sum)

args, rest = parser.parse_known_args()
args.func(rest)

但是帮助和错误文本不正确。

另一种方法是将每个子命令拆分为一个“prepare argparse”模块，该模块会知道所有参数，但会延迟调用 do_sum 函数，但我不喜欢它，因为它将参数声明与其使用分开。而且它仍然会加载很多不必要的文件。

是否有延迟加载的子命令？

Answer 1

简单的方法

与导入许多模块的成本相比，创建解析器和子解析器的成本很便宜。我解决这个问题的方法是在

main.py

中创建解析器和子解析器，并将实际操作留在模块中。

让我们从主模块开始：

# main.py
import argparse
import importlib


def parse_command_line():
    parser = argparse.ArgumentParser()
    sub_parsers = parser.add_subparsers(dest="command", required=True)

    sum_parser = sub_parsers.add_parser("sum")
    sum_parser.add_argument("a", type=int)
    sum_parser.add_argument("b", type=int)

    return parser.parse_args()


def main():
    args = parse_command_line()
    command = importlib.import_module(args.command)
    command.main(args)


if __name__ == "__main__":
    main()

在本模块中，我们创建了解析器和一个子解析器（sum）。我们可以稍后添加更多子解析器。

解析命令行后，

args.command

包含解析器的名称，例如“和”。计划是有一个具有该名称的模块，即

sum.py

。在该模块中，我们应该有一个函数

main()

来执行以下操作：

# sum.py
import argparse

def main(args: argparse.Namespace):
    print(f"{args.a} + {args.b} = {args.a + args.b}")

我喜欢这种方法的原因：

• 我们不会导入所有模块，只导入我们需要的模块。这意味着加载时间很快。
• 所有命令行处理都集中在一处，使其更易于阅读和故障排除。
• 随着我们添加更多模块，只有

  parse_command_line()

  功能会增长，

  main()

  功能保持不变。
• 每个模块中的代码更简单：每个模块只需要实现一个

  main()

  功能而不用担心解析命令行。

我不喜欢的：

• 当我们添加新模块时，我们需要在两个地方进行编辑：模块本身和

  parse_command_line()

  函数
• 如果我们只有几个模块，这种方法在性能增益方面并不能提供太多

插件方法

在这种方法中，让我们建立一个简单的“插件”系统。除了主模块之外，我们还有插件，每个插件都会满足这些要求：

• 文件名的格式应为

  sub_<name>.py

  。例如：

  sub_sum.py

  、

  sub_dbl.py

  、...
• 每个插件都会解析命令行并执行一些操作。

让我们从主模块开始：

# main.py
import argparse
import pathlib
import subprocess
import sys


def parse_command_line():
    # Determine the choices
    main_dir = pathlib.Path(__file__).parent
    choices = [
        path.stem.removeprefix("plugin_") for path in main_dir.glob("plugin_*.py")
    ]

    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=choices)

    return parser.parse_known_args()


def main():
    args, remainder = parse_command_line()
    command = [sys.executable, f"plugin_{args.command}.py"] + remainder
    subprocess.run(command)


if __name__ == "__main__":
    main()

以下是“sum”插件的内容：

# plugin_sum.py
import argparse


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("a", type=int)
    parser.add_argument("b", type=int)
    args = parser.parse_args()

    print(f"{args.a} + {args.b} = {args.a + args.b}")


if __name__ == "__main__":
    main()

那么它是如何运作的呢？主模块仅收集插件的名称并使用

subprocess.run()

调用正确的插件。每个插件都是一个独立的脚本，能够独立运行。

我喜欢什么：

• 添加更多模块（插件）不需要更改主模块
• 每个插件都是独立的脚本，更容易测试

我不喜欢的：

• 对于主脚本，

  --help

  标志没有提供太多帮助，因为它不了解插件
• 每个插件都必须自己进行命令行解析

结论

我更喜欢简单的方法，因为它简单。插件方法对于更复杂的结构更有效。这两种方法都不会加载所有模块才能工作。