Python argparse の action API を使ってみる

Pythonで引数を取得するのに便利なargparseモジュール。
add_argumentにactionを指定すると、引数があった時どのように値を格納するかを割り当てられる。
たぶん、よく使われるのはstore_trueを指定して単にフラグを立てるのに使うとか。

>>> import argparse
>>> parser = argparse.ArgumentParser()
>>> parser.add_argument("--hoge", action="store_true")
>>> args = parser.parse_args(["--hoge"])
>>> args.hoge
True

で、このactionにAPIを実装したオブジェクトを渡すと任意の下ごしらえができるそうなので使ってみた。

argparse.Action

argparseのドキュメントに従ってargparse.Actionを継承し、__call__メソッドを持つクラスを書く。

print_action.py

import argparse

class PrintAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        print "[{}] option add value [{}] as attr [{}]".format(
                option_string, values, self.dest)
        setattr(namespace, self.dest, values)

__call__メソッドでは次の4つの引数を受け取るようにする。

• parser : add_argument する ArgumentParser オブジェクト
• namespace : parse_args() で返される namespace オブジェクト
• values : 引数の値。add_argumentのtypeが指定されていた場合は先に型変換される。
• option_string : オプション引数な場合のオプション文字列。位置引数の場合は渡されない。

上記のPrintActionを使ってみる。

>>> import argparse
>>> from print_action import PrintAction
>>>
>>> parser = argparse.ArgumentParser()
>>> parser.add_argument("--hoge", action=PrintAction)
>>> args = parser.parse_args("--hoge huga".split(' '))
[--hoge] option add value [huga] as attr [hoge]
>>> args.hoge
'huga'

基本的にはvaluesを適当に加工した上で、引数名(self.dest)を属性名に使ってnamespaceに追加すればよい。

追記: 最近のドキュメントには、必要であれば__init__もオーバーライドするよう書いてある。argparse.Action.__init__では、「2つの位置引数と、action それ自身を除く ArgumentParser.add_argument() に渡されるすべてのキーワード引数を受け付けなければなりません。」と書いてあるが、要するにadd_argumentにおけるactionの使い方が適切かどうか__init__内でチェックできるようだ。

例

引数として整数、もしくは<int>[KMG]を取り、後者の場合は整数に変換するaction。

calc_suffix_action.py

import argparse

class CalcSuffixAction(argparse.Action):
    def __init__(self, option_strings, dest, nargs=None, **kwargs):
        """
        省略可。`add_argument`の他の引数がactionにとって適切か判定するならここで
        今回は`__call__`メソッドで受け取る`values`がリストでないことを期待しているので
        `nargs`を制限してみる。
        """
        if nargs is not None and nargs != '?':
            raise ValueError("Invalid `nargs`: multiple arguments not allowed")
        super(CalcSuffixAction, self).__init__(option_strings, dest, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        try:
            if values.endswith('K'):
                digit = int(values[:-1]) * 1000
            elif values.endswith('M'):
                digit = int(values[:-1]) * 1000000
            elif values.endswith('G'):
                digit = int(values[:-1]) * 1000000000
            elif not isinstance(values, int):
                digit = int(values)
            else:
                digit = values

        except ValueError:
            parser.error("Invalid argument ({})".format(values))

        setattr(namespace, self.dest, digit)

使ってみる。

>>> import argparse
>>> from calc_suffix_action import CalcSuffixAction
>>>
>>> parser = argparse.ArgumentParser()
>>> parser.add_argument("--hoge", action=CalcSuffixAction)
>>>
>>> parser.parse_args("--hoge 1024".split(' '))
Namespace(hoge=1024)
>>> parser.parse_args("--hoge 4K".split(' '))
Namespace(hoge=4000)
>>> parser.parse_args("--hoge 32M".split(' '))
Namespace(hoge=32000000)
>>> parser.parse_args("--hoge 1G".split(' '))
Namespace(hoge=1000000000)
>>> parser.parse_args("--hoge 3T".split(' '))
usage: [-h] [--hoge HOGE]
: error: Invalid argument (3T)