Pythonで関数の引数をargparseに自動登録する

はじめに

argparseはpythonで オプション付き引数を持ったコマンドラインスクリプトを作るために便利なパッケージだ。便利なパッケージなんだけど、関数で引数をしっかりと作り、さらにargparseでもう一度、書くということを何十回とやって、この作業無駄だなという思いに駆られた。ということで、自動化してみよう。

Clickとの違い

コメントでclickと同じことをしているんじゃないかと言われて、しばらくまた車輪の再発明をしてしまったと思っていたのだけども、clickを使ってみたところ、どうやら違うようなので追記しておく。clickの機能は、argparseで書くオプションや引数の定義を、デコレータとして各関数に分散して書けるということであって、関数の引数を自動的にオプションとして登録してくれるわけではない。結局、add_argumentの代わりに、@click.optionと書くだけで実は書いている内容はそんなに変わらない。とはいっても、clickがダメだと言っているわけではなく、目的が違うよということだ。整理して書く、可読性を上げるという点においてはclickは素晴らしいパッケージだ。しかし、今回は乱暴に「作った関数の引数をすべて、コマンドラインの引数として自動登録にしてやれ！」という目的のために書いているので、簡単だけど変数が多いスクリプトでは下記の検討はまだ役に立つかもしれない・・。逆にclickを使ってみて、色々と他のシーンで役立ちそうな機能を多く見つけたので、今度、その内容は別途まとめようと思う。

関数の引数解析

関数の引数解析にはinspectというパッケージを使う。python2.7を使っているので、旧式のgetargspecを利用するが、python3なら、getfullargspecを利用すること。

from inspect import getargspec

def func(aaa=1, bbb="a", ccc=[1,2,3,]):
    pass

args, varargs, keywords, defaults = getargspec(func)

print 'args: ', args
print 'varargs', varargs
print 'keywords: ', keywords
print 'defaults: ', defaults

上記を実行すると、以下の出力を得る。

args:  ['aaa', 'bbb', 'ccc']
varargs None
keywords:  None
defaults:  (1, 'a', [1, 2, 3])

argparseへの引数の自動登録

この情報を利用して、argparseのオプションを自動的に登録する。

from inspect import getargspec

def func(aaa=1, bbb="a", ccc=[1,2,3,]):
    pass

args, varargs, keywords, defaults = getargspec(func)

import argparse

parser = argparse.ArgumentParser()
for arg, default in zip(args, defaults):
    parser.add_argument(
        '--' + arg,
        default=default,
        type=type(default)
    )

print parser.parse_args()

上記を実行すると以下のNamespaceオブジェクトが得られる。

Namespace(aaa=1, bbb='a', ccc=[1, 2, 3])

可変長引数への対応

実は上記だと、可変長引数への対応が十分ではない。少し長くなるが可変長引数への対応を追加したコードが以下となる。

from inspect import getargspec

def func(aaa=1, bbb="a", ccc=[1,2,3,]):
    pass

args, varargs, keywords, defaults = getargspec(func)

import argparse

parser = argparse.ArgumentParser()
for arg, default in zip(args, defaults):
    if type(default) is tuple or type(default) is list:
        parser.add_argument(
            '--' + arg,
            default=default,
            type=type(default),
            nargs='+'
        )
    else:
        parser.add_argument(
            '--' + arg,
            default=default,
            type=type(default)
        )

具体的には、引数がリストまたはタプルだった場合に、nargs引数をadd_argumentに追加する。

bool値、enable, disable...への対応

bool値の引数の場合、デフォルトではないとき一意に値が決まるので、True, Falseと指定するのではなく、disable, enableという頭文字を指定して、値をとらないオプションとするのが普通だ。この対応を含めると、以下のようになる。

from inspect import getargspec


def func(aaa=1, bbb="a", ccc=[1,2,3,], ddd=True, eee=False):
    pass

args, varargs, keywords, defaults = getargspec(func)

import argparse

parser = argparse.ArgumentParser()
for arg, default in zip(args, defaults):
    if type(default) is tuple or type(default) is list:
        parser.add_argument(
            '--' + arg,
            default=default,
            type=type(default),
            nargs='+'
        )
    elif type(default) is bool and default:
        parser.add_argument(
            '--disable_' + arg,
            dest=arg,
            default=default,
            action="store_false"
        )
    elif type(default) is bool and not default:
        parser.add_argument(
            '--enable_' + arg,
            dest=arg,
            default=default,
            action="store_true"
        )
    else:
        parser.add_argument(
            '--' + arg,
            default=default,
            type=type(default)
        )

デフォルト値をもたない引数対応

最後になったが、デフォルト値を持たない引数も対応を行う。ただし、この引数は型が不明であるため、文字列として認識されることを覚えておこう。

from inspect import getargspec

def func(x, aaa=1, bbb="a", ccc=[1,2,3,], ddd=True, eee=False):
    pass

args, varargs, keywords, defaults = getargspec(func)

print 'args: ', args
print 'varargs', varargs
print 'keywords: ', keywords
print 'defaults: ', defaults

parser = argparse.ArgumentParser()

while len(args) > len(defaults):
    l = list(defaults)
    l.insert(0, None)
    defaults = tuple(l)

for arg, default in zip(args, defaults):
    if default is None:
        parser.add_argument(dest=arg)
    elif type(default) is tuple or type(default) is list:
        parser.add_argument(
            '--' + arg,
            default=default,
            type=type(default),
            nargs='+'
        )
    elif type(default) is bool and default:
        parser.add_argument(
            '--disable_' + arg,
            dest=arg,
            default=default,
            action="store_false"
        )
    elif type(default) is bool and not default:
        parser.add_argument(
            '--enable_' + arg,
            dest=arg,
            default=default,
            action="store_true"
        )
    else:
        parser.add_argument(
            '--' + arg,
            default=default,
            type=type(default)
        )

print parser.parse_args()

結論

inspectとargparseを用いることで、関数の引数を自動でargparseの引数に設定することができた。最終版のコードを以下に添付する。

extract_arguments.py

#!/usr/bin/env python

import argparse
from inspect import getargspec

def set_arguments(parser, func):
    args, varargs, keywords, defaults = getargspec(func)

    print 'args: ', args
    print 'varargs', varargs
    print 'keywords: ', keywords
    print 'defaults: ', defaults

    while len(args) > len(defaults):
        l = list(defaults)
        l.insert(0, None)
        defaults = tuple(l)

    for arg, default in zip(args, defaults):
        if default is None:
            parser.add_argument(dest=arg)
        elif type(default) is tuple or type(default) is list:
            parser.add_argument(
                '--' + arg,
                default=default,
                type=type(default),
                nargs='+'
            )
        elif type(default) is bool and default:
            parser.add_argument(
                '--disable_' + arg,
                dest=arg,
                default=default,
                action="store_false"
            )
        elif type(default) is bool and not default:
            parser.add_argument(
                '--enable_' + arg,
                dest=arg,
                default=default,
                action="store_true"
            )
        else:
            parser.add_argument(
                '--' + arg,
                default=default,
                type=type(default)
            )
    return parser

def func(x, aaa=1, bbb="a", ccc=[1,2,3,], ddd=True, eee=False):
    pass


parser = argparse.ArgumentParser()
parser = set_arguments(parser, func)
print parser.parse_args()

利便性を兼ねて、デコレータにしたものが以下、

register_arguments.py

#!/usr/bin/env python

import argparse
from inspect import getargspec

def register_arguments(parser, func):
    args, varargs, keywords, defaults = getargspec(func)
    while len(args) > len(defaults):
        l = list(defaults)
        l.insert(0, None)
        defaults = tuple(l)

    for arg, default in zip(args, defaults):
        if default is None:
            parser.add_argument(dest=arg)
        elif type(default) is tuple or type(default) is list:
            parser.add_argument(
                '--' + arg,
                default=default,
                type=type(default),
                nargs='+'
            )
        elif type(default) is bool and default:
            parser.add_argument(
                '--disable_' + arg,
                dest=arg,
                default=default,
                action="store_false"
            )
        elif type(default) is bool and not default:
            parser.add_argument(
                '--enable_' + arg,
                dest=arg,
                default=default,
                action="store_true"
            )
        else:
            parser.add_argument(
                '--' + arg,
                default=default,
                type=type(default)
            )
    return parser

def register_argparse(parser):
    def wrapper(f):
        register_arguments(parser, f)
        return f
    return wrapper

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    @register_argparse(parser)
    def func(x, aaa=1, bbb="a", ccc=[1,2,3,], ddd=True, eee=False):
        pass

    print parser.parse_args()

上記を実行すると以下の出力を出す。

$ python ./register_argments.py a
Namespace(aaa=1, bbb='a', ccc=[1, 2, 3], ddd=True, eee=False, x='a')