Pythonで指定行数のcsvファイルを出力する

はじめに

先日、数千行のCSVファイルを作りたい場面がありました。
Excelやスプレッドシートで作るのは大変なので、ChatGPTに相談しながら、Pythonで指定行数のファイルが出力できるようにしました。

概要

• 動作確認した環境は Python 3.10.12
• 出力行数と出力先のパスは、引数で指定
• 出力ファイルの文字コードはutf-8
• 1行目はヘッダ行
• フィールドに test[5桁0埋めの連番]@example.com を出力

コード

少し長いので折りたたんでいます。

作成したコード

sample.py

import csv
import sys
import datetime

# インデックスを使用してテスト用のメールアドレスを生成する関数
def generate_email_address(index):
    return f"test{index:05d}@example.com"

# 指定された行数のCSVファイルを生成する関数
def generate_csv_file(output_path, num_rows, log_interval=1000):
    # ヘッダ行の作成
    header = "email_address"

    # 実行開始時刻の記録
    start_time = datetime.datetime.now()
    print(f"Execution started at {start_time.strftime('%Y-%m-%d %H:%M:%S')}")

    # CSVファイルの書き込み
    with open(output_path, mode='w', encoding='utf-8', newline='') as file:
        writer = csv.writer(file)

        # ヘッダ行を書き込み
        writer.writerow(header.split(','))

        # 指定された行数分のデータを生成して書き込み
        for i in range(1, num_rows + 1):
            email_address = generate_email_address(i)
            row_data = [email_address]
            writer.writerow(row_data)

            # ログ出力
            if i % log_interval == 0:
                print(f"Generated {i} rows")

    # 実行終了時刻の記録
    end_time = datetime.datetime.now()
    print(f"Execution finished at {end_time.strftime('%Y-%m-%d %H:%M:%S')}")
    elapsed_time = end_time - start_time
    print(f"Total execution time: {elapsed_time}")

if __name__ == "__main__":
    # コマンドライン引数の確認
    if len(sys.argv) != 3:
        print("Usage: python script.py <num_rows> <output_file_path>")
        sys.exit(1)

    try:
        # 指定された行数を取得
        num_rows = int(sys.argv[1])

        # 出力ファイルのパスを取得
        output_file = sys.argv[2]

        # 出力ファイル名と指定された行数でCSVファイルを生成
        generate_csv_file(output_file, num_rows)

        # 成功メッセージの表示
        print(f"CSV file generated successfully: {output_file}")
    except ValueError:
        # エラー処理：不正な入力の場合
        print("Invalid input. Please provide a valid number of rows.")
        sys.exit(1)

実行例

# 3000行を「output.csv」として出力する
python3 script.py 3000 output.csv

1000行出力ごとにログが出ます。最後には所要時間も出ます。

実行ログ

$ python3 script.py 3000 output.csv
Execution started at 2023-12-25 12:35:12
Generated 1000 rows
Generated 2000 rows
Generated 3000 rows
Execution finished at 2023-12-25 12:35:12
Total execution time: 0:00:00.011421
CSV file generated successfully: output.csv
$

出力内容

output.csv

email_address
[e-mail]
[e-mail]
[e-mail]
・・・
・・・
[e-mail]

おわりに

もし似たような場面があれば、お役に立てると嬉しいです。

アドベントカレンダー、他の記事もぜひご覧ください