msizの日記

ソフトウェア関係の覚え書きが中心になる予定

python の csv.writer で日本語(encoding指定)

前回の記事で、csv.readerのencoding指定版を作ってみたので、csv.writerも試してみました。

import csv

class UnicodeCsvWriter:
  def __init__(self, fileobj, dialect='excel', encoding='utf-8', *args, **kwds):
    self.writer = csv.writer(fileobj, dialect, *args, **kwds)
    self.dialect = dialect
    self.encoding = encoding
  def writerow(self, seq):
    row_seq = []
    for elem in seq:
      if isinstance(elem, unicode):
        row_seq.append(elem.encode(self.encoding))
      else:
        row_seq.append(elem)
    self.writer.writerow(row_seq)
  def writerows(self, seq_of_seq):
    for seq in seq_of_seq:
      self.writerow(seq)

class UnicodeDictCsvWriter(csv.DictWriter):
  def __init__(self, f, fieldnames, restval="", extrasaction="raise",
               dialect="excel", encoding='utf-8', *args, **kwds):
    csv.DictWriter.__init__(self, f, fieldnames, restval, extrasaction, dialect, *args, **kwds)
    self.writer = UnicodeCsvWriter(f, dialect, encoding=encoding)

sys.getdefaultencoding()がasciiでも、以下のコードでエラーが出ないはず。
(インタープリター起動時のLANGはen_GB.utf-8)

>>> import StringIO.StringIO
>>> s_io1 = StringIO.StringIO()
>>> u_writer = UnicodeCsvWriter(s_io, encoding='utf-8')
>>> u_writer = UnicodeCsvWriter(s_io1, encoding='utf-8')
>>> u_writer.writerow([1, u'あ', 'い'])
>>> u_writer.writerows([[1, u'あ', 'い'], [2, u'あああ', 22]])
>>> print(s_io1.getvalue())
1,あ,い
1,あ,い
2,あああ,22

>>>
>>> s_io2 = StringIO.StringIO()
>>> ud_writer = UnicodeDictCsvWriter(s_io2, [1,2,3], encoding='utf-8')
>>> ud_writer.writerow({ 1: 1, 2: u'あ', 3: 'い' })
>>> ud_writer.writerows([{ 1: 1, 2: u'あ', 3: 'い' }, {1: 11, 2: u'ああ', 3: 'いい' }])
>>>
>>>
>>> print(s_io2.getvalue())
1,あ,い
1,あ,い
11,ああ,いい