ファイルに１行挿入する Python スクリプト作成中

PythonRecipe の 163 番目用に考えている最中かつ未テストのコード。Ruby レシピブックと同じく、一時ファイルを使用したもの。

import tempfile
import os
import stat

def _rename(f1, f2):
    if os.name == 'nt':
        os.remove(f2)
        os.rename(f1, f2)
    else:
        os.rename(f1, f2)


def _insert_line(org_name, temp_file, temp_name, start_line, data):
    org_f = open(org_name, 'r')
    try:
        for now_line, line in enumerate(org_f):
            if now_line == start_line: # 今何行目を処理中か確認、到達したら data を書く
                temp_file.write(data)
                temp_file.write(u'\n')
            temp_file.write(line)
    finally:
        org_f.close()
    temp_file.flush()


def insert_line(org_name, start_line, data):
    st = os.stat(org_name)
    mode = stat.S_IMODE(st.st_mode)

    temp = tempfile.mkstemp(text=True, dir=u'.')
    temp_file = os.fdopen(temp[0], 'w')
    temp_name = temp[1]
    try:
        _insert_line(org_name, temp_file, temp_name, start_line, data)
        temp_file.close()
        _rename(temp_name, org_name)
    except:
        temp_file.close()
        os.remove(temp_name)
        raise

    if hasattr(os, 'chmod'):
        os.chmod(org_name, mode)

insert_line(org_name, start_line, data) 関数。 org_name という名前のファイルの start_line 行目に data を挿入する。

課題がいくつか。１行処理するたびに if 文で今何行目か判定して到達したら data 挿入としている。やっぱり while 文で start 行数文書いて、挿入データ書いて、残り書くという Ruby レシピブックと同じ方法のほうがいいかなぁ。

アップする前にちゃんとコメントを入れよう、各関数ごとに一言説明を。

Windows の os.rename はすでにファイルが存在すると上書きできずに WindowsError になる。リファレンスマニュアル読み直しても「名前変更の原子的操作を実装する手段がないから」だとある。これでは今回の目標、行挿入のための「上書き」を達成できない。しかたないので _rename 関数を作って場合分け。今はとりあえず消してリネームするだけ。ここを事故防止のために元ファイルをリネームして取っておくようにするかも。いや、レシピ用だし例外処理詰め込みすぎて本来の処理が埋もれるのは避けるべきか。完成コードには含める、趣旨を説明するコードからは省くのが無難か。

if hasattr(os, 'chmod') も混乱の元か。 shutil.copymode を参考にした勢いで混じったが、これはきっと Mac OS 9 以前用の対策。

あとは・・・一時ファイルをおなじディレクトリに作らないと転ぶ環境の対策だ。 tempfile.mkstemp(text=True, dir=u'.') じゃだめだ。 dir 引数を直さないと。