銀月の符号

Python 使い見習いの日記・雑記

ファイルに1行挿入する Python スクリプト投稿

PythonRecipe163:ファイルに一行挿入する として書いてきました。ここはおかしいとか、こう書くべきというアイディアがある方、どうぞ書き換えてしまってください。よろしくお願いします。

import tempfile
import os
import stat

def _rename(f1, f2):
    u"""
    f1 を f2 にリネームする。
    Windows 上ではリネーム先にファイルが存在すると
    os.rename が WindowsError 例外となる。
    これを防ぐためリネーム先ファイルを削除する。
    """
    if os.name == 'nt':
        os.remove(f2)
        os.rename(f1, f2)
    else:
        os.rename(f1, f2)


def _insert_line(temp_file, org_name, start_line, data):
    u"""
    一時ファイルに挿入位置の前、つまり start_line 行目までをコピーする。
    その後 data を挿入し、残りデータをコピーする。

    start_line が元ファイルの行数より大きい場合、 data は挿入されない。
    start_line が負の値の場合 0 とみなされる。
    """
    org_f = open(org_name, 'r')
    try:
        for i in xrange(start_line):
            temp_file.write(org_f.next())
        temp_file.write(data)
        while True:
            temp_file.write(org_f.next())
    except StopIteration:
        pass
    finally:
        org_f.close()
    temp_file.flush()


def insert_line(org_name, start_line, data):
    u"""
    org_name ファイルの start_line 行目に data を挿入する。

    まず、元ファイル org_name の属性を保持する。
    そして元ファイルの start_line 行目に data を挿入してある一時ファイルを作成する。
    完成したらその一時ファイルを元のファイルにリネームし、属性を元に戻す。
    """
    st = os.stat(org_name)
    mode = stat.S_IMODE(st.st_mode)

    temp = tempfile.mkstemp(text=True, dir=os.path.dirname(org_name))
    temp_file = os.fdopen(temp[0], 'w')
    temp_name = temp[1]

    try:
        _insert_line(temp_file, org_name, start_line, data)
        temp_file.close()
        _rename(temp_name, org_name)
    except:
        # 処理に失敗したら一時ファイルを削除する
        temp_file.close()
        os.remove(temp_name)
        raise

    os.chmod(org_name, mode)