読者です 読者をやめる 読者になる 読者になる

PythonでCSVのFixtureをYAMLに変換

Python

たぶん、こんな感じじゃないかなぁ。

import sys
import csv
import yaml

def main(src, dest):
    reader = csv.reader(file(src))
    data = list(reader)

    result = []
    cols = data[0]
    rows = data[1:]
    for row in rows:
        result.append(dict([(k, v) for k, v in zip(cols, row)]))

    yaml.safe_dump(result, file(dest, 'wb'),
                   default_flow_style=False,
                   encoding='utf-8',
                   allow_unicode=True)

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print 'Usage: %s [filename]' % sys.argv[0]
        sys.exit(1)

    src = sys.argv[1]
    dest = src.replace('.csv', '.yml', src.rindex('.'))
    main(src, dest)

自分で文字列を組み立ててYAMLを作るよりはPyYAMLを使う。インストールは、

$ easy_install PyYAML

や、

$ pip install PyYAML

などで。

data = list(reader)

はちょっと富豪的なやり方かもしれないが、Fixtureのデータ量なんてたかが知れていると思い、こう書いている。

dumpでなくて、safe_dumpを使っているのは、普通にdumpを使うと、

>>> yaml.dump(dict(spam="スパム"))
'{spam: !!python/str "\\u30B9\\u30D1\\u30E0"}\n'

のようになって悲しい思いをするため。元データを読み取る時に適切な型に変換していけば、safe_dumpでなくdumpでもいけそうな気がしたが、今回のケースではsafe_dumpが手っ取り早いと思う。