argparseでsubcommandを作るためのユーティリティ

argparse python

サブコマンドを作る時に何らかのライブラリに依存して良いならclickがオススメではあるけれど。 使いたくない場合もあったりする。そういう時にどうするとまだましになるかみたいな事を考えたりしていた。

argparseでのサブコマンド

argparseでのサブコマンドの定義の仕方は以下のような感じ。

def foo():
    print("foo")


def bar():
    print("bar")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="subcommand")
    subparsers.required = True

    foo_parser = subparsers.add_parser("foo")
    foo_parser.set_defaults(fn=foo)

    bar_parser = subparsers.add_parser("bar")
    bar_parser.set_defaults(fn=bar)

    args = parser.parse_args()
    args.fn()

こうするとfooとbarが使える

$ python 00subcommands.py  -h
usage: 00subcommands.py [-h] {foo,bar} ...

positional arguments:
  {foo,bar}

optional arguments:
  -h, --help  show this help message and exit

引数の数が違う場合

引数の数が違う場合にちょっと大変になる。

例えばこういう風にfooではx,yでbarにはx,y,zの引数が用意されていた場合など。

    foo_parser = subparsers.add_parser("foo")
    foo_parser.add_argument("-x")
    foo_parser.add_argument("-y")
    foo_parser.set_defaults(fn=foo)

    bar_parser = subparsers.add_parser("bar")
    bar_parser.add_argument("-x")
    bar_parser.add_argument("-y")
    bar_parser.add_argument("-z")
    bar_parser.set_defaults(fn=bar)

if文とかで頑張らないとだめ。

if args.fn == foo:
    return foo(args.x, args.y)
elif args.fn == bar:
    return bar(args.x, args.y, args.z)

ユーティリティ

以下の様なユーティリティを書くとマシかもしれない。面倒なのは parser.parse_args() で手に入った Namespace オブジェクトの取扱い。 add_argumentで返されるactionのdestの値がシステム名的なものなのでそれを利用する関数を登録する。

クラスなどにしてラップしても良いけれど。contextmanagerを使うのが楽な気がした。

import contextlib

@contextlib.contextmanager
def subparser(subparsers, fn, *args, **kwargs):
    parser = subparsers.add_parser(fn.__name__, *args, **kwargs)
    dests = []
    arrived = set()

    def add_argument(*args, **kwargs):
        ac = parser.add_argument(*args, **kwargs)
        if ac.dest not in arrived:
            arrived.add(ac.dest)
            dests.append(ac.dest)
        return ac

    yield add_argument

    def run(args):
        return fn(**{name: getattr(args, name) for name in dests})

    parser.set_defaults(fn=run)

こんな感じで使う。

def foo(*, x, y):
    print("foo", x, y)


def bar(*, x, y, z):
    print("bar", x, y, z)


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_subparsers
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="subcommand")
    subparsers.required = True

    with subparser(subparsers, foo) as add_argument:
        add_argument("-x", default=None)
        add_argument("-y", default=None)

    with subparser(subparsers, bar) as add_argument:
        add_argument("-x", default=None)
        add_argument("-y", default=None)
        add_argument("-z", default=None)

    args = parser.parse_args()
    return args.fn(args)

$ python 02* foo -h
usage: 02withutil.py foo [-h] [-x X] [-y Y]

optional arguments:
  -h, --help  show this help message and exit
  -x X
  -y Y
$ python 02* bar -h
usage: 02withutil.py bar [-h] [-x X] [-y Y] [-z Z]

optional arguments:
  -h, --help  show this help message and exit
  -x X
  -y Y
  -z Z

# 引数多すぎ
$ python 02* foo -x 10 -y 20 -z 30
usage: 02withutil.py [-h] {foo,bar} ...
02withutil.py: error: unrecognized arguments: -z 30

# ok
$ python 02* foo -x 10 -y 20
foo 10 20

補足

foo,barの引数の定義にkeyword only argumentsを使うと良い。引数の数が会わない時のエラーがわかりやすいので。

例えば、fooにも -z オプションを追加してしまった場合には以下の様になる。

$ python 02* foo
Traceback (most recent call last):
  File "02withutil.py", line 56, in <module>
    main()
  File "02withutil.py", line 52, in main
    return args.fn(args)
  File "02withutil.py", line 20, in run
    return fn(**{name: getattr(args, name) for name in dests})
TypeError: foo() got an unexpected keyword argument 'z'

あー。zオプションなんて定義していなかったと気付ける。あと引数の順序を気にする必要がなくなるし。