Pythonのargparse.ArgumentParser.add_argumentでEnumをchoicesする

動機

PythonでCLIスクリプトを作成するとき、起動時のオプション解釈には argparse を使う。
argparse では引数に制限を加えたい場合、 choices オプションが使える。

from argparse import ArgumentParser
p = ArgumentParser()
p.add_argument('mode', choices=['walk','run','stop'])

p.parse_args(['-h'])
# positional arguments:
#   {walk,run,stop}

p.parse_args(['walk'])
# Namespace(mode='walk')

p.parse_args(['stand'])
# python: error: argument mode: invalid choice: 'stand' (choose from 'walk', 'run', 'stop')

ただ、やはりマジックコードをスクリプトに埋め込むのはよくない。
そこで列挙型である enum と組み合わせたい。

問題

しかし、そのまま何もしないとヘルプの表示が崩れるうえ、そもそもちゃんと解析できない。

from enum import Enum
class ACTION(Enum):
    WALK = 'walk'
    RUN = 'run'
    STOP = 'stop'

from argparse import ArgumentParser
p = ArgumentParser()
p.add_argument('mode', choices=ACTION)

p.parse_args(['-h'])
# positional arguments:
#   {ACTION.WALK,ACTION.RUN,ACTION.STOP} ←変！

p.parse_args(['walk'])
# TypeError: unsupported operand type(s) for 'in': 'str' and 'EnumMeta'

解決策

下記2点で解決する。

• ヘルプの表示を直すために ACTION.__str__ を定義して、オプション文字列のみを返す
• 解析エラーを直すために add_argument に type を指定する

具体的には下記の通り。

from enum import Enum
class ACTION(Enum):
    WALK = 'walk'
    RUN = 'run'
    STOP = 'stop'

    def __str__(self): # これを定義！
        return self.value

from argparse import ArgumentParser
p = ArgumentParser()
p.add_argument('mode', choices=ACTION, type=ACTION) # type を指定！

p.parse_args(['-h'])
# positional arguments:
#   {walk,run,stop} ←いいかんじ！

p.parse_args(['walk'])
# => Namespace(mode=<ACTION.WALK: 'walk'>) ← ちゃんと解析できた！

補足知識

Enum は値で呼び出すと、その要素を返してくれる。
(ドキュメントに書かれていない動作？)
リファレンスではなくてチュートリアルに記載がありました: 列挙型メンバーおよびそれらの属性へのプログラム的アクセス

from enum import Enum
class ACTION(Enum):
    WALK = 'walk'
    RUN = 'run'
    STOP = 'stop'

ACTION('walk')
# => <ACTION.WALK: 'walk'>