コマンドライン引数によって、オプション引数が変わるコマンドの書き方のメモ(2 phase parse?)

コマンドライン引数によって、オプション引数が変わる?

例えば --fn というオプションを取り、このオプションが fのときとgのときとで引数が変わるようなもの

--fn無し
--fn f
--fn g

# --fn 無し
$ python 03cli.py
usage: 03cli.py [--verbose] [--name NAME] --fn {f,g}

optional arguments:
  --verbose
  --name NAME
  --fn {f,g}
03cli.py: error: the following arguments are required: --fn

# --fn f
$ python 03cli.py --fn f
usage: f [-h] [--z Z] [--verbose VERBOSE] [--name NAME] --y Y

optional arguments:
  -h, --help         show this help message and exit
  --z Z
  --verbose VERBOSE  (default=True)
  --name NAME
  --y Y
f: error: the following arguments are required: --y

# --fn g
$ python 03cli.py --fn g
usage: g [-h] [--z Z] [--verbose VERBOSE] [--name NAME] --i I

optional arguments:
  -h, --help         show this help message and exit
  --z Z              (default=100)
  --verbose VERBOSE  (default=True)
  --name NAME
  --i I
g: error: the following arguments are required: --i

f のときには、--y オプションが、g のときには --i オプションが現れる。こういうコマンドを作る方法についてのメモ

以下の記事の続きでもある。

functools.partialの元を辿る - podhmo's diary

2 phase parse

基本的には、コマンドライン引数を２回parseする。どうやるのかと言うと、ArgumentParser.parse_known_argsを使う。あんまり一般的なメソッドではないかもしれないけれど。通常のparse_args()的な挙動に加えて、余った残りの引数も返してくれる。あとはもう一度別のparserでparseしてあげれば良い。

args, rest_argv = parser.parse_known_args(argv)

subargs = subparser.parse_args(rest_argv)

注意点として２段目のparserが動的に変わるという点。

コード

コードは以下の様な感じ（後で詳細を説明する)。

import typing as t
import inspect
import itertools
from functools import partial
from collections import ChainMap
import argparse


def find_original_with_arguments(fn):
    args = ()
    kwargs = {}
    if isinstance(fn, partial):
        args = fn.args
        kwargs = fn.keywords
        fn = fn.func
    if inspect.isclass(fn):
        fn = fn.__init__
    return fn, args, kwargs


def _make_parser(*args, **kwargs):
    parser = argparse.ArgumentParser(*args, **kwargs)
    parser.print_usage = parser.print_help
    return parser


def main(argv: t.Optional[t.List[str]] = None) -> None:
    parser = _make_parser(add_help=False)
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--name", required=False)
    parser.add_argument("--fn", required=True, choices=[x.__name__ for x in [f, g]])
    args, rest_args = parser.parse_known_args(argv)

    cli_kwargs = vars(args).copy()
    fn = globals()[cli_kwargs.pop("fn")]

    original, fn_args, default_kwargs = find_original_with_arguments(fn)
    spec = inspect.getfullargspec(original)
    second_parser = _make_parser(fn.__name__)
    seen = set()
    for name, default, required in itertools.chain(
        ((k, v, False) for k, v in spec.kwonlydefaults.items() or ()),
        ((k, None, True) for k in spec.kwonlyargs or ()),
    ):
        if name in seen:
            continue
        seen.add(name)

        clitype = None
        typ = spec.annotations.get(name)
        if typ is not None:
            if typ in (int, float):
                clitype = typ  # todo: supporting only int,

        default = default_kwargs.get(name) or default
        help_message = None
        if default is not None:
            help_message = f"(default={default_kwargs.get(name) or default!r})"

        if name in cli_kwargs:
            default = cli_kwargs[name]
            required = False

        second_parser.add_argument(
            f"--{name.replace('_', '-')}",
            type=clitype,
            required=required,
            default=default,
            help=help_message
        )

    second_args = second_parser.parse_args(rest_args)
    return original(*fn_args, **ChainMap(default_kwargs, cli_kwargs, vars(second_args)))


def f(*, name: str, y: str, z: t.Optional[str] = None, verbose: bool = True):
    r = (name, y, z)
    return r, [type(x) for x in r]


def g(*, name: str, i: int, z: t.Optional[int] = 100, verbose: bool = True):
    r = (name, i, z)
    return r, [type(x) for x in r]

if __name__ == "__main__":
    main()

詳細

関数に対するディスパッチャ

コマンドライン引数の文字列を直接関数にマッピングするのには globals()[fnname] が便利。

globals()["f"] # => 定義されている関数fが返る

functools.partialの扱い

さっき書いた

オプション引数のデフォルトと型

型アノテーションがついていれば、良い感じにマッピングできるかもしれない。上の例ではint,floatにしか対応していないけれど。inspect.getfullargspec()で取り出した値を見れば良い感じにできそうな感じはする。

デフォルト値は関数のそれと、functools.partial()に対応したそれを見る。ついでにヘルプメッセージにdefault値がわかるなら表示してあげられると親切。

optional arguments:
  -h, --help         show this help message and exit
  --z Z              (default=100)  # <- これ
  --verbose VERBOSE  (default=True)  # <- これ
  --name NAME
  --i I

1段目のオプションで既に使われたオプションを２段目でも使う

ヘルプメッセージの表示を考えると、関数の引数として使われているものは、一段目のparserで指定されたものであってもhelpに表示されて欲しい(例えば上の例で言う--nameなどがそう)。

optional arguments:
  -h, --help         show this help message and exit
  --z Z              (default=100)
  --verbose VERBOSE  (default=True)
  --name NAME  # <- これ
  --i I

requiredをfalseにしてdefault値として指定してあげれば、parse後の値として入る。

`-h` でのヘルプメッセージを抑制する

--fn f -h などの場合には２段目のparserのヘルプメッセージを表示してあげたい。これはshow_helpをFalseにしてあげれば良い。

$ python 03cli.py -h
usage: 03cli.py [--verbose] [--name NAME] --fn {f,g}

optional arguments:
  --verbose
  --name NAME
  --fn {f,g}
03cli.py: error: the following arguments are required: --fn

--fn f がある場合のヘルプメッセージ。

$ python 03cli.py -h --name f
usage: f [-h] [--z Z] [--verbose VERBOSE] [--name NAME] --y Y

optional arguments:
  -h, --help         show this help message and exit
  --z Z
  --verbose VERBOSE  (default=True)
  --name NAME
  --y Y

失敗した時のヘルプメッセージを豪華に

parser.print_usage = parser.print_help

というhack。以下の様な簡潔な表示より。

usage: 03cli.py [--verbose] [--name NAME] --fn {f,g}
03cli.py: error: the following arguments are required: --fn

冗長な表現の方が嬉しいときもある。

usage: 03cli.py [--verbose] [--name NAME] --fn {f,g}

optional arguments:
  --verbose
  --name NAME
  --fn {f,g}
03cli.py: error: the following arguments are required: --fn