Pythonのargparseでrsyncの--include/--excludeみたいなオプションを実装する

Motivation

Pythonのコマンドライン引数処理はargparseモジュールが優秀で大抵は事足りるのだけれど、rsyncの--includeと--excludeのような複数回指定できて、かつ書かれた順番も取得したいオプションを実装したくなった。
例えば与えられた文字列の集合をワイルドカードを使ったフィルターで取捨選択する例を考える。

names.txt

Aaron
Abel
Adam
Adolph
Adrian
Alan
Albert
...

この時、--include "A*" --exclude "?d*" --include "*m"というオプションが与えられたら次のような処理をしたい。

1. Aから始まる文字列を対象にする
2. Aから始まる文字列でも2文字目がdな文字列は対象外にする
3. さらにその対象外にされた文字列の中でmで終わる文字列は対象にする

※ rsyncの挙動とは異なります

そうすると、--includeと--excludeオプションは複数回出現することが許されなければならない上に書かれた順番も取得する必要がある。前者についてはadd_argumentのactionに"append"を指定すれば実現できるが、与えられた引数はオプションごとにまとめられてしまうため、--includeと--excludeがどのような順序で指定されたかという情報は失われてしまう。

Solution

actionを自分で定義して、同じ属性にラベル付きで格納する。
argparse.Actionについてはこちら。

filter.py

import argparse


def make_multistate_append_action(key):
    class _MultistateAppendAction(argparse.Action):
        def __call__(self, parser, namespace, values, option_string=None):
            args = getattr(namespace, self.dest)
            args = [] if args is None else args
            args.append((key, values))
            setattr(namespace, self.dest, args)

    return _MultistateAppendAction

この関数は、与えられたコマンドライン引数をkeyと共にリストに格納するactionを返す。

filter.py

parser = argparse.ArgumentParser()

parser.add_argument("inputfile")
parser.add_argument("-i", "--include", nargs='+', dest="filters",
                    action=make_multistate_append_action(True))
parser.add_argument("-e", "--exclude", nargs='+', dest="filters",
                    action=make_multistate_append_action(False))

args = parser.parse_args()

ポイントは--includeと--excludeのdestに同じ値を設定すること。意外にもArgumentParserはdestが重複するオプションを追加してもエラーにはならない。
この状態でargsを確認するとこんな感じ。

$ python3 test.py names.txt --include "A*" --exclude "?d*" --include "*m"
Namespace(filters=[(True, ['A*']), (False, ['?d*']), (True, ['*m'])], inputfile='names.txt')

Example

フィルターまで実装してみるとこんな感じ。reを使うほどでもないけどワイルドカードは扱いたいという時にはfnmatchモジュールがとても便利。
この実装だと重複のあるリストにはそのまま使えないのでご勘弁。

filter.py

import argparse
import fnmatch
from itertools import groupby, chain


def make_multistate_append_action(key):
    class _MultistateAppendAction(argparse.Action):
        def __call__(self, parser, namespace, values, option_string=None):
            args = getattr(namespace, self.dest)
            args = [] if args is None else args
            args.append((key, values))
            setattr(namespace, self.dest, args)

    return _MultistateAppendAction


def map_filters(strings, filters):
    strings = set(strings)

    if filters is None:
        return strings

    include_strings = set()
    for to_include, values in groupby(filters, key=lambda f: f[0]):
        patterns = set(chain(*(f[1] for f in values)))
        filtered_strings = set.union(*(set(fnmatch.filter(strings, p)) for p in patterns))

        if not to_include:
            include_strings |= strings - filtered_strings
        strings = filtered_strings

    if to_include:
        include_strings |= strings

    return include_strings


def _main():
    parser = argparse.ArgumentParser()

    parser.add_argument("inputfile")
    parser.add_argument("-i", "--include", nargs='+', dest="filters",
                        action=make_multistate_append_action(True))
    parser.add_argument("-e", "--exclude", nargs='+', dest="filters",
                        action=make_multistate_append_action(False))

    args = parser.parse_args()

    with open(args.inputfile) as f:
        strings = [l.strip() for l in f]

    for txt in map_filters(strings, args.filters):
        print(txt)


if __name__ == "__main__":
    _main()

$ python3 filter.py names.txt --include "A*" --exclude "?d*" --include "*m"
Abel
Aaron
Albert
Adam
Alan