Pythonで簡単サブコマンド

こんにちは、zinです🦑

データ基盤まわりのコードを書いていると、
「普段はワークフローエンジンから実行するけど、トラブル時には手元でも動かせるようにしておきたいなぁ」
と思うこと、ありませんか？

そんなとき、標準ライブラリの argparse だけで簡単に実現できます。
argparseは高機能で使いこなそうと思うと奥が深いですが、今回はできるだけシンプルにサクッと使える最小構成を目指します。

さっそく実装編

30行程度でそれなりのものができます。parserに渡す引数は好みに応じて調整してください。
利用可能なオプションの一覧を最初に受け取り、add_handler時に参照している点だけ注意が必要です。

cli_helper.py

import argparse
import typing as t


class SubCommand:
    def __init__(self, options: dict[str, dict[str, t.Any]]):
        self.__parser = argparse.ArgumentParser(
            formatter_class=argparse.RawDescriptionHelpFormatter
        )
        self.__subparsers = self.__parser.add_subparsers()
        self.__options = options

    def add_handler(self, handler: t.Callable, keys: list[str]) -> None:
        cmd = self.__subparsers.add_parser(
            name=handler.__name__,
            description=handler.__doc__,
            help=handler.__doc__,
        )
        for key in keys:
            cmd.add_argument(key, **self.__options[key])
        cmd.set_defaults(handler=handler)

    def call(self) -> int:
        args = vars(self.__parser.parse_args())
        handler = args.pop("handler", None)
        if callable(handler):
            handler(**args)
            return 0
        self.__parser.print_help()
        return 1

サンプルスクリプトを用意

キーワード引数を使っている関数であれば、既存コードに手を加えずにコマンド化できるのがミソです。
サブコマンド間で同じオプションを用意することが多かったので、optionsを事前に定義して使いまわせるようにしました。

main.py

import json

import cli_helper


def sugoi_shori(**kwargs) -> None:
    """これはすごい処理だ、本当にすごいことが起きる"""
    print(json.dumps(kwargs))


def yabai_shori(**kwargs) -> None:
    """これを実行すると大変なことになる"""
    print(json.dumps(kwargs))


if __name__ == "__main__":
    # https://docs.python.org/ja/3/library/argparse.html#the-add-argument-method
    options = {
        "--stage": {
            "required": True,
            "help": "環境",
            "choices": ["prod", "staging"],
        },
        "--force": {"action": "store_true", "help": "強制的に何かをする"},
        "--count": {"default": 0, "type": int, "help": "何かの数"},
    }
    cmd = cli_helper.SubCommand(options)
    cmd.add_handler(sugoi_shori, ["--stage"])
    cmd.add_handler(yabai_shori, ["--stage", "--force", "--count"])
    exit(cmd.call())

実行例

動作例はこちら。

# help表示 docstringがそのままコマンドの説明になります
$ python3 main.py
usage: main.py [-h] {sugoi_shori,yabai_shori} ...

positional arguments:
  {sugoi_shori,yabai_shori}
    sugoi_shori         これはすごい処理だ、本当にすごいことが起きる
    yabai_shori         これを実行すると大変なことになる

options:
  -h, --help            show this help message and exit


# ヤバイ処理を動かしたい
$ python3 main.py yabai_shori
usage: main.py yabai_shori [-h] --stage {prod,staging} [--force]
                           [--count COUNT]
main.py yabai_shori: error: the following arguments are required: --stage


# あっ
$ python3 main.py yabai_shori --stage proda
usage: main.py yabai_shori [-h] --stage {prod,staging} [--force]
                           [--count COUNT]
main.py yabai_shori: error: argument --stage: invalid choice: 'proda' (choose from prod, staging)


# 動きました
$ python3 main.py yabai_shori --stage prod --force
{"stage": "prod", "force": true, "count": 0}

おわりに

薄いラッパー関数ですが、毎回 ArgumentParser インスタンスを作るところから始めると
どうしてもごちゃつきがちな部分を、ほどよく簡潔にまとめられるようになりました。