14 : 言語処理100本ノックでPythonのお勉強

第二章 : 14 先頭からN行目を出力

自然数Nをコマンドライン引数などの手段で受け取り,入力のうち先頭のN行だけを表示せよ.確認にはheadコマンドを用いよ.

コードはこんな感じ。

file = 'hightemp.txt'
N = 4

with open(file) as f:
    for i, line in enumerate(f):
        if i < N:
            line = line.rstrip()
            print(line)

enumerate()関数を使うとfor文回すときに中身とインデックスを両方取得できるので便利ですね。 line.rstrip()のところは改行コードを削除するために使っています。

unixコマンドではこんな感じで確認できます。
head -4 hightemp.txt'

コマンドラインからファイルと表示したい行数を指定して実行する場合はこんな感じになります。 sysモジュールを使うとargvのリストとしてコマンドライン引数を受け取ることができます。 第2引数は行数を指定するために渡していますが、文字列として受け取っているので、関数の中で整数に変換して使っています。

import sys

def get_head_line(file, N):
    N = int(N)
    with open(file) as f:
        for i, line in enumerate(f):
            if i < N:
                line = line.rstrip()
                print(line)

if __name__ == '__main__':
    get_head_line(sys.argv[1], sys.argv[2])