pythonのargparseでのファイル入出力のちょっとだけ気になる点

python argparse

pythonのargparseでのファイル入出力には、ちょっとだけ気になる点があります。これはそれらについてのメモです。

ファイル入力と標準入力を一緒に使う場合

以下のような形で入力を使いたい場合があると思います。

  • 引数としてファイルを取る
  • ファイルが渡されなかった場合には標準入力を見る

以下の様な形です。

# <file>を見る
$ <cmd> --input-file <file>

# 標準入力を見る
$ echo '{"msg": "hello"}'|  <cmd>

このような場合分けに対応しようと思った場合に、変にifで分岐しなければいけないのが嫌ですね。

import sys
import argparse
import json

parser = argparse.ArgumentParser()
parser.add_argument("--input-file")
args = parser.parse_args()

# ここで条件分岐
if args.input_file is None:
    data = json.load(sys.stdin)
else:
    with open(args.input_file) as rf:
        data = json.load(rf)
print(data)

あと、標準入力を見るのに、sysモジュールをimportしなければいけないのは面倒です。

argparse.FileTypeを使う

argparse.ArgumentParserのadd_argument()はtypeという引数をとれます。これにFileTypeを渡してあげると良い感じにファイルIOになってくれます(close自体ができていないのですが。まぁcpythonでやる上ではgcに回収されるタイミングでfdは開放されるので。。とは言ってもという人には付録を書きました)。defaultに"-"を渡してあげるとsys.stdinやsys.stdoutに良い感じで変換されます。

なので以下の様なコードでOKです(条件分岐はなくなります)。

import argparse
import json

parser = argparse.ArgumentParser()
parser.add_argument("--input-file", type=argparse.FileType("r"), default="-")
args = parser.parse_args()
data = json.load(args.input_file)
print(data)

やりましたね。

$ python 01*.py --input-file data.json
{'msg': 'hello'}
$ cat data.json | python 01*.py
{'msg': 'hello'}

同様のことは出力にも？

同様のことは出力にも言えます。逆に保存するファイルを指定するようなコードを考えてみましょう。以下のような動作をするコードです。

$ python 10*.py
{
  "msg": "hello"
}
$ python 10*.py --output-file 10.json
$ cat 10.json
{
  "msg": "hello"
}

こちらもFileTypeを使うと便利です。

import argparse
import json

parser = argparse.ArgumentParser()
parser.add_argument("--output-file", type=argparse.FileType("w"), default="-")
args = parser.parse_args()

data = {"msg": "hello"}
json.dump(data, args.output_file, indent=2)

なんですけれど。存在しないディレクトリに所属するファイルのパスを指定したりするとダメです。事前にmkdir -p $(dirname <file>) のようなことをやってほしいのですが。ダメです。

$ python 10*.py --output-file /tmp/foo/bar/10.json
usage: 10write.py [-h] [--output-file OUTPUT_FILE]
10write.py: error: argument --output-file: can't open '/tmp/foo/bar/10.json': [Errno 2] No such file or directory: '/tmp/foo/bar/10.json'

かなしい。

mkdir -p 的なことがしたい

かなしいですが、仕方がないので継承したクラスを作ってみましょう。mkdir -p的なことはos.path.makedirsを使えばできます。

import argparse
import json
import os.path


class MkdirpFileType(argparse.FileType):
    def __call__(self, string):
        if string and "w" in self._mode:  # privateな紳士協定を破っているので行儀は良くない
            if os.path.dirname(string):
                os.makedirs(os.path.dirname(string), exist_ok=True)
        return super().__call__(string)


parser = argparse.ArgumentParser()
parser.add_argument("--output-file", type=MkdirpFileType("w"), default="-")
args = parser.parse_args()

data = {"msg": "hello"}
json.dump(data, args.output_file, indent=2)

ようやく存在しないファイルパスが渡されたときにはディレクトリを勝手に作ってくれるようになりました。

$ python 11*.py --output-file /tmp/foo/bar/11.json
$ cat /tmp/foo/bar/11.json
{
  "msg": "hello"
}

これで上手くいくのですが、ちょっとした機能の拡張が必要ということはライブラリを作るか毎回定義するかをしなければいけないので微妙です。この微妙なのがなんとかならないかな−というのがこの記事の主たる内容でした。

付録: しっかりwithを使ってclosingしたいひとはcontextlib

標準入力の所でwithを省略したコードを挙げていましたが、openをwithを使わず使うのが気になる人もいると思います。そのような人はcontextlib.closingを使うと便利です。 これはopenが返す値はcontext managerですがsys.stdinはcontext managerではないのですが、両者は共にclose()メソッドを持っているので大丈夫という話です。

import argparse
import contextlib
import json

parser = argparse.ArgumentParser()
parser.add_argument("--input-file", type=argparse.FileType("r"), default="-")
args = parser.parse_args()
with contextlib.closing(args.input_file) as rf:
    data = json.load(rf)
print(data)

実質、importの数が変わらないのでまぁ微妙かもですけれど。

別解としてcontextlib.ExitStackを使うという方法もあります。ExitStackはgoのdeferっぽいことをやってくれたり、今回のようなwith-syntaxを使った部分をifに置き換えたりということができたりします。

import argparse
import contextlib
import json

parser = argparse.ArgumentParser()
parser.add_argument("--input-file", type=argparse.FileType("r"), default="-")
args = parser.parse_args()
with contextlib.ExitStack() as s:
    if hasattr(args.input_file, "name"):
        s.enter_context(args.input_file)  # or s.callback(args.input_file.close)
    data = json.load(args.input_file)
print(data)

参考